Add ssoLoginURL in activate account outputs

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

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

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<852ecee082c322aba273f04575353fbd>>
+ * @generated SignedSource<<e2bdaf924ff82d11f3c2ef4a8353c8e8>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -18,6 +18,7 @@ export type ActivateAccountPageMutation$variables = {
 export type ActivateAccountPageMutation$data = {
   readonly activateAccount: {
     readonly createPasswordToken: string | null | undefined;
+    readonly ssoLoginUrl: string | null | undefined;
   } | null | undefined;
 };
 export type ActivateAccountPageMutation = {
@@ -54,6 +55,13 @@ v1 = [
         "kind": "ScalarField",
         "name": "createPasswordToken",
         "storageKey": null
+      },
+      {
+        "alias": null,
+        "args": null,
+        "kind": "ScalarField",
+        "name": "ssoLoginUrl",
+        "storageKey": null
       }
     ],
     "storageKey": null
@@ -77,16 +85,16 @@ return {
     "selections": (v1/*: any*/)
   },
   "params": {
-    "cacheID": "388c90522ecf3a3365e171a623fb2d2e",
+    "cacheID": "ab62d7332fce4af79b46c8627e551370",
     "id": null,
     "metadata": {},
     "name": "ActivateAccountPageMutation",
     "operationKind": "mutation",
-    "text": "mutation ActivateAccountPageMutation(\n  $input: ActivateAccountInput!\n) {\n  activateAccount(input: $input) {\n    createPasswordToken\n  }\n}\n"
+    "text": "mutation ActivateAccountPageMutation(\n  $input: ActivateAccountInput!\n) {\n  activateAccount(input: $input) {\n    createPasswordToken\n    ssoLoginUrl\n  }\n}\n"
   }
 };
 })();
 
-(node as any).hash = "c27b723383934423bf92926820426653";
+(node as any).hash = "fd91da524bb46256a9d8e1ae31590dc2";
 
 export default node;

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

@@ -16,6 +16,7 @@ const activateAccountMutation = graphql`
   ) {
     activateAccount(input: $input) {
       createPasswordToken
+      ssoLoginUrl
     }
   }
 `;
@@ -73,6 +74,14 @@ export default function ActivateAccountPage() {
           throw new Error("mutation data missing");
         }
 
+        if (activateAccount.ssoLoginUrl) {
+          const url = new URL(activateAccount.ssoLoginUrl);
+          url.search = searchParams.toString();
+
+          window.location.href = url.toString();
+          return;
+        }
+
         if (activateAccount.createPasswordToken) {
           const search = new URLSearchParams([
             ["token", activateAccount.createPasswordToken],
@@ -85,12 +94,14 @@ export default function ActivateAccountPage() {
             },
             { replace: true },
           );
-        } else {
+          return;
-          void navigate({
-            pathname: safeContinueUrl.pathname,
-            search: safeContinueUrl.search,
-          }, { replace: true });
         }
+
+        const search = new URLSearchParams([["continue", safeContinueUrl.toString()]]);
+        void navigate({
+          pathname: "/auth/password-login",
+          search: "?" + search.toString(),
+        }, { replace: true });
       },
       onError: (e) => {
         toast({
@@ -100,7 +111,7 @@ export default function ActivateAccountPage() {
         });
       },
     });
