Implement redirect path for saml

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

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