Improve 2FA flow

backend/src/server/routes/v1/user-router.ts

@@ -255,7 +255,9 @@ export const registerUserRouter = async (server: FastifyZodProvider) => {
         totp: z.string()
       }),
       response: {
-        200: z.object({})
+        200: z.object({
+          recoveryCodes: z.string().array()
+        })
       }
     },
     onRequest: verifyAuth([AuthMode.JWT], {

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

@@ -143,4 +143,65 @@ export const registerMfaRouter = async (server: FastifyZodProvider) => {
       };
     }
   });
+
+  server.route({
+    url: "/mfa/verify/recovery-code",
+    method: "POST",
+    config: {
+      rateLimit: mfaRateLimit
+    },
+    schema: {
+      body: z.object({
+        recoveryCode: z.string().trim().length(8, "Recovery code must be 8 characters")
+      }),
+      response: {
+        200: z.object({
+          encryptionVersion: z.number().default(1).nullable().optional(),
+          protectedKey: z.string().nullish(),
+          protectedKeyIV: z.string().nullish(),
+          protectedKeyTag: z.string().nullish(),
+          publicKey: z.string().nullish(),
+          encryptedPrivateKey: z.string().nullish(),
+          iv: z.string().nullish(),
+          tag: z.string().nullish(),
+          token: z.string()
+        })
+      }
+    },
+    handler: async (req, res) => {
+      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: 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

@@ -684,7 +684,8 @@ export const authLoginServiceFactory = ({
     mfaJwtToken,
     ip,
     userAgent,
-    orgId
+    orgId,
+    isRecoveryCode = false
   }: TVerifyMfaTokenDTO) => {
     const appCfg = getConfig();
     const user = await userDAL.findById(userId);
@@ -698,16 +699,21 @@ export const authLoginServiceFactory = ({
           code: mfaToken
         });
       } else if (mfaMethod === MfaMethod.TOTP) {
-        if (mfaToken.length === 6) {
-          await totpService.verifyUserTotp({
-            userId,
-            totp: mfaToken
-          });
-        } else {
+        if (isRecoveryCode) {
           await totpService.verifyWithUserRecoveryCode({
             userId,
             recoveryCode: mfaToken
           });
+        } else {
+          if (mfaToken.length !== 6) {
+            throw new BadRequestError({
+              message: "Invalid TOTP code. Please use a valid recovery code."
+            });
+          }
+          await totpService.verifyUserTotp({
+            userId,
+            totp: mfaToken
+          });
         }
       }
     } catch (err) {

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

@@ -24,6 +24,7 @@ export type TVerifyMfaTokenDTO = {
   ip: string;
   userAgent: string;
   orgId?: string;
+  isRecoveryCode?: boolean;
 };
 
 export type TOauthLoginDTO = {

backend/src/services/totp/totp-service.ts

@@ -131,15 +131,20 @@ export const totpServiceFactory = ({ totpConfigDAL, kmsService, userDAL }: TTotp
       secret
     });
 
-    if (isValid) {
-      await totpConfigDAL.updateById(totpConfig.id, {
-        isVerified: true
-      });
-    } else {
+    if (!isValid) {
       throw new BadRequestError({
         message: "Invalid TOTP token"
       });
     }
+
+    await totpConfigDAL.updateById(totpConfig.id, {
+      isVerified: true
+    });
+
+    const recoveryCodes = decryptWithRoot(totpConfig.encryptedRecoveryCodes).toString().split(",");
+    return {
+      recoveryCodes
+    };
   };
 
   const verifyUserTotp = async ({ userId, totp }: TVerifyUserTotpDTO) => {