|
| 1 | +import { db } from '@sim/db' |
| 2 | +import { chat, user, verification, workflow, workspace } from '@sim/db/schema' |
| 3 | +import { readTestRedisUrl } from '@sim/db/testing/test-infrastructure' |
| 4 | +import { deleteWorkspaceFixture, insertWorkspaceFixture } from '@sim/db/testing/workspace-fixtures' |
| 5 | +import { createRouteContext } from '@sim/testing/helpers/http' |
| 6 | +import { generateId } from '@sim/utils/id' |
| 7 | +import { eq, like } from 'drizzle-orm' |
| 8 | +import Redis from 'ioredis' |
| 9 | +import { NextRequest } from 'next/server' |
| 10 | +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' |
| 11 | + |
| 12 | +const { storage, redisClient } = vi.hoisted(() => ({ |
| 13 | + storage: { backend: 'database' as 'database' | 'redis' }, |
| 14 | + redisClient: { current: undefined as Redis | undefined }, |
| 15 | +})) |
| 16 | +vi.mock('@/lib/core/storage', () => ({ getStorageMethod: () => storage.backend })) |
| 17 | +vi.mock('@/lib/core/config/redis', async (importOriginal) => ({ |
| 18 | + ...(await importOriginal<typeof import('@/lib/core/config/redis')>()), |
| 19 | + getRedisClient: () => redisClient.current, |
| 20 | +})) |
| 21 | + |
| 22 | +import { MAX_OTP_ATTEMPTS, storeOTP } from '@/lib/core/security/otp' |
| 23 | +import { PUT } from '@/app/api/chat/[identifier]/otp/route' |
| 24 | + |
| 25 | +const redisUrl = readTestRedisUrl() |
| 26 | +const ownerId = generateId() |
| 27 | +const workspaceId = generateId() |
| 28 | +const workflowId = generateId() |
| 29 | +const chatId = generateId() |
| 30 | +const identifier = `otp-${chatId}` |
| 31 | +const email = `${generateId()}@otp.invalid` |
| 32 | +const CODE = '246810' |
| 33 | +const NO_CODE = 'No verification code found, request a new one' |
| 34 | +const INVALID = 'Invalid verification code' |
| 35 | +const LOCKED = 'Too many failed attempts. Please request a new code.' |
| 36 | + |
| 37 | +async function verify(otp: string) { |
| 38 | + const request = new NextRequest(`http://localhost:3000/api/chat/${identifier}/otp`, { |
| 39 | + method: 'PUT', |
| 40 | + headers: { 'content-type': 'application/json' }, |
| 41 | + body: JSON.stringify({ email, otp }), |
| 42 | + }) |
| 43 | + const response = await PUT(request, createRouteContext({ identifier })) |
| 44 | + const body = (await response.json()) as { error?: string; message?: string } |
| 45 | + return { |
| 46 | + status: response.status, |
| 47 | + error: body.error ?? body.message, |
| 48 | + cookie: response.headers.get('set-cookie'), |
| 49 | + } |
| 50 | +} |
| 51 | + |
| 52 | +/** Distinct wrong six-digit guesses that never equal {@link CODE}. */ |
| 53 | +function wrongGuesses(count: number) { |
| 54 | + return Array.from({ length: count }, (_, index) => String(100000 + index)) |
| 55 | +} |
| 56 | + |
| 57 | +beforeAll(async () => { |
| 58 | + await db.insert(user).values({ |
| 59 | + id: ownerId, |
| 60 | + name: 'OTP fixture', |
| 61 | + email: `${ownerId}@otp.invalid`, |
| 62 | + emailVerified: true, |
| 63 | + createdAt: new Date(), |
| 64 | + updatedAt: new Date(), |
| 65 | + }) |
| 66 | + await insertWorkspaceFixture(db, { |
| 67 | + id: workspaceId, |
| 68 | + name: 'OTP fixture', |
| 69 | + ownerId, |
| 70 | + billedAccountUserId: ownerId, |
| 71 | + }) |
| 72 | + await db.insert(workflow).values({ |
| 73 | + id: workflowId, |
| 74 | + userId: ownerId, |
| 75 | + workspaceId, |
| 76 | + name: 'OTP fixture', |
| 77 | + lastSynced: new Date(), |
| 78 | + createdAt: new Date(), |
| 79 | + updatedAt: new Date(), |
| 80 | + }) |
| 81 | + await db.insert(chat).values({ |
| 82 | + id: chatId, |
| 83 | + workflowId, |
| 84 | + userId: ownerId, |
| 85 | + identifier, |
| 86 | + title: 'OTP fixture', |
| 87 | + authType: 'email', |
| 88 | + allowedEmails: [email], |
| 89 | + }) |
| 90 | +}) |
| 91 | + |
| 92 | +afterAll(async () => { |
| 93 | + await db.delete(verification).where(like(verification.identifier, `chat-otp:${chatId}:%`)) |
| 94 | + await deleteWorkspaceFixture(db, eq(workspace.id, workspaceId)) |
| 95 | + await db.delete(user).where(eq(user.id, ownerId)) |
| 96 | + await redisClient.current?.quit() |
| 97 | +}) |
| 98 | + |
| 99 | +describe.each(['database', 'redis'] as const)('%s chat OTP verification', (backend) => { |
| 100 | + describe.runIf(backend === 'database' || Boolean(redisUrl))('real storage', () => { |
| 101 | + beforeAll(async () => { |
| 102 | + if (backend === 'redis' && !redisClient.current) { |
| 103 | + redisClient.current = new Redis(redisUrl!, { lazyConnect: true, maxRetriesPerRequest: 0 }) |
| 104 | + await redisClient.current.connect() |
| 105 | + } |
| 106 | + }) |
| 107 | + beforeEach(async () => { |
| 108 | + storage.backend = backend |
| 109 | + await storeOTP('chat', chatId, email, CODE) |
| 110 | + }) |
| 111 | + |
| 112 | + it('compares at most the attempt limit of guesses however many arrive concurrently', async () => { |
| 113 | + const results = await Promise.all(wrongGuesses(30).map(verify)) |
| 114 | + |
| 115 | + const compared = results.filter((result) => result.error !== NO_CODE) |
| 116 | + expect(compared).toHaveLength(MAX_OTP_ATTEMPTS) |
| 117 | + expect(compared.filter((result) => result.error === INVALID)).toHaveLength( |
| 118 | + MAX_OTP_ATTEMPTS - 1 |
| 119 | + ) |
| 120 | + expect(compared.filter((result) => result.error === LOCKED)).toHaveLength(1) |
| 121 | + expect(await verify(CODE)).toMatchObject({ status: 400, error: NO_CODE, cookie: null }) |
| 122 | + }) |
| 123 | + |
| 124 | + it('accepts the correct code on the last remaining attempt and mints the auth cookie', async () => { |
| 125 | + for (const guess of wrongGuesses(MAX_OTP_ATTEMPTS - 1)) { |
| 126 | + expect(await verify(guess)).toMatchObject({ status: 400, error: INVALID }) |
| 127 | + } |
| 128 | + const accepted = await verify(CODE) |
| 129 | + expect(accepted.status).toBe(200) |
| 130 | + expect(accepted.cookie).toContain(`chat_auth_${chatId}=`) |
| 131 | + expect(await verify(CODE)).toMatchObject({ status: 400, error: NO_CODE }) |
| 132 | + }) |
| 133 | + |
| 134 | + it('keeps the code expiry when a failed attempt is recorded', async () => { |
| 135 | + const readExpiry = async () => |
| 136 | + backend === 'redis' |
| 137 | + ? redisClient.current?.ttl(`otp:${email}:${chatId}`) |
| 138 | + : ( |
| 139 | + await db |
| 140 | + .select({ expiresAt: verification.expiresAt }) |
| 141 | + .from(verification) |
| 142 | + .where(eq(verification.identifier, `chat-otp:${chatId}:${email}`)) |
| 143 | + )[0]?.expiresAt.getTime() |
| 144 | + const before = await readExpiry() |
| 145 | + expect(await verify('999999')).toMatchObject({ status: 400, error: INVALID }) |
| 146 | + const after = await readExpiry() |
| 147 | + if (backend === 'redis') expect(after).toBeGreaterThan(0) |
| 148 | + else expect(after).toBe(before) |
| 149 | + }) |
| 150 | + |
| 151 | + it.runIf(backend === 'database')('rejects an expired code', async () => { |
| 152 | + await db |
| 153 | + .update(verification) |
| 154 | + .set({ expiresAt: new Date(Date.now() - 1000) }) |
| 155 | + .where(eq(verification.identifier, `chat-otp:${chatId}:${email}`)) |
| 156 | + expect(await verify(CODE)).toMatchObject({ status: 400, error: NO_CODE, cookie: null }) |
| 157 | + }) |
| 158 | + }) |
| 159 | +}) |
0 commit comments