Implement redirect-path for password method

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

apps/console/src/components/OrganizationErrorBoundary.tsx

@@ -1,5 +1,5 @@
 import { AssumptionRequiredError, UnAuthenticatedError } from "@probo/relay";
-import { Navigate, useLocation, useRouteError } from "react-router";
+import { Navigate, useRouteError } from "react-router";
 
 import { useOrganizationId } from "#/hooks/useOrganizationId";
 
@@ -7,15 +7,26 @@ import { PageError } from "./PageError";
 
 export function OrganizationErrorBoundary() {
   const error = useRouteError();
-  const location = useLocation();
   const organizationId = useOrganizationId();
 
+  const search = new URLSearchParams([
+    ["organization-id", organizationId],
+    ["redirect-path", window.location.href],
+  ]);
+
   if (error instanceof UnAuthenticatedError) {
-    return <Navigate to="/auth/login" state={{ from: location.pathname }} />;
+    return <Navigate to={{ pathname: "/auth/login", search: "?" + search.toString() }} />;
   }
 
   if (error instanceof AssumptionRequiredError) {
-    return <Navigate to={`/organizations/${organizationId}/assume`} state={{ from: location.pathname }} />;
+    return (
+      <Navigate
+        to={{
+          pathname: `/organizations/${organizationId}/assume`,
+          search: "?" + search.toString(),
+        }}
+      />
+    );
   }
 
   return <PageError error={error instanceof Error ? error : new Error("unknown error")} />;

apps/console/src/components/RootErrorBoundary.tsx

@@ -1,14 +1,15 @@
 import { UnAuthenticatedError } from "@probo/relay";
-import { Navigate, useLocation, useRouteError } from "react-router";
+import { Navigate, useRouteError } from "react-router";
 
 import { PageError } from "./PageError";
 
 export function RootErrorBoundary() {
   const error = useRouteError();
-  const location = useLocation();
+
+  const search = new URLSearchParams([["redirect-path", window.location.href]]);
 
   if (error instanceof UnAuthenticatedError) {
-    return <Navigate to="/auth/login" state={{ from: location.pathname }} />;
+    return <Navigate to={{ pathname: "/auth/login", search: "?" + search.toString() }} />;
   }
 
   return <PageError error={error instanceof Error ? error : new Error("unknown error")} />;

apps/console/src/hooks/iam/useAssume.ts

@@ -9,6 +9,7 @@ import type { useAssumeMutation } from "#/__generated__/iam/useAssumeMutation.gr
 import { useOrganizationId } from "../useOrganizationId";
 
 interface UseAssumeParameters {
+  afterAssumePath: string;
   onSuccess: () => void;
 }
 
@@ -41,7 +42,7 @@ const assumeMutation = graphql`
 `;
 
 export function useAssume(params: UseAssumeParameters) {
-  const { onSuccess } = params;
+  const { afterAssumePath, onSuccess } = params;
 
   const organizationId = useOrganizationId();
   const navigate = useNavigate();
@@ -55,7 +56,12 @@ export function useAssume(params: UseAssumeParameters) {
       },
       onError: (error) => {
         if (error instanceof UnAuthenticatedError) {
-          void navigate("/auth/login");
+          const search = new URLSearchParams([
+            ["organization-id", organizationId],
+            ["redirect-path", afterAssumePath],
+          ]);
+
+          void navigate({ pathname: "/auth/login", search: "?" + search.toString() });
           return;
         }
       },
@@ -69,7 +75,9 @@ export function useAssume(params: UseAssumeParameters) {
 
         switch (result.__typename) {
           case "PasswordRequired":
-            search.set("organizationId", organizationId);
+            search.set("organization-id", organizationId);
+            search.set("redirect-path", afterAssumePath);
+
             void navigate({ pathname: "/auth/passord-login", search: "?" + search.toString() });
             break;
           case "SAMLAuthenticationRequired":
@@ -80,7 +88,7 @@ export function useAssume(params: UseAssumeParameters) {
         }
       },
     });
-  }, [onSuccess, navigate, assumeOrganizationSession, organizationId]);
+  }, [afterAssumePath, organizationId, onSuccess, navigate, assumeOrganizationSession]);
 
   return;
 }

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

