Add an account activation step in the signing request flow when needed

Signed-off-by: Émile Ré <emile@getprobo.com>

apps/console/src/hooks/useSafeContinueUrl.ts

@@ -0,0 +1,22 @@
+import { useSearchParams } from "react-router";
+
+export function useSafeContinueUrl(fallbackUrl?: URL): URL {
+  fallbackUrl = new URL(fallbackUrl ?? window.location.origin, window.location.origin);
+  const [searchParams] = useSearchParams();
+
+  const continueUrlParam = searchParams.get("continue");
+  let safeContinueUrl: URL;
+  if (continueUrlParam) {
+    let continueUrl: URL;
+    try {
+      continueUrl = new URL(continueUrlParam, window.location.origin);
+    } catch {
+      continueUrl = fallbackUrl;
+    }
+    safeContinueUrl = new URL(continueUrl.pathname + continueUrl.search, window.location.origin);
+  } else {
+    safeContinueUrl = fallbackUrl;
+  }
+
+  return safeContinueUrl;
+}

apps/console/src/pages/iam/auth/ActivateAccountPage.tsx

@@ -8,6 +8,7 @@ import { Link, useNavigate, useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 
 import type { ActivateAccountPageMutation$data, ActivateAccountPageMutation } from "#/__generated__/iam/ActivateAccountPageMutation.graphql";
+import { useSafeContinueUrl } from "#/hooks/useSafeContinueUrl";
 
 const activateAccountMutation = graphql`
   mutation ActivateAccountPageMutation(
@@ -22,9 +23,10 @@ const activateAccountMutation = graphql`
 export default function ActivateAccountPage() {
   const { __ } = useTranslate();
   const { toast } = useToast();
-  const navigate = useNavigate();
   const [searchParams] = useSearchParams();
+  const navigate = useNavigate();
   const submittedRef = useRef<boolean>(false);
+  const safeContinueUrl = useSafeContinueUrl();
 
   usePageTitle(__("Activate Account"));
 
@@ -44,6 +46,7 @@ export default function ActivateAccountPage() {
               window.location.href = "/";
               return;
             }
+            // FIXME: If already activated redirect too
           }
           toast({
             title: __("Activation failed"),
@@ -69,12 +72,22 @@ export default function ActivateAccountPage() {
         }
 
         if (activateAccount.createPasswordToken) {
+          const search = new URLSearchParams([
+            ["token", activateAccount.createPasswordToken],
+            ["continue", safeContinueUrl.toString()],
+          ]);
           void navigate(
-            { pathname: "/auth/create-password", search: `?token=${activateAccount.createPasswordToken}` },
+            {
+              pathname: "/auth/create-password",
+              search: "?" + search.toString(),
+            },
             { replace: true },
           );
         } else {
-          void navigate("/", { replace: true });
+          void navigate({
+            pathname: safeContinueUrl.pathname,
+            search: safeContinueUrl.search,
+          }, { replace: true });
         }
       },
       onError: (e) => {
@@ -85,7 +98,7 @@ export default function ActivateAccountPage() {
         });
       },
     });
