add new test files and fix migration

prisma/migrations/22_add_2fa/migration.sql

@@ -67,6 +67,3 @@ CREATE UNIQUE INDEX "two_factor_otp_used_user_id_otp_key" ON "two_factor_otp_use
 
 -- CreateIndex
 CREATE UNIQUE INDEX "two_factor_rate_limit_user_id_key" ON "two_factor_rate_limit"("user_id");
-
--- CreateIndex
-CREATE INDEX "session_replay_visit_id_idx" ON "session_replay"("visit_id");

src/app/api/2fa/setup/confirm/route.test.ts

@@ -0,0 +1,153 @@
+import { beforeEach, expect, test, vi } from 'vitest';
+import { POST } from './route';
+
+const mocks = vi.hoisted(() => {
+  const tx = {
+    twoFactorAuth: {
+      update: vi.fn(),
+    },
+    twoFactorBackupCode: {
+      deleteMany: vi.fn(),
+      createMany: vi.fn(),
+    },
+  };
+
+  return {
+    parseRequest: vi.fn(),
+    findUnique: vi.fn(),
+    transaction: vi.fn(),
+    tx,
+    generateBackupCodes: vi.fn(),
+    decryptSecret: vi.fn(),
+    checkRateLimit: vi.fn(),
+    recordFailedAttempt: vi.fn(),
+    resetRateLimit: vi.fn(),
+    isOtpReplayed: vi.fn(),
+    markOtpUsed: vi.fn(),
+    verifyTotp: vi.fn(),
+  };
+});
+
+vi.mock('@/lib/request', () => ({
+  parseRequest: mocks.parseRequest,
+}));
+
+vi.mock('@/lib/prisma', () => ({
+  default: {
+    client: {
+      twoFactorAuth: {
+        findUnique: mocks.findUnique,
+      },
+    },
+    transaction: mocks.transaction,
+  },
+}));
+
+vi.mock('@/lib/two-factor/backup-codes', () => ({
+  generateBackupCodes: mocks.generateBackupCodes,
+}));
+
+vi.mock('@/lib/two-factor/crypto', () => ({
+  decryptSecret: mocks.decryptSecret,
+}));
+
+vi.mock('@/lib/two-factor/rate-limit', () => ({
+  checkRateLimit: mocks.checkRateLimit,
+  recordFailedAttempt: mocks.recordFailedAttempt,
+  resetRateLimit: mocks.resetRateLimit,
+}));
+
+vi.mock('@/lib/two-factor/replay-prevention', () => ({
+  isOtpReplayed: mocks.isOtpReplayed,
+  markOtpUsed: mocks.markOtpUsed,
+}));
+
+vi.mock('@/lib/two-factor/totp', () => ({
+  verifyTotp: mocks.verifyTotp,
+}));
+
+beforeEach(() => {
+  mocks.parseRequest.mockReset();
+  mocks.findUnique.mockReset();
+  mocks.transaction.mockReset();
+  mocks.tx.twoFactorAuth.update.mockReset();
+  mocks.tx.twoFactorBackupCode.deleteMany.mockReset();
+  mocks.tx.twoFactorBackupCode.createMany.mockReset();
+  mocks.generateBackupCodes.mockReset();
+  mocks.decryptSecret.mockReset();
+  mocks.checkRateLimit.mockReset();
+  mocks.recordFailedAttempt.mockReset();
+  mocks.resetRateLimit.mockReset();
+  mocks.isOtpReplayed.mockReset();
+  mocks.markOtpUsed.mockReset();
+  mocks.verifyTotp.mockReset();
+
+  mocks.parseRequest.mockResolvedValue({
+    auth: { user: { id: 'user-1' } },
+    body: { token: '123456' },
+    error: undefined,
+  });
+  mocks.findUnique.mockResolvedValue({ userId: 'user-1', isEnabled: false, secret: 'encrypted' });
+  mocks.transaction.mockImplementation(async callback => callback(mocks.tx));
+  mocks.generateBackupCodes.mockResolvedValue({
+    plaintext: ['code-1', 'code-2'],
+    hashed: ['hash-1', 'hash-2'],
+  });
+  mocks.decryptSecret.mockReturnValue('plain-secret');
+  mocks.checkRateLimit.mockResolvedValue({ allowed: true });
+  mocks.recordFailedAttempt.mockResolvedValue({ lockedUntil: undefined });
+  mocks.resetRateLimit.mockResolvedValue(undefined);
+  mocks.isOtpReplayed.mockResolvedValue(false);
+  mocks.markOtpUsed.mockResolvedValue(undefined);
+  mocks.verifyTotp.mockResolvedValue(true);
+  mocks.tx.twoFactorAuth.update.mockResolvedValue(undefined);
+  mocks.tx.twoFactorBackupCode.deleteMany.mockResolvedValue(undefined);
+  mocks.tx.twoFactorBackupCode.createMany.mockResolvedValue(undefined);
+});
+
+test('POST confirms setup, enables 2FA, stores backup codes, and resets the rate limit', async () => {
+  const response = await POST(
+    new Request('http://localhost/api/2fa/setup/confirm', { method: 'POST' }),
+  );
+
+  expect(mocks.checkRateLimit).toHaveBeenCalledWith('user-1');
+  expect(mocks.decryptSecret).toHaveBeenCalledWith('encrypted');
+  expect(mocks.verifyTotp).toHaveBeenCalledWith('123456', 'plain-secret');
+  expect(mocks.tx.twoFactorAuth.update).toHaveBeenCalledWith({
+    where: { userId: 'user-1' },
+    data: { isEnabled: true },
+  });
+  expect(mocks.tx.twoFactorBackupCode.deleteMany).toHaveBeenCalledWith({
+    where: { userId: 'user-1' },
+  });
+  expect(mocks.tx.twoFactorBackupCode.createMany).toHaveBeenCalledWith({
+    data: [
+      { userId: 'user-1', codeHash: 'hash-1' },
+      { userId: 'user-1', codeHash: 'hash-2' },
+    ],
+  });
+  expect(mocks.markOtpUsed).toHaveBeenCalledWith('user-1', '123456', mocks.tx);
+  expect(mocks.resetRateLimit).toHaveBeenCalledWith('user-1');
+  await expect(response.json()).resolves.toEqual({
+    backupCodes: ['code-1', 'code-2'],
+  });
+  expect(response.status).toBe(200);
+});
+
+test('POST records a failed attempt and skips writes when the token is invalid', async () => {
+  mocks.verifyTotp.mockResolvedValue(false);
+
+  const response = await POST(
+    new Request('http://localhost/api/2fa/setup/confirm', { method: 'POST' }),
+  );
+
+  expect(mocks.recordFailedAttempt).toHaveBeenCalledWith('user-1');
+  expect(mocks.transaction).not.toHaveBeenCalled();
+  expect(mocks.resetRateLimit).not.toHaveBeenCalled();
+  await expect(response.json()).resolves.toMatchObject({
+    error: {
+      code: 'two-factor-error-invalid-code',
+    },
+  });
+  expect(response.status).toBe(400);
+});