@@ -41,7 +41,7 @@ export default function PasswordSignInPage() {
           email: emailValue,
           password: passwordValue,
           // Assume when signing in
-          organizationId: searchParams.get("organizationId"),
+          organizationId: searchParams.get("organization-id"),
         },
       },
       onCompleted: (_, error) => {
@@ -57,7 +57,7 @@ export default function PasswordSignInPage() {
           return;
         }
 
-        window.location.href = "/";
+        window.location.href = searchParams.get("redirect-path") ?? "/";
       },
       onError: (e) => {
         toast({

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

@@ -1,10 +1,12 @@
 import { useTranslate } from "@probo/i18n";
 import { Button } from "@probo/ui";
-import { Link } from "react-router";
+import { Link, useLocation } from "react-router";
 
 export default function SignInPage() {
   const { __ } = useTranslate();
 
+  const location = useLocation();
+
   return (
     <div className="space-y-6 w-full max-w-md mx-auto pt-8">
       <h1 className="text-center text-2xl font-bold">
@@ -14,7 +16,10 @@ export default function SignInPage() {
         {__("Choose your login method")}
       </p>
 
-      <Button className="w-xs h-10 mx-auto" to="/auth/password-login">
+      <Button
+        className="w-xs h-10 mx-auto"
+        to={{ pathname: "/auth/password-login", search: location.search }}
+      >
         {__("Login with Email")}
       </Button>
 
@@ -27,14 +32,21 @@ export default function SignInPage() {
         </span>
       </div>
 
-      <Button variant="secondary" className="w-xs h-10 mx-auto" to="/auth/sso-login">
+      <Button
+        variant="secondary"
+        className="w-xs h-10 mx-auto"
+        to={{ pathname: "/auth/sso-login", search: location.search }}
+      >
         {__("Login with SSO")}
       </Button>
 
       <div className="text-center mt-6 text-sm text-txt-secondary">
         {__("Don't have an account ?")}
         {" "}
-        <Link to="/auth/register" className="underline hover:text-txt-primary">
+        <Link
+          to={{ pathname: "/auth/register", search: location.search }}
+          className="underline hover:text-txt-primary"
+        >
           {__("Register")}
         </Link>
       </div>

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

@@ -1,23 +1,19 @@
 import { useTranslate } from "@probo/i18n";
-import { useLocation, useNavigate } from "react-router";
+import { useNavigate, useSearchParams } from "react-router";
 
 import { useAssume } from "#/hooks/iam/useAssume";
+import { IAMRelayProvider } from "#/providers/IAMRelayProvider";
 
 import AuthLayout from "../auth/AuthLayout";
 
-interface State {
-  from: string;
-}
-
-export default function AssumePage() {
+function AssumePageInner() {
   const navigate = useNavigate();
-  const location = useLocation();
-  const state = location.state as State;
-
+  const [searchParams] = useSearchParams();
   const { __ } = useTranslate();
 
   useAssume({
-    onSuccess: () => void navigate(state.from),
+    afterAssumePath: searchParams.get("redirect-path") ?? "/",
+    onSuccess: () => void navigate(searchParams.get("redirect-path") ?? "/"),
   });
 
   return (
@@ -33,3 +29,11 @@ export default function AssumePage() {
     </AuthLayout>
   );
 }
+
+export default function AssumePage() {
+  return (
+    <IAMRelayProvider>
+      <AssumePageInner />
+    </IAMRelayProvider>
+  );
+}

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

@@ -1,6 +1,7 @@
 import { Skeleton } from "@probo/ui";
 import { Suspense, useCallback } from "react";
 import { useQueryLoader } from "react-relay";
+import { useLocation } from "react-router";
 
 import type { ViewerMembershipLayoutQuery } from "#/__generated__/iam/ViewerMembershipLayoutQuery.graphql";
 import { useAssume } from "#/hooks/iam/useAssume";
@@ -14,6 +15,7 @@ import {
 
 function ViewerMembershipLayoutQueryLoader() {
   const organizationId = useOrganizationId();
+  const location = useLocation();
 
   const [queryRef, loadQuery] = useQueryLoader<ViewerMembershipLayoutQuery>(
     viewerMembershipLayoutQuery,
@@ -28,7 +30,10 @@ function ViewerMembershipLayoutQueryLoader() {
     [loadQuery, organizationId],
   );
 
-  useAssume({ onSuccess: onAssumeSuccess });
+  useAssume({
+    afterAssumePath: location.pathname,
+    onSuccess: onAssumeSuccess,
+  });
 
   if (!queryRef) {
     return <Skeleton className="w-full h-screen" />;

apps/console/src/pages/organizations/employee/EmployeeLayoutLoader.tsx

@@ -1,6 +1,7 @@
 import { Skeleton } from "@probo/ui";
 import { Suspense, useCallback } from "react";
 import { useQueryLoader } from "react-relay";
+import { useLocation } from "react-router";
 
 import type { ViewerMembershipLayoutQuery } from "#/__generated__/iam/ViewerMembershipLayoutQuery.graphql";
 import { useAssume } from "#/hooks/iam/useAssume";
@@ -17,6 +18,7 @@ function EmployeeLayoutQueryLoader() {
   const [queryRef, loadQuery] = useQueryLoader<ViewerMembershipLayoutQuery>(
     viewerMembershipLayoutQuery,
   );
+  const location = useLocation();
 
   const onAssumeSuccess = useCallback(
     () =>
@@ -27,7 +29,10 @@ function EmployeeLayoutQueryLoader() {
     [loadQuery, organizationId],
   );
 
-  useAssume({ onSuccess: onAssumeSuccess });
+  useAssume({
+    afterAssumePath: location.pathname,
+    onSuccess: onAssumeSuccess,
+  });
 
   if (!queryRef) {
     return <Skeleton className="w-full h-screen" />;

apps/console/src/routes.tsx

@@ -122,6 +122,7 @@ const routes = [
   },
   {
     path: "/organizations/:organizationId",
+    ErrorBoundary: OrganizationErrorBoundary,
     children: [
       {
         path: "assume",
@@ -132,7 +133,6 @@ const routes = [
         Component: lazy(
           () => import("./pages/organizations/employee/EmployeeLayoutLoader"),
         ),
-        ErrorBoundary: OrganizationErrorBoundary,
         children: [
           {
             index: true,
@@ -154,9 +154,9 @@ const routes = [
         Component: lazy(
           () => import("./pages/iam/organizations/ViewerMembershipLayoutLoader"),
         ),
-        ErrorBoundary: OrganizationErrorBoundary,
         children: [
           {
+            path: "",
             Component: () => {
               const { role } = use(CurrentUser);
               switch (role) {

packages/ui/src/Atoms/Button/Button.tsx

@@ -4,7 +4,7 @@ import type {
   FC,
   PropsWithChildren,
 } from "react";
-import { Link } from "react-router";
+import { Link, type To } from "react-router";
 import { tv, type VariantProps } from "tailwind-variants";
 
 import { Slot } from "../Slot";
@@ -49,7 +49,7 @@ type Props = PropsWithChildren<
       | "tertiary"
       | "quaternary"
       | "danger";
-    to?: string;
+    to?: To;
     asChild?: boolean;
   } & VariantProps<typeof button>
 >

pkg/iam/session_service.go

@@ -321,12 +321,11 @@ func (s SessionService) OpenPasswordChildSessionForOrganization(
 	organizationID gid.GID,
 ) (*coredata.Session, *coredata.Membership, error) {
 	var (
-		now          = time.Now()
-		rootSession  = &coredata.Session{}
-		identity     = &coredata.Identity{}
-		membership   = &coredata.Membership{}
-		childSession = &coredata.Session{}
-		scope        = coredata.NewScopeFromObjectID(organizationID)
+		now         = time.Now()
+		rootSession = &coredata.Session{}
+		identity    = &coredata.Identity{}
+		membership  = &coredata.Membership{}
+		scope       = coredata.NewScopeFromObjectID(organizationID)
 	)
 
 	err := s.pg.WithTx(
@@ -366,7 +365,7 @@ func (s SessionService) OpenPasswordChildSessionForOrganization(
 			}
 
 			tenantID := scope.GetTenantID()
-			childSession = &coredata.Session{
+			childSession := &coredata.Session{
 				ID:              gid.New(tenantID, coredata.SessionEntityType),
 				IdentityID:      rootSession.IdentityID,
 				TenantID:        &tenantID,
@@ -399,7 +398,7 @@ func (s SessionService) OpenPasswordChildSessionForOrganization(
 		return nil, nil, err
 	}
 
-	return childSession, membership, nil
+	return rootSession, membership, nil
 }
 
 // OpenSAMLChildSessionForOrganization creates a SAML-authenticated child session for the given