Implement redirect path for saml

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

apps/console/src/__generated__/iam/AssumePageMutation.graphql.ts

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<096c8befe586a961b8ebf410eda0e019>>
+ * @generated SignedSource<<8717e91a7e1602c06f6c67fa76580ecf>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -12,6 +12,7 @@ import { ConcreteRequest } from 'relay-runtime';
 export type ReauthenticationReason = "POLICY_REQUIREMENT" | "SENSITIVE_ACTION" | "SESSION_EXPIRED";
 export type AssumeOrganizationSessionInput = {
   organizationId: string;
+  redirectPath: string;
 };
 export type AssumePageMutation$variables = {
   input: AssumeOrganizationSessionInput;

apps/console/src/__generated__/iam/MembershipsDropdownMenuItem_assumeMutation.graphql.ts

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<4beff809268fb32711cc02699683b623>>
+ * @generated SignedSource<<f829648ef4b9a5c0b8c4d62ac844bca0>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -12,6 +12,7 @@ import { ConcreteRequest } from 'relay-runtime';
 export type ReauthenticationReason = "POLICY_REQUIREMENT" | "SENSITIVE_ACTION" | "SESSION_EXPIRED";
 export type AssumeOrganizationSessionInput = {
   organizationId: string;
+  redirectPath: string;
 };
 export type MembershipsDropdownMenuItem_assumeMutation$variables = {
   input: AssumeOrganizationSessionInput;

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

@@ -3,7 +3,7 @@ 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, useSearchParams } from "react-router";
+import { Link, useLocation, useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 
 import type { PasswordSignInPageMutation } from "#/__generated__/iam/PasswordSignInPageMutation.graphql";
@@ -19,6 +19,7 @@ const signInMutation = graphql`
 `;
 
 export default function PasswordSignInPage() {
+  const location = useLocation();
   const [searchParams] = useSearchParams();
 
   const { __ } = useTranslate();
@@ -72,7 +73,7 @@ export default function PasswordSignInPage() {
   return (
     <form className="space-y-6 w-full max-w-md mx-auto pt-4" onSubmit={handlePasswordLogin}>
       <Link
-        to="/auth/login"
+        to={{ pathname: "/auth/login", search: location.search }}
         className="flex items-center gap-2 text-txt-secondary hover:text-txt-primary transition-colors mb-4"
       >
         <IconChevronLeft size={20} />
@@ -112,7 +113,7 @@ export default function PasswordSignInPage() {
       <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/auth/sign-in/SSOSignInPage.tsx

@@ -6,7 +6,7 @@ import {
   usePreloadedQuery,
   useQueryLoader,
 } from "react-relay";
-import { Link, useNavigate } from "react-router";
+import { Link, useLocation, useNavigate, useSearchParams } from "react-router";
 import { graphql } from "relay-runtime";
 
 import type { SSOSignInPageQuery } from "#/__generated__/iam/SSOSignInPageQuery.graphql";
@@ -18,6 +18,7 @@ const ssoAvailabilityQuery = graphql`
 `;
 
 export default function SSOSignInPage() {
+  const location = useLocation();
   const { __ } = useTranslate();
 
   const [queryRef, loadQuery]
@@ -39,7 +40,7 @@ export default function SSOSignInPage() {
     <>
       <form className="space-y-6 w-full max-w-md mx-auto pt-4" onSubmit={handleSSOCheck}>
         <Link
-          to="/auth/login"
+          to={{ pathname: "/auth/register", search: location.search }}
           className="flex items-center gap-2 text-txt-secondary hover:text-txt-primary transition-colors mb-4"
         >
           <IconChevronLeft size={20} />
@@ -70,7 +71,7 @@ export default function SSOSignInPage() {
           {__("Don't have an account ?")}
           {" "}
           <Link
-            to="/auth/register"
+            to={{ pathname: "/auth/register", search: location.search }}
             className="underline hover:text-txt-primary"
           >
             {__("Register")}
@@ -96,6 +97,7 @@ function NavigateToSSOLoginURL(props: {
 
   const { __ } = useTranslate();
   const { toast } = useToast();
+  const [searchParams] = useSearchParams();
   const navigate = useNavigate();
 
   const { ssoLoginURL } = usePreloadedQuery<SSOSignInPageQuery>(
@@ -127,8 +129,11 @@ function NavigateToSSOLoginURL(props: {
       return;
     }
 
-    window.location.href = ssoLoginURL.value;
+    const url = new URL(ssoLoginURL.value);
-  }, [__, navigate, ssoLoginURL, toast]);
+    url.search = "?" + searchParams.toString();
+
+    window.location.href = url.toString();
+  }, [__, navigate, ssoLoginURL, toast, searchParams]);
 
   return null;
 }

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

@@ -38,19 +38,18 @@ function AssumePageInner() {
 
   const [assumeOrganizationSession] = useMutation<AssumePageMutation>(assumeMutation);
 
+  const redirectPath = searchParams.get("redirect-path") ?? `/organizations/${organizationId}`;
+
   useEffect(() => {
     assumeOrganizationSession({
       variables: {
-        input: { organizationId },
+        input: { organizationId, redirectPath },
       },
       onError: (error) => {
         if (error instanceof UnAuthenticatedError) {
           const search = new URLSearchParams([
             ["organization-id", organizationId],
-            [
+            ["redirect-path", redirectPath],
-              "redirect-path",
-              searchParams.get("redirect-path") ?? window.location.pathname + window.location.search,
-            ],
           ]);
 
           void navigate({ pathname: "/auth/login", search: "?" + search.toString() });
@@ -64,26 +63,27 @@ function AssumePageInner() {
 
         const { result } = assumeOrganizationSession;
         const search = new URLSearchParams();
+        let samlSSOLoginURL: URL;
 
         switch (result.__typename) {
           case "PasswordRequired":
             search.set("organization-id", organizationId);
-            search.set(
+            search.set("redirect-path", redirectPath);
-              "redirect-path",
-              searchParams.get("redirect-path") ?? window.location.pathname + window.location.search,
-            );
 
             void navigate({ pathname: "/auth/passord-login", search: "?" + search.toString() });
             break;
           case "SAMLAuthenticationRequired":
-            window.location.href = result.redirectUrl;
+            samlSSOLoginURL = new URL(result.redirectUrl);
+            samlSSOLoginURL.search = "?" + searchParams.toString();
+
+            window.location.href = samlSSOLoginURL.toString();
             break;
           default:
-            void navigate(searchParams.get("redirect-path") ?? window.location.pathname + window.location.search);
+            void navigate(redirectPath);
         }
       },
     });
-  }, [organizationId, navigate, assumeOrganizationSession, searchParams]);
+  }, [organizationId, navigate, assumeOrganizationSession, redirectPath, searchParams]);
 
   return (
     <AuthLayout>

pkg/iam/saml/service.go

@@ -107,6 +107,7 @@ func (s *Service) GenerateSpMetadata() ([]byte, error) {
 func (s *Service) InitiateLogin(
 	ctx context.Context,
 	configID gid.GID,
+	redirectPath string,
 ) (*url.URL, error) {
 	var (
 		now           = time.Now()
@@ -131,7 +132,7 @@ func (s *Service) InitiateLogin(
 				return NewSAMLDisabledError()
 			}
 
-			sp, err := s.serviceProvider(ctx, config)
+			sp, err := s.serviceProvider(config)
 			if err != nil {
 				return fmt.Errorf("cannot build service provider: %w", err)
 			}
@@ -152,7 +153,11 @@ func (s *Service) InitiateLogin(
 				return fmt.Errorf("cannot insert SAML request: %w", err)
 			}
 
-			redirect, err = req.Redirect(config.ID.String(), sp)
+			relayState := url.Values{}
+			relayState.Add("config-id", config.ID.String())
+			relayState.Add("redirect-path", redirectPath)
+
+			redirect, err = req.Redirect(url.QueryEscape(relayState.Encode()), sp)
 			if err != nil {
 				return fmt.Errorf("cannot generate redirect URL: %w", err)
 			}
@@ -196,7 +201,7 @@ func (s *Service) HandleAssertion(
 				return NewSAMLDisabledError()
 			}
 
-			sp, err := s.serviceProvider(ctx, config)
+			sp, err := s.serviceProvider(config)
 			if err != nil {
 				return fmt.Errorf("cannot create service provider: %w", err)
 			}

pkg/iam/saml/sp.go

@@ -15,7 +15,6 @@
 package saml
 
 import (
-	"context"
 	"encoding/base64"
 	"fmt"
 
@@ -24,7 +23,6 @@ import (
 )
 
 func (s *Service) serviceProvider(
-	ctx context.Context,
 	config *coredata.SAMLConfiguration,
 ) (*saml.ServiceProvider, error) {
 	cert, err := config.GetIdPCertificate()

pkg/iam/session_service.go

@@ -490,6 +490,7 @@ func (s SessionService) AssumeOrganizationSession(
 	ctx context.Context,
 	sessionID gid.GID,
 	organizationID gid.GID,
+	redirectPath string,
 ) (*coredata.Session, *coredata.Membership, error) {
 	var (
 		now          = time.Now()
@@ -558,7 +559,7 @@ func (s SessionService) AssumeOrganizationSession(
 
 			if err == nil && samlConfig.EnforcementPolicy == coredata.SAMLEnforcementPolicyRequired {
 				if rootSession.AuthMethod != coredata.AuthMethodSAML {
-					redirectURL, err := s.SAMLService.InitiateLogin(ctx, samlConfig.ID)
+					redirectURL, err := s.SAMLService.InitiateLogin(ctx, samlConfig.ID, redirectPath)
 					if err != nil {
 						return fmt.Errorf("cannot initiate SAML login: %w", err)
 					}

pkg/server/api/connect/v1/saml_handler.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"net/url"
 
 	"github.com/go-chi/chi/v5"
 	"go.gearno.de/kit/httpserver"
@@ -58,18 +59,36 @@ func (h *SAMLHandler) ConsumeHandler(w http.ResponseWriter, r *http.Request) {
 	samlResponse := r.FormValue("SAMLResponse")
 	relayState := r.FormValue("RelayState")
 
-	configID, err := gid.ParseGID(relayState)
+	values, err := url.ParseQuery(relayState)
 	if err != nil {
 		httpserver.RenderError(w, http.StatusBadRequest, errors.New("invalid relay state"))
 		return
 	}
 
+	configIDStr := values.Get("config-id")
+	if configIDStr == "" {
+		httpserver.RenderError(w, http.StatusBadRequest, errors.New("missing config ID"))
+		return
+	}
+
+	configID, err := gid.ParseGID(configIDStr)
+	if err != nil {
+		httpserver.RenderError(w, http.StatusBadRequest, errors.New("invalid config ID"))
+		return
+	}
+
+	redirectPath := values.Get("redirect-path")
+
 	user, membership, err := h.iam.SAMLService.HandleAssertion(ctx, samlResponse, configID)
 	if err != nil {
 		httpserver.RenderError(w, http.StatusUnauthorized, err)
 		return
 	}
 
+	if redirectPath == "" {
+		redirectPath = "/organizations/" + membership.OrganizationID.String()
+	}
+
 	rootSession := authn.SessionFromContext(ctx)
 
 	switch {
@@ -105,7 +124,7 @@ func (h *SAMLHandler) ConsumeHandler(w http.ResponseWriter, r *http.Request) {
 
 	h.sessionCookie.Set(w, rootSession)
 
-	redirectURL := h.baseURL.WithPath("/organizations/" + membership.OrganizationID.String()).MustString()
+	redirectURL := h.baseURL.WithPath(redirectPath).MustString()
 	http.Redirect(w, r, redirectURL, http.StatusFound)
 }
 
@@ -118,13 +137,15 @@ func (h *SAMLHandler) LoginHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	redirectPathQueryParam := r.URL.Query().Get("redirect-path")
+
 	samlConfigID, err := gid.ParseGID(samlConfigIDParam)
 	if err != nil {
 		httpserver.RenderError(w, http.StatusBadRequest, errors.New("invalid SAML config ID"))
 		return
 	}
 
-	url, err := h.iam.SAMLService.InitiateLogin(ctx, samlConfigID)
+	url, err := h.iam.SAMLService.InitiateLogin(ctx, samlConfigID, redirectPathQueryParam)
 	if err != nil {
 		panic(fmt.Errorf("cannot initiate SAML login: %w", err))
 	}

pkg/server/api/connect/v1/schema.graphql

@@ -669,6 +669,7 @@ input ChangeEmailInput {
 
 input AssumeOrganizationSessionInput {
   organizationId: ID!
+  redirectPath: String!
 }
 
 input RevokeSessionInput {

pkg/server/api/connect/v1/schema/schema.go

@@ -3111,6 +3111,7 @@ input ChangeEmailInput {
 
 input AssumeOrganizationSessionInput {
   organizationId: ID!
+  redirectPath: String!
 }
 
 input RevokeSessionInput {
@@ -14930,7 +14931,7 @@ func (ec *executionContext) unmarshalInputAssumeOrganizationSessionInput(ctx con
 		asMap[k] = v
 	}
 
-	fieldsInOrder := [...]string{"organizationId"}
+	fieldsInOrder := [...]string{"organizationId", "redirectPath"}
 	for _, k := range fieldsInOrder {
 		v, ok := asMap[k]
 		if !ok {
@@ -14944,6 +14945,13 @@ func (ec *executionContext) unmarshalInputAssumeOrganizationSessionInput(ctx con
 				return it, err
 			}
 			it.OrganizationID = data
+		case "redirectPath":
+			ctx := graphql.WithPathContext(ctx, graphql.NewPathWithField("redirectPath"))
+			data, err := ec.unmarshalNString2string(ctx, v)
+			if err != nil {
+				return it, err
+			}
+			it.RedirectPath = data
 		}
 	}
 

pkg/server/api/connect/v1/types/types.go

@@ -36,6 +36,7 @@ type AcceptInvitationPayload struct {
 
 type AssumeOrganizationSessionInput struct {
 	OrganizationID gid.GID `json:"organizationId"`
+	RedirectPath   string  `json:"redirectPath"`
 }
 
 type AssumeOrganizationSessionPayload struct {

pkg/server/api/connect/v1/v1_resolver.go

@@ -680,7 +680,7 @@ func (r *mutationResolver) ChangeEmail(ctx context.Context, input types.ChangeEm
 func (r *mutationResolver) AssumeOrganizationSession(ctx context.Context, input types.AssumeOrganizationSessionInput) (*types.AssumeOrganizationSessionPayload, error) {
 	rootSession := authn.SessionFromContext(ctx)
 
-	childSession, membership, err := r.iam.SessionService.AssumeOrganizationSession(ctx, rootSession.ID, input.OrganizationID)
+	childSession, membership, err := r.iam.SessionService.AssumeOrganizationSession(ctx, rootSession.ID, input.OrganizationID, input.RedirectPath)
 	if err != nil {
 		var (
 			errMembershipNotFound         *iam.ErrMembershipNotFound