diff --git a/prisma/migrations/22_add_2fa/migration.sql b/prisma/migrations/22_add_2fa/migration.sql index b3accc2df..04bb955e4 100644 --- a/prisma/migrations/22_add_2fa/migration.sql +++ b/prisma/migrations/22_add_2fa/migration.sql @@ -66,7 +66,4 @@ CREATE INDEX "two_factor_backup_code_user_id_idx" ON "two_factor_backup_code"("u CREATE UNIQUE INDEX "two_factor_otp_used_user_id_otp_key" ON "two_factor_otp_used"("user_id", "otp"); -- 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"); +CREATE UNIQUE INDEX "two_factor_rate_limit_user_id_key" ON "two_factor_rate_limit"("user_id"); \ No newline at end of file diff --git a/src/app/api/2fa/setup/confirm/route.test.ts b/src/app/api/2fa/setup/confirm/route.test.ts new file mode 100644 index 000000000..1dfefbfa2 --- /dev/null +++ b/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); +}); diff --git a/src/app/api/2fa/setup/initiate/route.test.ts b/src/app/api/2fa/setup/initiate/route.test.ts new file mode 100644 index 000000000..04d605e2b --- /dev/null +++ b/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); +}); diff --git a/src/app/api/2fa/verify/route.test.ts b/src/app/api/2fa/verify/route.test.ts new file mode 100644 index 000000000..c9c63f331 --- /dev/null +++ b/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); +}); diff --git a/src/app/api/2fa/verify/route.ts b/src/app/api/2fa/verify/route.ts index eea621b2b..1a1790d94 100644 --- a/src/app/api/2fa/verify/route.ts +++ b/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);