Address PR comments

backend/src/server/routes/v2/mfa-router.ts

@@ -1,5 +1,7 @@
+import { FastifyReply, FastifyRequest } from "fastify";
 import { z } from "zod";
 
+import { TUsers } from "@app/db/schemas";
 import { getConfig } from "@app/lib/config/env";
 import { crypto } from "@app/lib/crypto";
 import { BadRequestError, NotFoundError } from "@app/lib/errors";
@@ -7,6 +9,49 @@ import { mfaRateLimit } from "@app/server/config/rateLimiter";
 import { addAuthOriginDomainCookie } from "@app/server/lib/cookie";
 import { AuthModeMfaJwtTokenPayload, AuthTokenType, MfaMethod } from "@app/services/auth/auth-type";
 
+const handleMfaVerification = async (
+  req: FastifyRequest & { mfa: { userId: string; orgId?: string; user: TUsers } },
+  res: FastifyReply,
+  server: FastifyZodProvider,
+  mfaToken: string,
+  mfaMethod: MfaMethod,
+  isRecoveryCode?: boolean
+) => {
+  const userAgent = req.headers["user-agent"];
+  const mfaJwtToken = req.headers.authorization?.replace("Bearer ", "");
+  if (!userAgent) throw new Error("user agent header is required");
+  if (!mfaJwtToken) throw new Error("authorization header is required");
+  const appCfg = getConfig();
+
+  const { user, token } = await server.services.login.verifyMfaToken({
+    userAgent,
+    mfaJwtToken,
+    ip: req.realIp,
+    userId: req.mfa.userId,
+    orgId: req.mfa.orgId,
+    mfaToken,
+    mfaMethod,
+    isRecoveryCode
+  });
+
+  void res.setCookie("jid", token.refresh, {
+    httpOnly: true,
+    path: "/",
+    sameSite: "strict",
+    secure: appCfg.HTTPS_ENABLED
+  });
+
+  addAuthOriginDomainCookie(res);
+
+  return {
+    ...user,
+    token: token.access,
+    protectedKey: user.protectedKey || null,
+    protectedKeyIV: user.protectedKeyIV || null,
+    protectedKeyTag: user.protectedKeyTag || null
+  };
+};
+
 export const registerMfaRouter = async (server: FastifyZodProvider) => {
   const cfg = getConfig();
 
@@ -109,38 +154,7 @@ export const registerMfaRouter = async (server: FastifyZodProvider) => {
       }
     },
     handler: async (req, res) => {
-      const userAgent = req.headers["user-agent"];
+      return handleMfaVerification(req, res, server, req.body.mfaToken, req.body.mfaMethod);
-      const mfaJwtToken = req.headers.authorization?.replace("Bearer ", "");
-      if (!userAgent) throw new Error("user agent header is required");
-      if (!mfaJwtToken) throw new Error("authorization header is required");
-      const appCfg = getConfig();
-
-      const { user, token } = await server.services.login.verifyMfaToken({
-        userAgent,
-        mfaJwtToken,
-        ip: req.realIp,
-        userId: req.mfa.userId,
-        orgId: req.mfa.orgId,
-        mfaToken: req.body.mfaToken,
-        mfaMethod: req.body.mfaMethod
-      });
-
-      void res.setCookie("jid", token.refresh, {
-        httpOnly: true,
-        path: "/",
-        sameSite: "strict",
-        secure: appCfg.HTTPS_ENABLED
-      });
-
-      addAuthOriginDomainCookie(res);
-
-      return {
-        ...user,
-        token: token.access,
-        protectedKey: user.protectedKey || null,
-        protectedKeyIV: user.protectedKeyIV || null,
-        protectedKeyTag: user.protectedKeyTag || null
-      };
     }
   });
 
@@ -169,39 +183,7 @@ export const registerMfaRouter = async (server: FastifyZodProvider) => {
       }
     },
     handler: async (req, res) => {
-      const userAgent = req.headers["user-agent"];
+      return handleMfaVerification(req, res, server, req.body.recoveryCode, MfaMethod.TOTP, true);
-      const mfaJwtToken = req.headers.authorization?.replace("Bearer ", "");
-      if (!userAgent) throw new Error("user agent header is required");
-      if (!mfaJwtToken) throw new Error("authorization header is required");
-      const appCfg = getConfig();
-
-      const { user, token } = await server.services.login.verifyMfaToken({
-        userAgent,
-        mfaJwtToken,
-        ip: req.realIp,
-        userId: req.mfa.userId,
-        orgId: req.mfa.orgId,
-        mfaToken: req.body.recoveryCode,
-        mfaMethod: MfaMethod.TOTP,
-        isRecoveryCode: true
-      });
-
-      void res.setCookie("jid", token.refresh, {
-        httpOnly: true,
-        path: "/",
-        sameSite: "strict",
-        secure: appCfg.HTTPS_ENABLED
-      });
-
-      addAuthOriginDomainCookie(res);
-
-      return {
-        ...user,
-        token: token.access,
-        protectedKey: user.protectedKey || null,
-        protectedKeyIV: user.protectedKeyIV || null,
-        protectedKeyTag: user.protectedKeyTag || null
-      };
     }
   });
 };

backend/src/services/auth/auth-login-service.ts

@@ -707,7 +707,7 @@ export const authLoginServiceFactory = ({
         } else {
           if (mfaToken.length !== 6) {
             throw new BadRequestError({
-              message: "Invalid TOTP code. Please use a valid recovery code."
+              message: "Please use a valid TOTP code."
             });
           }
           await totpService.verifyUserTotp({

frontend/src/components/mfa/TotpRegistration.tsx

@@ -35,7 +35,7 @@ const TotpRegistration = ({ onComplete, shouldCenterQr }: Props) => {
         type: "success"
       });
 
-      if (result.recoveryCodes) {
+      if (result.recoveryCodes && result.recoveryCodes.length > 0) {
         setRecoveryCodes(result.recoveryCodes);
         setShowRecoveryModal(true);
       } else if (onComplete) {