Implement mfa login flow

ericallam · ericallam · commit e7cb0d678d1f · 2025-07-06T19:54:34.000+01:00
diff --git a/apps/webapp/app/routes/auth.github.callback.tsx b/apps/webapp/app/routes/auth.github.callback.tsx
@@ -1,25 +1,53 @@
 import type { LoaderFunction } from "@remix-run/node";
+import { redirect } from "@remix-run/node";
 import { authenticator } from "~/services/auth.server";
 import { redirectCookie } from "./auth.github";
+import { getUserSession, commitSession } from "~/services/sessionStorage.server";
 import { logger } from "~/services/logger.server";
+import { MfaRequiredError } from "~/services/mfa/multiFactorAuthentication.server";
 
 export let loader: LoaderFunction = async ({ request }) => {
-  const cookie = request.headers.get("Cookie");
-  const redirectValue = await redirectCookie.parse(cookie);
-  const redirectTo = redirectValue ?? "/";
+  try {
+    const cookie = request.headers.get("Cookie");
+    const redirectValue = await redirectCookie.parse(cookie);
+    const redirectTo = redirectValue ?? "/";
 
-  logger.debug("auth.github.callback loader", {
-    redirectTo,
-  });
+    logger.debug("auth.github.callback loader", {
+      redirectTo,
+    });
 
-  const authuser = await authenticator.authenticate("github", request, {
-    successRedirect: redirectTo,
-    failureRedirect: "/login",
-  });
+    const authuser = await authenticator.authenticate("github", request, {
+      successRedirect: undefined, // Don't auto-redirect, we'll handle it
+      failureRedirect: undefined, // Don't auto-redirect on failure either
+    });
 
-  logger.debug("auth.github.callback authuser", {
-    authuser,
-  });
+    logger.debug("auth.github.callback authuser", {
+      authuser,
+    });
 
-  return authuser;
+    // If we get here, user doesn't have MFA - complete login normally
+    return redirect(redirectTo);
+  } catch (error) {
+    // Check if this is an MFA_REQUIRED error
+    if (error instanceof MfaRequiredError) {
+      // User has MFA enabled - store pending user ID and redirect to MFA page
+      const session = await getUserSession(request);
+      session.set("pending-mfa-user-id", error.userId);
+
+      const cookie = request.headers.get("Cookie");
+      const redirectValue = await redirectCookie.parse(cookie);
+      const redirectTo = redirectValue ?? "/";
+      session.set("pending-mfa-redirect-to", redirectTo);
+
+      return redirect("/login/mfa", {
+        headers: {
+          "Set-Cookie": await commitSession(session),
+        },
+      });
+    }
+
+    // Regular authentication failure, redirect to login page
+    logger.debug("auth.github.callback error", { error });
+    return redirect("/login");
+  }
 };
diff --git a/apps/webapp/app/routes/login.mfa/route.tsx b/apps/webapp/app/routes/login.mfa/route.tsx
@@ -15,7 +15,10 @@ import { InputOTP, InputOTPGroup, InputOTPSlot } from "~/components/primitives/I
 import { Paragraph } from "~/components/primitives/Paragraph";
 import { Spinner } from "~/components/primitives/Spinner";
 import { authenticator } from "~/services/auth.server";