src/app/api/2fa/setup/initiate/route.test.ts

@@ -0,0 +1,99 @@
+import { beforeEach, expect, test, vi } from 'vitest';
+import { POST } from './route';
+
+const mocks = vi.hoisted(() => ({
+  parseRequest: vi.fn(),
+  getUser: vi.fn(),
+  findUnique: vi.fn(),
+  upsert: vi.fn(),
+  generateTotpSecret: vi.fn(),
+  encryptSecret: vi.fn(),
+  generateOtpAuthUri: vi.fn(),
+  generateQrCodeDataUrl: vi.fn(),
+}));
+
+vi.mock('@/lib/request', () => ({
+  parseRequest: mocks.parseRequest,
+}));
+
+vi.mock('@/queries/prisma/user', () => ({
+  getUser: mocks.getUser,
+}));
+
+vi.mock('@/lib/prisma', () => ({
+  default: {
+    client: {
+      twoFactorAuth: {
+        findUnique: mocks.findUnique,
+        upsert: mocks.upsert,
+      },
+    },
+  },
+}));
+
+vi.mock('@/lib/two-factor/crypto', () => ({
+  encryptSecret: mocks.encryptSecret,
+}));
+
+vi.mock('@/lib/two-factor/totp', () => ({
+  generateOtpAuthUri: mocks.generateOtpAuthUri,
+  generateQrCodeDataUrl: mocks.generateQrCodeDataUrl,
+  generateTotpSecret: mocks.generateTotpSecret,
+}));
+
+beforeEach(() => {
+  mocks.parseRequest.mockReset();
+  mocks.getUser.mockReset();
+  mocks.findUnique.mockReset();
+  mocks.upsert.mockReset();
+  mocks.generateTotpSecret.mockReset();
+  mocks.encryptSecret.mockReset();
+  mocks.generateOtpAuthUri.mockReset();
+  mocks.generateQrCodeDataUrl.mockReset();
+
+  mocks.parseRequest.mockResolvedValue({
+    auth: { user: { id: 'user-1' } },
+    error: undefined,
+  });
+  mocks.getUser.mockResolvedValue({ id: 'user-1', username: 'alice' });
+  mocks.findUnique.mockResolvedValue(null);
+  mocks.generateTotpSecret.mockReturnValue('plain-secret');
+  mocks.encryptSecret.mockReturnValue('encrypted-secret');
+  mocks.generateOtpAuthUri.mockReturnValue('otpauth://alice');
+  mocks.generateQrCodeDataUrl.mockResolvedValue('data:image/png;base64,qr');
+  mocks.upsert.mockResolvedValue(undefined);
+});
+
+test('POST creates a pending 2FA setup and returns the manual key and QR data', async () => {
+  const response = await POST(
+    new Request('http://localhost/api/2fa/setup/initiate', { method: 'POST' }),
+  );
+
+  expect(mocks.findUnique).toHaveBeenCalledWith({ where: { userId: 'user-1' } });
+  expect(mocks.upsert).toHaveBeenCalledWith({
+    where: { userId: 'user-1' },
+    update: { secret: 'encrypted-secret', isEnabled: false },
+    create: { userId: 'user-1', secret: 'encrypted-secret', isEnabled: false },
+  });
+  await expect(response.json()).resolves.toEqual({
+    manualKey: 'plain-secret',
+    qrCodeDataUrl: 'data:image/png;base64,qr',
+  });
+  expect(response.status).toBe(200);
+});
+
+test('POST rejects setup when 2FA is already enabled for the user', async () => {
+  mocks.findUnique.mockResolvedValue({ userId: 'user-1', isEnabled: true });
+
+  const response = await POST(
+    new Request('http://localhost/api/2fa/setup/initiate', { method: 'POST' }),
+  );
+
+  expect(mocks.upsert).not.toHaveBeenCalled();
+  await expect(response.json()).resolves.toMatchObject({
+    error: {
+      code: 'two-factor-error-already-enabled',
+    },
+  });
+  expect(response.status).toBe(400);
+});