-  }, [__, toast, activateAccount, navigate]);
+  }, [__, toast, activateAccount, navigate, safeContinueUrl]);
 
   useEffect(() => {
     const token = searchParams.get("token");

apps/console/src/pages/iam/auth/CreatePasswordPage.tsx

@@ -3,7 +3,7 @@ import { usePageTitle } from "@probo/hooks";
 import { useTranslate } from "@probo/i18n";
 import { Button, Field, useToast } from "@probo/ui";
 import { useMutation } from "react-relay";
-import { Link, useNavigate, useSearchParams } from "react-router";
+import { Link, useLocation, useNavigate, useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 import { z } from "zod";
 
@@ -27,6 +27,7 @@ export default function CreatePasswordPage() {
   const { toast } = useToast();
   const [searchParams] = useSearchParams();
   const navigate = useNavigate();
+  const location = useLocation();
 
   usePageTitle(__("Create Password"));
 
@@ -61,7 +62,12 @@ export default function CreatePasswordPage() {
           description: __("Account created successfully"),
           variant: "success",
         });
-        void navigate("/auth/login", { replace: true });
+        void navigate({
+          pathname: "/auth/login",
+          search: location.search,
+        }, {
+          replace: true,
+        });
       },
       onError: (e) => {
         toast({

apps/console/src/pages/iam/auth/sign-in/PasswordSignInPage.tsx

@@ -3,10 +3,11 @@ import { useTranslate } from "@probo/i18n";
 import { Button, Field, IconChevronLeft, useToast } from "@probo/ui";
 import type { FormEventHandler } from "react";
 import { useMutation } from "react-relay";
-import { Link, matchPath, useLocation, useSearchParams } from "react-router";
+import { Link, matchPath, useLocation } from "react-router";
 import { graphql } from "relay-runtime";
 
 import type { PasswordSignInPageMutation } from "#/__generated__/iam/PasswordSignInPageMutation.graphql";
+import { useSafeContinueUrl } from "#/hooks/useSafeContinueUrl";
 
 const signInMutation = graphql`
   mutation PasswordSignInPageMutation($input: SignInInput!) {
@@ -20,7 +21,7 @@ const signInMutation = graphql`
 
 export default function PasswordSignInPage() {
   const location = useLocation();
-  const [searchParams] = useSearchParams();
+  const safeContinueUrl = useSafeContinueUrl();
 
   const { __ } = useTranslate();
   const { toast } = useToast();
@@ -36,20 +37,6 @@ export default function PasswordSignInPage() {
 
     if (!emailValue || !passwordValue) return;
 
-    const continueUrlParam = searchParams.get("continue");
-    let safeContinueUrl: URL;
-    if (continueUrlParam) {
-      let continueUrl: URL;
-      try {
-        continueUrl = new URL(continueUrlParam, window.location.origin);
-      } catch {
-        continueUrl = new URL(window.location.origin);
-      }
-      safeContinueUrl = new URL(continueUrl.pathname + continueUrl.search, window.location.origin);
-    } else {
-      safeContinueUrl = new URL(window.location.origin);
-    }
-
     const match = matchPath(
       { path: "/organizations/:organizationId", caseSensitive: false, end: false },
       safeContinueUrl.pathname,

apps/console/src/pages/iam/organizations/AssumePage.tsx

@@ -8,6 +8,7 @@ import { graphql } from "relay-runtime";
 import type { AssumePageMutation } from "#/__generated__/iam/AssumePageMutation.graphql";
 import type { AssumePageQuery } from "#/__generated__/iam/AssumePageQuery.graphql";
 import { useOrganizationId } from "#/hooks/useOrganizationId";
+import { useSafeContinueUrl } from "#/hooks/useSafeContinueUrl";
 
 import AuthLayout from "../auth/AuthLayout";
 
@@ -48,28 +49,23 @@ export function AssumePage(props: { queryRef: PreloadedQuery<AssumePageQuery> })
   const [searchParams] = useSearchParams();
   const { __ } = useTranslate();
 
+  const safeContinueUrl = useSafeContinueUrl(
+    new URL(window.location.origin + `/organizations/${organizationId}`),
+  );
+
   const { viewer } = usePreloadedQuery<AssumePageQuery>(assumePageQuery, queryRef);
   const [assumeOrganizationSession] = useMutation<AssumePageMutation>(assumeMutation);
 
-  const continueUrlParam = searchParams.get("continue");
-  let safeContinueUrl: string;
-  if (continueUrlParam) {
-    const continueUrl = new URL(continueUrlParam);
-    safeContinueUrl = window.location.origin + continueUrl.pathname + continueUrl.search;
-  } else {
-    safeContinueUrl = window.location.origin + `/organizations/${organizationId}`;
-  }
-
   useEffect(() => {
     assumeOrganizationSession({
       variables: {
-        input: { organizationId, continue: safeContinueUrl },
+        input: { organizationId, continue: safeContinueUrl.toString() },
       },
       onError: (error) => {
         if (error instanceof UnAuthenticatedError) {
           const search = new URLSearchParams([
             ["organization-id", organizationId],
-            ["continue", safeContinueUrl],
+            ["continue", safeContinueUrl.toString()],
           ]);
 
           void navigate({ pathname: "/auth/login", search: "?" + search.toString() });
@@ -88,7 +84,7 @@ export function AssumePage(props: { queryRef: PreloadedQuery<AssumePageQuery> })
         switch (result.__typename) {
           case "PasswordRequired":
             search.set("organization-id", organizationId);
-            search.set("continue", safeContinueUrl);
+            search.set("continue", safeContinueUrl.toString());
 
             void navigate({ pathname: "/auth/password-login", search: "?" + search.toString() });
             break;
@@ -102,7 +98,7 @@ export function AssumePage(props: { queryRef: PreloadedQuery<AssumePageQuery> })
             window.location.href = samlSSOLoginURL.toString();
             break;
           default:
-            window.location.href = safeContinueUrl;
+            window.location.href = safeContinueUrl.toString();
         }
       },
     });