-  }, [__, toast, activateAccount, navigate, safeContinueUrl]);
+  }, [__, toast, activateAccount, navigate, safeContinueUrl, searchParams]);
 
   useEffect(() => {
     const token = searchParams.get("token");

pkg/iam/auth_service.go

@@ -135,7 +135,7 @@ func (req CreateIdentityWithPasswordRequest) Validate() error {
 func (s *AuthService) ActivateAccount(
 	ctx context.Context,
 	req *ActivateAccountRequest,
-) (*coredata.MembershipProfile, *string, error) {
+) (*coredata.Identity, *coredata.MembershipProfile, error) {
 	if err := req.Validate(); err != nil {
 		return nil, nil, fmt.Errorf("invalid request: %w", err)
 	}
@@ -151,10 +151,9 @@ func (s *AuthService) ActivateAccount(
 		profile    *coredata.MembershipProfile
 		identity   *coredata.Identity
 		now        = time.Now()
-		createPasswordToken *string
 	)
 
-	err = s.pg.WithTx(
+	if err = s.pg.WithTx(
 		ctx,
 		func(tx pg.Conn) error {
 			err := invitation.LoadByID(ctx, tx, scope, payload.Data.InvitationID)
@@ -229,36 +228,25 @@ func (s *AuthService) ActivateAccount(
 
 			return nil
 		},
-	)
+	); err != nil {
-
-	if err != nil {
 		return nil, nil, err
 	}
 
-	count, err := s.AccountService.CountSAMLConfigurationsForEmail(ctx, identity.EmailAddress)
+	return identity, profile, nil
-	if err != nil {
-		return nil, nil, fmt.Errorf("cannot count SAML configurations: %w", err)
 }
 
-	if count > 0 {
+func (s AuthService) GetResetPasswordToken(ctx context.Context, email mail.Addr) (string, error) {
-		return profile, nil, nil
-	}
-
-	if identity.HashedPassword == nil {
 	token, err := statelesstoken.NewToken(
 		s.tokenSecret,
 		TokenTypePasswordReset,
 		s.passwordResetTokenValidity,
-			PasswordResetData{Email: identity.EmailAddress},
+		PasswordResetData{Email: email},
 	)
 	if err != nil {
-			return nil, nil, fmt.Errorf("cannot generate password create token: %w", err)
+		return "", fmt.Errorf("cannot generate password create token: %w", err)
 	}
 
-		createPasswordToken = &token
+	return token, nil
-	}
-
-	return profile, createPasswordToken, nil
 }
 
 func (s AuthService) ResetPassword(
@@ -313,12 +301,7 @@ func (s AuthService) SendPasswordResetInstructionByEmail(
 	ctx context.Context,
 	email mail.Addr,
 ) error {
-	token, err := statelesstoken.NewToken(
+	token, err := s.GetResetPasswordToken(ctx, email)
-		s.tokenSecret,
-		TokenTypePasswordReset,
-		s.passwordResetTokenValidity,
-		PasswordResetData{Email: email},
-	)
 	if err != nil {
 		return fmt.Errorf("cannot generate password reset token: %w", err)
 	}

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

@@ -788,6 +788,7 @@ type SignOutPayload {
 
 type ActivateAccountPayload {
   createPasswordToken: String
+  ssoLoginUrl: String
   profile: Profile
 }
 

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

@@ -65,6 +65,7 @@ type ComplexityRoot struct {
 	ActivateAccountPayload struct {
 		CreatePasswordToken func(childComplexity int) int
 		Profile             func(childComplexity int) int
+		SsoLoginURL         func(childComplexity int) int
 	}
 
 	AssumeOrganizationSessionPayload struct {
@@ -623,6 +624,12 @@ func (e *executableSchema) Complexity(ctx context.Context, typeName, field strin
 		}
 
 		return e.ComplexityRoot.ActivateAccountPayload.Profile(childComplexity), true
+	case "ActivateAccountPayload.ssoLoginUrl":
+		if e.ComplexityRoot.ActivateAccountPayload.SsoLoginURL == nil {
+			break
+		}
+
+		return e.ComplexityRoot.ActivateAccountPayload.SsoLoginURL(childComplexity), true
 
 	case "AssumeOrganizationSessionPayload.result":
 		if e.ComplexityRoot.AssumeOrganizationSessionPayload.Result == nil {
@@ -3206,6 +3213,7 @@ type SignOutPayload {
 
 type ActivateAccountPayload {
   createPasswordToken: String
+  ssoLoginUrl: String
   profile: Profile
 }
 
@@ -4195,6 +4203,35 @@ func (ec *executionContext) fieldContext_ActivateAccountPayload_createPasswordTo
 	return fc, nil
 }
 
+func (ec *executionContext) _ActivateAccountPayload_ssoLoginUrl(ctx context.Context, field graphql.CollectedField, obj *types.ActivateAccountPayload) (ret graphql.Marshaler) {
+	return graphql.ResolveField(
+		ctx,
+		ec.OperationContext,
+		field,
+		ec.fieldContext_ActivateAccountPayload_ssoLoginUrl,
+		func(ctx context.Context) (any, error) {
+			return obj.SsoLoginURL, nil
+		},
+		nil,
+		ec.marshalOString2ᚖstring,
+		true,
+		false,
+	)
+}
+
+func (ec *executionContext) fieldContext_ActivateAccountPayload_ssoLoginUrl(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
+	fc = &graphql.FieldContext{
+		Object:     "ActivateAccountPayload",
+		Field:      field,
+		IsMethod:   false,
+		IsResolver: false,
+		Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) {
+			return nil, errors.New("field of type String does not have child fields")
+		},
+	}
+	return fc, nil
+}
+
 func (ec *executionContext) _ActivateAccountPayload_profile(ctx context.Context, field graphql.CollectedField, obj *types.ActivateAccountPayload) (ret graphql.Marshaler) {
 	return graphql.ResolveField(
 		ctx,
@@ -6342,6 +6379,8 @@ func (ec *executionContext) fieldContext_Mutation_activateAccount(ctx context.Co
 			switch field.Name {
 			case "createPasswordToken":
 				return ec.fieldContext_ActivateAccountPayload_createPasswordToken(ctx, field)
+			case "ssoLoginUrl":
+				return ec.fieldContext_ActivateAccountPayload_ssoLoginUrl(ctx, field)
 			case "profile":
 				return ec.fieldContext_ActivateAccountPayload_profile(ctx, field)
 			}
@@ -16444,6 +16483,8 @@ func (ec *executionContext) _ActivateAccountPayload(ctx context.Context, sel ast
 			out.Values[i] = graphql.MarshalString("ActivateAccountPayload")
 		case "createPasswordToken":
 			out.Values[i] = ec._ActivateAccountPayload_createPasswordToken(ctx, field, obj)
+		case "ssoLoginUrl":
+			out.Values[i] = ec._ActivateAccountPayload_ssoLoginUrl(ctx, field, obj)
 		case "profile":
 			out.Values[i] = ec._ActivateAccountPayload_profile(ctx, field, obj)
 		default:

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

@@ -31,6 +31,7 @@ type ActivateAccountInput struct {
 
 type ActivateAccountPayload struct {
 	CreatePasswordToken *string  `json:"createPasswordToken,omitempty"`
+	SsoLoginURL         *string  `json:"ssoLoginUrl,omitempty"`
 	Profile             *Profile `json:"profile,omitempty"`
 }
 

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

@@ -371,7 +371,7 @@ func (r *mutationResolver) ActivateAccount(ctx context.Context, input types.Acti
 		r.sessionCookie.Clear(w)
 	}
 
-	user, createPasswordToken, err := r.iam.AuthService.ActivateAccount(
+	identity, user, err := r.iam.AuthService.ActivateAccount(
 		ctx,
 		&iam.ActivateAccountRequest{
 			InvitationToken: input.Token,
@@ -400,8 +400,42 @@ func (r *mutationResolver) ActivateAccount(ctx context.Context, input types.Acti
 		return nil, gqlutils.Internal(ctx)
 	}
 
+	var ssoLoginURL *string
+	samlConfigs, err := r.iam.AccountService.ListSAMLConfigurationsForEmail(ctx, user.EmailAddress)
+	if err != nil {
+		r.logger.ErrorCtx(ctx, "cannot list saml configurations", log.Error(err))
+		return nil, gqlutils.Internal(ctx)
+	}
+
+	for _, samlConfig := range samlConfigs {
+		if samlConfig.OrganizationID != user.OrganizationID {
+			continue
+		}
+
+		ssoLoginURL = new(r.SSOLoginURL(samlConfig.ID))
+	}
+
+	if ssoLoginURL != nil {
+		return &types.ActivateAccountPayload{
+			CreatePasswordToken: nil,
+			SsoLoginURL:         ssoLoginURL,
+			Profile:             types.NewProfile(user),
+		}, nil
+	}
+
+	var createPasswordToken *string
+	if identity.HashedPassword == nil {
+		token, err := r.iam.AuthService.GetResetPasswordToken(ctx, identity.EmailAddress)
+		if err != nil {
+			return nil, fmt.Errorf("cannot generate password create token: %w", err)
+		}
+
+		createPasswordToken = &token
+	}
+
 	return &types.ActivateAccountPayload{
 		CreatePasswordToken: createPasswordToken,
+		SsoLoginURL:         nil,
 		Profile:             types.NewProfile(user),
 	}, nil
 }