checkafter redirect, continue after setup

apps/login/app/(login)/mfa/set/page.tsx

@@ -20,7 +20,7 @@ export default async function Page({
 }: {
   searchParams: Record<string | number | symbol, string | undefined>;
 }) {
-  const { loginName, altPassword, authRequestId, organization, sessionId } =
+  const { loginName, checkAfter, authRequestId, organization, sessionId } =
     searchParams;
 
   const sessionFactors = sessionId
@@ -95,6 +95,7 @@ export default async function Page({
             organization={organization}
             loginSettings={loginSettings}
             userMethods={sessionFactors.authMethods ?? []}
+            checkAfter={checkAfter === "true"}
           ></ChooseSecondFactorToSetup>
         ) : (
           <Alert>No second factors available to setup.</Alert>

apps/login/app/(login)/otp/[method]/set/page.tsx

@@ -7,11 +7,14 @@ import {
   server,
 } from "#/lib/zitadel";
 import Alert from "#/ui/Alert";
+import { Button, ButtonVariants } from "#/ui/Button";
 import DynamicTheme from "#/ui/DynamicTheme";
+import { Spinner } from "#/ui/Spinner";
 import TOTPRegister from "#/ui/TOTPRegister";
 import UserAvatar from "#/ui/UserAvatar";
 import { getMostRecentCookieWithLoginname } from "#/utils/cookies";
 import { RegisterTOTPResponse } from "@zitadel/server";
+import Link from "next/link";
 import { ClientError } from "nice-grpc";
 
 export default async function Page({
@@ -21,7 +24,8 @@ export default async function Page({
   searchParams: Record<string | number | symbol, string | undefined>;
   params: Record<string | number | symbol, string | undefined>;
 }) {
-  const { loginName, organization, sessionId, authRequestId } = searchParams;
+  const { loginName, organization, sessionId, authRequestId, checkAfter } =
+    searchParams;
   const { method } = params;
 
   const branding = await getBrandingSettings(server, organization);
@@ -65,6 +69,34 @@ export default async function Page({
     });
   }
 
+  const paramsToContinue = new URLSearchParams({});
+  let urlToContinue = "/accounts";
+
+  if (authRequestId && sessionId) {
+    if (sessionId) {
+      paramsToContinue.append("sessionId", sessionId);
+    }
+    if (authRequestId) {
+      paramsToContinue.append("authRequestId", authRequestId);
+    }
+    if (organization) {
+      paramsToContinue.append("organization", organization);
+    }
+    urlToContinue = `/login?` + paramsToContinue;
+  } else if (loginName) {
+    if (loginName) {
+      paramsToContinue.append("loginName", loginName);
+    }
+    if (authRequestId) {
+      paramsToContinue.append("authRequestId", authRequestId);
+    }
+    if (organization) {
+      paramsToContinue.append("organization", organization);
+    }
+
+    urlToContinue = `/signedin?` + paramsToContinue;
+  }
+
   return (
     <DynamicTheme branding={branding}>
       <div className="flex flex-col items-center space-y-4">
@@ -107,17 +139,39 @@ export default async function Page({
                 sessionId={sessionId}
                 authRequestId={authRequestId}
                 organization={organization}
+                checkAfter={checkAfter === "true"}
               ></TOTPRegister>
             </div>{" "}
           </>
         ) : (
-          <p className="ztdl-p">
-            {method === "email"
-              ? "Code via email was successfully added."
-              : method === "sms"
-              ? "Code via SMS was successfully added."
-              : ""}
-          </p>
+          <>
+            <p className="ztdl-p">
+              {method === "email"
+                ? "Code via email was successfully added."
+                : method === "sms"
+                ? "Code via SMS was successfully added."
+                : ""}
+            </p>
+
+            <div className="mt-8 flex w-full flex-row items-center">
+              <span className="flex-grow"></span>
+              <Link
+                href={
+                  checkAfter
+                    ? `/otp/${method}?` + new URLSearchParams()
+                    : urlToContinue
+                }
+              >
+                <Button
+                  type="submit"
+                  className="self-end"
+                  variant={ButtonVariants.Primary}
+                >
+                  continue
+                </Button>
+              </Link>
+            </div>
+          </>
         )}
       </div>
     </DynamicTheme>

apps/login/ui/ChooseSecondFactorToSetup.tsx

@@ -17,6 +17,7 @@ type Props = {
   organization?: string;
   loginSettings: LoginSettings;
   userMethods: AuthenticationMethodType[];
+  checkAfter: boolean;
 };
 
 export default function ChooseSecondFactorToSetup({
@@ -26,6 +27,7 @@ export default function ChooseSecondFactorToSetup({
   organization,
   loginSettings,
   userMethods,
+  checkAfter,
 }: Props) {
   const cardClasses = (alreadyAdded: boolean) =>
     clsx(
@@ -47,6 +49,9 @@ export default function ChooseSecondFactorToSetup({
   if (organization) {
     params.append("organization", organization);
   }
+  if (checkAfter) {
+    params.append("checkAfter", "true");
+  }
 
   const TOTP = (alreadyAdded: boolean) => {
     return (

apps/login/ui/PasswordForm.tsx

@@ -120,6 +120,7 @@ export default function PasswordForm({
       } else if (loginSettings?.forceMfa && !resp.authFactors?.length) {
         const params = new URLSearchParams({
           loginName: resp.factors.user.loginName,
+          checkAfter: "true", // this defines if the check is directly made after the setup
         });
 
         if (authRequestId) {

apps/login/ui/TOTPRegister.tsx

@@ -23,6 +23,7 @@ type Props = {
   sessionId?: string;
   authRequestId?: string;
   organization?: string;
+  checkAfter?: boolean;
 };
 export default function TOTPRegister({
   uri,
@@ -31,6 +32,7 @@ export default function TOTPRegister({
   sessionId,
   authRequestId,
   organization,
+  checkAfter,
 }: Props) {
   const [error, setError] = useState<string>("");
   const [loading, setLoading] = useState<boolean>(false);
@@ -48,22 +50,13 @@ export default function TOTPRegister({
     return verifyTOTP(values.code, loginName, organization)
       .then((response) => {
         setLoading(false);
-        if (authRequestId && sessionId) {
-          const params = new URLSearchParams({
-            sessionId: sessionId,
-            authRequest: authRequestId,
-          });
+        // if attribute is set, validate MFA after it is setup, otherwise proceed as usual (when mfa is enforced to login)
+        if (checkAfter) {
+          const params = new URLSearchParams({});
 
-          if (organization) {
-            params.append("organization", organization);
+          if (loginName) {
+            params.append("loginName", loginName);
           }
-
-          return router.push(`/login?` + params);
-        } else if (loginName) {
-          const params = new URLSearchParams({
-            loginName,
-          });
-
           if (authRequestId) {
             params.append("authRequestId", authRequestId);
           }
@@ -71,7 +64,33 @@ export default function TOTPRegister({
             params.append("organization", organization);
           }
 
-          return router.push(`/signedin?` + params);
+          return router.push(`/otp/time-based?` + params);
+        } else {
+          if (authRequestId && sessionId) {
+            const params = new URLSearchParams({
+              sessionId: sessionId,
+              authRequest: authRequestId,
+            });
+
+            if (organization) {
+              params.append("organization", organization);
+            }
+
+            return router.push(`/login?` + params);
+          } else if (loginName) {
+            const params = new URLSearchParams({
+              loginName,
+            });
+
+            if (authRequestId) {
+              params.append("authRequestId", authRequestId);
+            }
+            if (organization) {
+              params.append("organization", organization);
+            }
+
+            return router.push(`/signedin?` + params);
+          }
         }
       })
       .catch((e) => {