-import { commitSession, getUserSession } from "~/services/sessionStorage.server";
+import { commitSession, getUserSession, sessionStorage } from "~/services/sessionStorage.server";
+import { MultiFactorAuthenticationService } from "~/services/mfa/multiFactorAuthentication.server";
+import { redirectWithErrorMessage } from "~/models/message.server";
+import { ServiceValidationError } from "~/v3/services/baseService.server";
 
 export const meta: MetaFunction = ({ matches }) => {
   const parentMeta = matches
@@ -37,11 +40,20 @@ export const meta: MetaFunction = ({ matches }) => {
 };
 
 export async function loader({ request }: LoaderFunctionArgs) {
+  // Check if user is already fully authenticated
   await authenticator.isAuthenticated(request, {
     successRedirect: "/",
   });
 
   const session = await getUserSession(request);
+  
+  // Check if there's a pending MFA user ID
+  const pendingUserId = session.get("pending-mfa-user-id");
+  if (!pendingUserId) {
+    // No pending MFA, redirect to login
+    return redirect("/login");
+  }
+
   const error = session.get("auth:error");
 
   let mfaError: string | undefined;
@@ -64,42 +76,100 @@ export async function loader({ request }: LoaderFunctionArgs) {
 }
 
 export async function action({ request }: ActionFunctionArgs) {
-  const clonedRequest = request.clone();
+  try {
+    const session = await getUserSession(request);
+    const pendingUserId = session.get("pending-mfa-user-id");
+    
+    if (!pendingUserId) {
+      return redirect("/login");
+    }
 
-  const payload = Object.fromEntries(await clonedRequest.formData());
+    const payload = Object.fromEntries(await request.formData());
 
-  const { action } = z
-    .object({
-      action: z.enum(["verify-recovery", "verify-mfa"]),
-    })
-    .parse(payload);
+    const { action } = z
+      .object({
+        action: z.enum(["verify-recovery", "verify-mfa"]),
+      })
+      .parse(payload);
 
-  if (action === "verify-recovery") {
-    // TODO: Implement recovery code verification logic
-    const recoveryCode = payload.recoveryCode;
+    const mfaService = new MultiFactorAuthenticationService();
 
-    // For now, just redirect to dashboard
-    return redirect("/");
-  } else if (action === "verify-mfa") {
-    // TODO: Implement MFA code verification logic
-    const mfaCode = payload.mfaCode;
+    if (action === "verify-recovery") {
+      const recoveryCode = payload.recoveryCode as string;
+      
+      if (!recoveryCode) {
+        session.set("auth:error", { message: "Recovery code is required" });
+        return redirect("/login/mfa", {
+          headers: { "Set-Cookie": await commitSession(session) },
+        });
+      }
 
-    // For now, just redirect to dashboard
-    return redirect("/");
-  } else {
-    const session = await getUserSession(request);
-    session.unset("triggerdotdev:magiclink");
+      const result = await mfaService.verifyRecoveryCodeForLogin(pendingUserId, recoveryCode);
+      
+      if (!result.success) {
+        session.set("auth:error", { message: result.error });
+        return redirect("/login/mfa", {
+          headers: { "Set-Cookie": await commitSession(session) },
+        });
+      }
 
-    return redirect("/login/magic", {
-      headers: {
-        "Set-Cookie": await commitSession(session),
-      },
-    });
+      // Recovery code verified - complete the login
+      return await completeLogin(request, session, pendingUserId);
+
+    } else if (action === "verify-mfa") {
+      const mfaCode = payload.mfaCode as string;
+      
+      if (!mfaCode || mfaCode.length !== 6) {
+        session.set("auth:error", { message: "Valid 6-digit code is required" });
+        return redirect("/login/mfa", {
+          headers: { "Set-Cookie": await commitSession(session) },
+        });
+      }
+
+      const result = await mfaService.verifyTotpForLogin(pendingUserId, mfaCode);
+      
+      if (!result.success) {
+        session.set("auth:error", { message: result.error });
+        return redirect("/login/mfa", {
+          headers: { "Set-Cookie": await commitSession(session) },
+        });
+      }
+
+      // TOTP code verified - complete the login
+      return await completeLogin(request, session, pendingUserId);
+    }
+
+    return redirect("/login");
+    
+  } catch (error) {
+    if (error instanceof ServiceValidationError) {
+      return redirectWithErrorMessage("/login", request, error.message);
+    }
+    throw error;
   }
 }
 
+async function completeLogin(request: Request, session: any, userId: string) {
+  // Create a new authenticated session
+  const authSession = await sessionStorage.getSession(request.headers.get("Cookie"));
+  authSession.set(authenticator.sessionKey, { userId });
+  
+  // Get the redirect URL and clean up pending MFA data
+  const redirectTo = session.get("pending-mfa-redirect-to") ?? "/";
+  session.unset("pending-mfa-user-id");
+  session.unset("pending-mfa-redirect-to");
+  session.unset("auth:error");
+  
+  return redirect(redirectTo, {
+    headers: {
+      "Set-Cookie": await sessionStorage.commitSession(authSession),
+    },
+  });
+}
+
 export default function LoginMfaPage() {
-  const { mfaError } = useTypedLoaderData<typeof loader>();
+  const data = useTypedLoaderData<typeof loader>();
+  const mfaError = 'mfaError' in data ? data.mfaError : undefined;
   const navigate = useNavigation();
   const [showRecoveryCode, setShowRecoveryCode] = useState(false);
   const [mfaCode, setMfaCode] = useState("");
@@ -151,7 +221,7 @@ export default function LoginMfaPage() {
                     <span className="text-text-bright">Verify</span>
                   )}
                 </Button>
-                {mfaError && <FormError>{mfaError}</FormError>}
+                {typeof mfaError === 'string' && <FormError>{mfaError}</FormError>}
               </Fieldset>
               <Button
                 type="button"
@@ -203,7 +273,7 @@ export default function LoginMfaPage() {
                     <span className="text-text-bright">Verify</span>
                   )}
                 </Button>
-                {mfaError && <FormError>{mfaError}</FormError>}
+                {typeof mfaError === 'string' && <FormError>{mfaError}</FormError>}
               </Fieldset>
               <Button
                 type="button"
diff --git a/apps/webapp/app/routes/magic.tsx b/apps/webapp/app/routes/magic.tsx
@@ -1,12 +1,39 @@
 import type { LoaderFunctionArgs } from "@remix-run/server-runtime";
+import { redirect } from "@remix-run/server-runtime";
 import { authenticator } from "~/services/auth.server";
+import { MfaRequiredError } from "~/services/mfa/multiFactorAuthentication.server";
 import { getRedirectTo } from "~/services/redirectTo.server";
+import { getUserSession, commitSession } from "~/services/sessionStorage.server";
 
 export async function loader({ request }: LoaderFunctionArgs) {
-  const redirectTo = await getRedirectTo(request);
+  try {
+    // Attempt to authenticate the user with email-link
+    const authUser = await authenticator.authenticate("email-link", request, {
+      successRedirect: undefined, // Don't auto-redirect, we'll handle it
+      failureRedirect: undefined, // Don't auto-redirect on failure either
+    });
 
-  await authenticator.authenticate("email-link", request, {
-    successRedirect: redirectTo ?? "/",
-    failureRedirect: "/login/magic",
-  });
+    // If we get here, user doesn't have MFA - complete login normally
+    const redirectTo = await getRedirectTo(request);
+    return redirect(redirectTo ?? "/");
+  } catch (error) {
+    // Check if this is an MFA_REQUIRED error
+    if (error instanceof MfaRequiredError) {
+      // User has MFA enabled - store pending user ID and redirect to MFA page
+      const session = await getUserSession(request);
+      session.set("pending-mfa-user-id", error.userId);
+
+      const redirectTo = await getRedirectTo(request);
+      session.set("pending-mfa-redirect-to", redirectTo ?? "/");
+
+      return redirect("/login/mfa", {
+        headers: {
+          "Set-Cookie": await commitSession(session),
+        },
+      });
+    }
+
+    // Regular authentication failure, redirect to magic link page
+    return redirect("/login/magic");
+  }
 }
diff --git a/apps/webapp/app/services/emailAuth.server.tsx b/apps/webapp/app/services/emailAuth.server.tsx
@@ -6,6 +6,7 @@ import { env } from "~/env.server";
 import { sendMagicLinkEmail } from "~/services/email.server";
 import { postAuthentication } from "./postAuth.server";
 import { logger } from "./logger.server";
+import { MfaRequiredError } from "./mfa/multiFactorAuthentication.server";
 
 let secret = env.MAGIC_LINK_SECRET;
 if (!secret) throw new Error("Missing MAGIC_LINK_SECRET env variable.");
@@ -36,8 +37,19 @@ const emailStrategy = new EmailLinkStrategy(
 
       await postAuthentication({ user, isNewUser, loginMethod: "MAGIC_LINK" });
 
+      // Check if user has MFA enabled
+      if (user.mfaEnabledAt) {
+        // Throw a special error that will be caught by the magic route
+        throw new MfaRequiredError(user.id);
+      }
+
       return { userId: user.id };
     } catch (error) {
+      // Skip logging the error if it's a MfaRequiredError
+      if (error instanceof MfaRequiredError) {
+        throw error;
+      }
+
       logger.debug("Magic link user failed to authenticate", { error: JSON.stringify(error) });
       throw error;
     }
diff --git a/apps/webapp/app/services/gitHubAuth.server.ts b/apps/webapp/app/services/gitHubAuth.server.ts
@@ -5,6 +5,7 @@ import { findOrCreateUser } from "~/models/user.server";
 import type { AuthUser } from "./authUser";
 import { postAuthentication } from "./postAuth.server";
 import { logger } from "./logger.server";
+import { MfaRequiredError } from "./mfa/multiFactorAuthentication.server";
 
 export function addGitHubStrategy(
   authenticator: Authenticator<AuthUser>,
@@ -40,10 +41,21 @@ export function addGitHubStrategy(
 
         await postAuthentication({ user, isNewUser, loginMethod: "GITHUB" });
 
+        // Check if user has MFA enabled
+        if (user.mfaEnabledAt) {
+          // Throw a special error that will be caught by the callback route
+          throw new MfaRequiredError(user.id);
+        }
+
         return {
           userId: user.id,
         };
       } catch (error) {
+        // Skip logging the error if it's a MfaRequiredError
+        if (error instanceof MfaRequiredError) {
+          throw error;
+        }
+
         console.error(error);
         throw error;
       }
diff --git a/apps/webapp/app/services/mfa/multiFactorAuthentication.server.ts b/apps/webapp/app/services/mfa/multiFactorAuthentication.server.ts
