Fix non-MFA login

apps/webapp/app/routes/auth.github.callback.tsx

@@ -1,53 +1,58 @@
 import type { LoaderFunction } from "@remix-run/node";
 import { redirect } from "@remix-run/node";
+import { prisma } from "~/db.server";
+import { getSession, redirectWithErrorMessage } from "~/models/message.server";
 import { authenticator } from "~/services/auth.server";
+import { commitSession } from "~/services/sessionStorage.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 }) => {
-  try {
-    const cookie = request.headers.get("Cookie");
-    const redirectValue = await redirectCookie.parse(cookie);
-    const redirectTo = redirectValue ?? "/";
+  const cookie = request.headers.get("Cookie");
+  const redirectValue = await redirectCookie.parse(cookie);
+  const redirectTo = redirectValue ?? "/";
+
+  const auth = await authenticator.authenticate("github", request, {
+    failureRedirect: "/login", // If auth fails, the failureRedirect will be thrown as a Response
+  });
+
+  // manually get the session
+  const session = await getSession(request.headers.get("cookie"));
+
+  const userRecord = await prisma.user.findFirst({
+    where: {
+      id: auth.userId,
+    },
+    select: {
+      id: true,
+      mfaEnabledAt: true,
+    },
+  });
+
+  if (!userRecord) {
+    return redirectWithErrorMessage(
+      "/login",
+      request,
+      "Could not find your account. Please contact support."
+    );
+  }
 
-    logger.debug("auth.github.callback loader", {
-      redirectTo,
-    });
+  if (userRecord.mfaEnabledAt) {
+    session.set("pending-mfa-user-id", userRecord.id);
+    session.set("pending-mfa-redirect-to", redirectTo);
 
-    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
+    return redirect("/login/mfa", {
+      headers: {
+        "Set-Cookie": await commitSession(session),
+      },
     });
+  }
 
-    logger.debug("auth.github.callback authuser", {
-      authuser,
-    });
+  // and store the user data
+  session.set(authenticator.sessionKey, auth);
 
-    // 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");
-  }
+  return redirect(redirectTo, {
+    headers: {
+      "Set-Cookie": await commitSession(session),
+    },
+  });
 };

apps/webapp/app/routes/magic.tsx

@@ -1,39 +1,56 @@
 import type { LoaderFunctionArgs } from "@remix-run/server-runtime";
 import { redirect } from "@remix-run/server-runtime";
+import { prisma } from "~/db.server";
+import { redirectWithErrorMessage } from "~/models/message.server";
 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";
+import { commitSession, getSession } from "~/services/sessionStorage.server";
 
 export async function loader({ request }: LoaderFunctionArgs) {
-  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
-    });
+  const redirectTo = await getRedirectTo(request);
+
+  const auth = await authenticator.authenticate("email-link", request, {
+    failureRedirect: "/login/magic", // If auth fails, the failureRedirect will be thrown as a Response
+  });
+
+  // manually get the session
+  const session = await getSession(request.headers.get("cookie"));
+
+  const userRecord = await prisma.user.findFirst({
+    where: {
+      id: auth.userId,
+    },
+    select: {
+      id: true,
+      mfaEnabledAt: true,
+    },
+  });
+
+  if (!userRecord) {
+    return redirectWithErrorMessage(
+      "/login/magic",
+      request,
+      "Could not find your account. Please contact support."
+    );
+  }
+
+  if (userRecord.mfaEnabledAt) {
+    session.set("pending-mfa-user-id", userRecord.id);
+    session.set("pending-mfa-redirect-to", redirectTo ?? "/");
 
-    // 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");
+    return redirect("/login/mfa", {
+      headers: {
+        "Set-Cookie": await commitSession(session),
+      },
+    });
   }
+
+  // and store the user data
+  session.set(authenticator.sessionKey, auth);
+
+  return redirect(redirectTo ?? "/", {
+    headers: {
+      "Set-Cookie": await commitSession(session),
+    },
+  });
 }

apps/webapp/app/services/emailAuth.server.tsx

@@ -37,12 +37,6 @@ 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

apps/webapp/app/services/gitHubAuth.server.ts

@@ -41,12 +41,6 @@ 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,
         };

apps/webapp/app/services/mfa/multiFactorAuthentication.server.ts

@@ -15,15 +15,6 @@ const SecretSchema = z.object({
   secret: z.string(),
 });
 
-export class MfaRequiredError extends Error {
-  public readonly userId: string;
-
-  constructor(userId: string) {
-    super(`MFA is required for user ${userId}`);
-    this.userId = userId;
-  }
-}
-
 export class MultiFactorAuthenticationService {
   #prismaClient: PrismaClient;