src/app/api/2fa/verify/route.test.ts

@@ -0,0 +1,155 @@
+import { beforeEach, expect, test, vi } from 'vitest';
+
+const mocks = vi.hoisted(() => ({
+  getBearerToken: vi.fn(),
+  saveAuth: vi.fn(),
+  parseSecureToken: vi.fn(),
+  createSecureToken: vi.fn(),
+  getUser: vi.fn(),
+  getAllUserTeams: vi.fn(),
+  findTwoFactorAuth: vi.fn(),
+  findBackupCodes: vi.fn(),
+  updateBackupCodes: vi.fn(),
+  verifyBackupCode: vi.fn(),
+  decryptSecret: vi.fn(),
+  checkRateLimit: vi.fn(),
+  recordFailedAttempt: vi.fn(),
+  resetRateLimit: vi.fn(),
+  isOtpReplayed: vi.fn(),
+  markOtpUsed: vi.fn(),
+  verifyTotp: vi.fn(),
+  secret: vi.fn(),
+}));
+
+vi.mock('@/lib/auth', () => ({
+  getBearerToken: mocks.getBearerToken,
+  saveAuth: mocks.saveAuth,
+}));
+
+vi.mock('@/lib/crypto', () => ({
+  secret: mocks.secret,
+}));
+
+vi.mock('@/lib/jwt', () => ({
+  createSecureToken: mocks.createSecureToken,
+  parseSecureToken: mocks.parseSecureToken,
+}));
+
+vi.mock('@/queries/prisma', () => ({
+  getAllUserTeams: mocks.getAllUserTeams,
+  getUser: mocks.getUser,
+}));
+
+vi.mock('@/lib/prisma', () => ({
+  default: {
+    client: {
+      twoFactorAuth: {
+        findUnique: mocks.findTwoFactorAuth,
+      },
+      twoFactorBackupCode: {
+        findMany: mocks.findBackupCodes,
+        updateMany: mocks.updateBackupCodes,
+      },
+    },
+  },
+}));
+
+vi.mock('@/lib/two-factor/backup-codes', () => ({
+  verifyBackupCode: mocks.verifyBackupCode,
+}));
+
+vi.mock('@/lib/two-factor/crypto', () => ({
+  decryptSecret: mocks.decryptSecret,
+}));
+
+vi.mock('@/lib/two-factor/rate-limit', () => ({
+  checkRateLimit: mocks.checkRateLimit,
+  recordFailedAttempt: mocks.recordFailedAttempt,
+  resetRateLimit: mocks.resetRateLimit,
+}));
+
+vi.mock('@/lib/two-factor/replay-prevention', () => ({
+  isOtpReplayed: mocks.isOtpReplayed,
+  markOtpUsed: mocks.markOtpUsed,
+}));
+
+vi.mock('@/lib/two-factor/totp', () => ({
+  verifyTotp: mocks.verifyTotp,
+}));
+
+vi.mock('@/lib/redis', () => ({
+  default: {
+    enabled: false,
+  },
+}));
+
+import { POST } from './route';
+
+beforeEach(() => {
+  mocks.getBearerToken.mockReset();
+  mocks.saveAuth.mockReset();
+  mocks.parseSecureToken.mockReset();
+  mocks.createSecureToken.mockReset();
+  mocks.getUser.mockReset();
+  mocks.getAllUserTeams.mockReset();
+  mocks.findTwoFactorAuth.mockReset();
+  mocks.findBackupCodes.mockReset();
+  mocks.updateBackupCodes.mockReset();
+  mocks.verifyBackupCode.mockReset();
+  mocks.decryptSecret.mockReset();
+  mocks.checkRateLimit.mockReset();
+  mocks.recordFailedAttempt.mockReset();
+  mocks.resetRateLimit.mockReset();
+  mocks.isOtpReplayed.mockReset();
+  mocks.markOtpUsed.mockReset();
+  mocks.verifyTotp.mockReset();
+  mocks.secret.mockReset();
+
+  mocks.getBearerToken.mockReturnValue('partial-token');
+  mocks.secret.mockReturnValue('app-secret');
+  mocks.parseSecureToken.mockReturnValue({ type: 'partial-auth', userId: 'user-1' });
+  mocks.getUser.mockResolvedValue({
+    id: 'user-1',
+    username: 'alice',
+    role: 'admin',
+    createdAt: new Date('2026-07-23T00:00:00.000Z'),
+  });
+  mocks.getAllUserTeams.mockResolvedValue([]);
+  mocks.findTwoFactorAuth.mockResolvedValue({ userId: 'user-1', isEnabled: true, secret: 'encrypted' });
+  mocks.createSecureToken.mockReturnValue('full-auth-token');
+  mocks.decryptSecret.mockReturnValue('plain-secret');
+  mocks.checkRateLimit.mockResolvedValue({ allowed: true });
+  mocks.recordFailedAttempt.mockResolvedValue({ lockedUntil: undefined });
+  mocks.resetRateLimit.mockResolvedValue(undefined);
+  mocks.isOtpReplayed.mockResolvedValue(false);
+  mocks.markOtpUsed.mockResolvedValue(undefined);
+  mocks.verifyTotp.mockResolvedValue(true);
+  mocks.findBackupCodes.mockResolvedValue([]);
+  mocks.updateBackupCodes.mockResolvedValue({ count: 0 });
+  mocks.verifyBackupCode.mockResolvedValue(null);
+});
+
+test('POST accepts a token-only payload and completes 2FA verification', async () => {
+  const response = await POST(
+    new Request('http://localhost/api/2fa/verify', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: 'Bearer partial-token',
+      },
+      body: JSON.stringify({ token: '123456' }),
+    }),
+  );
+
+  expect(mocks.verifyTotp).toHaveBeenCalledWith('123456', 'plain-secret');
+  expect(mocks.markOtpUsed).toHaveBeenCalledWith('user-1', '123456');
+  expect(mocks.resetRateLimit).toHaveBeenCalledWith('user-1');
+  await expect(response.json()).resolves.toMatchObject({
+    token: 'full-auth-token',
+    user: {
+      id: 'user-1',
+      username: 'alice',
+    },
+  });
+  expect(response.status).toBe(200);
+});

src/app/api/2fa/verify/route.ts

@@ -16,8 +16,8 @@ import redis from '@/lib/redis';
 
 export async function POST(request: Request) {
   const schema = z.union([
-    z.object({ token: z.string().length(6), backupCode: z.undefined() }),
-    z.object({ backupCode: z.string().min(1), token: z.undefined() }),
+    z.object({ token: z.string().length(6) }).strict(),
+    z.object({ backupCode: z.string().min(1) }).strict(),
   ]);
 
   const rawToken = getBearerToken(request);