diff --git a/apps/web/src/lib/device-auth/device-auth.test.ts b/apps/web/src/lib/device-auth/device-auth.test.ts index 1b5ca6905d..983d28613b 100644 --- a/apps/web/src/lib/device-auth/device-auth.test.ts +++ b/apps/web/src/lib/device-auth/device-auth.test.ts @@ -253,6 +253,42 @@ describe('Device Auth', () => { expect(secondResult.token).toBeUndefined(); }); + test('concurrent polls mint at most one token from a single approval', async () => { + const { code } = await createDeviceAuthRequest({}); + await approveDeviceAuthRequest(code, testUserId); + + const results = await Promise.all([ + pollDeviceAuthRequest(code), + pollDeviceAuthRequest(code), + ]); + + const approved = results.filter(r => r.status === 'approved' && r.token); + const expired = results.filter(r => r.status === 'expired'); + expect(approved).toHaveLength(1); + expect(expired).toHaveLength(1); + }); + + test('does not mint a token when the user is blocked after approval', async () => { + const { code } = await createDeviceAuthRequest({}); + await approveDeviceAuthRequest(code, testUserId); + + // User becomes blocked between approval and the poll. + await db + .update(kilocode_users) + .set({ blocked_reason: 'blocked after approval' }) + .where(eq(kilocode_users.id, testUserId)); + + const result = await pollDeviceAuthRequest(code); + + expect(result.status).toBe('denied'); + expect(result.token).toBeUndefined(); + + // The approval is consumed, so retrying cannot mint a token either. + const retry = await pollDeviceAuthRequest(code); + expect(retry.status).toBe('expired'); + expect(retry.token).toBeUndefined(); + }); + test('normalizes responses - non-existent code returns expired', async () => { const result = await pollDeviceAuthRequest('FAKE-CODE'); expect(result.status).toBe('expired'); diff --git a/apps/web/src/lib/device-auth/device-auth.ts b/apps/web/src/lib/device-auth/device-auth.ts index 88578bad0f..b4471e72d0 100644 --- a/apps/web/src/lib/device-auth/device-auth.ts +++ b/apps/web/src/lib/device-auth/device-auth.ts @@ -3,6 +3,7 @@ import { db } from '@/lib/drizzle'; import { device_auth_requests, kilocode_users } from '@kilocode/db/schema'; import { eq, and, lt, sql } from 'drizzle-orm'; import { generateApiToken } from '@/lib/tokens'; +import { isUserBlacklistedByDomain } from '@/lib/user/server'; import { randomInt } from 'node:crypto'; const CODE_LENGTH = 8; @@ -176,26 +177,42 @@ export async function pollDeviceAuthRequest(code: string): Promise<{ return { status: request.status as 'pending' | 'denied' }; } - // For approved requests, fetch user and generate token + // Atomically consume the approval before minting: a single guarded + // compare-and-swap flips 'approved' -> 'expired' and returns the row only to + // the caller that won. This prevents concurrent polls from each minting a + // long-lived token from one approval. + const consumed = await db + .update(device_auth_requests) + .set({ status: 'expired' }) + .where(and(eq(device_auth_requests.code, code), eq(device_auth_requests.status, 'approved'))) + .returning({ kiloUserId: device_auth_requests.kilo_user_id }); + + if (consumed.length === 0 || !consumed[0].kiloUserId) { + // Lost the race to another poller (or the status changed) — normalize to + // expired so the code cannot be polled into a second token. + return { status: 'expired' }; + } + const kiloUserId = consumed[0].kiloUserId; + const [user] = await db .select() .from(kilocode_users) - .where(eq(kilocode_users.id, request.kilo_user_id)) + .where(eq(kilocode_users.id, kiloUserId)) .limit(1); if (!user) { throw new Error('User not found'); } - const token = generateApiToken(user, { deviceAuthRequestCode: code }); + // Re-check authorization at mint time. Authorization is verified when the + // user approves the request, but they may have been blocked or blacklisted + // between approval and this poll. The approval is already consumed above, so + // a now-blocked user cannot retry this code into a fresh long-lived token. + if (user.blocked_reason || (await isUserBlacklistedByDomain(user))) { + return { status: 'denied' }; + } - // Mark as consumed to enforce single-use - await db - .update(device_auth_requests) - .set({ - status: 'expired', - }) - .where(eq(device_auth_requests.code, code)); + const token = generateApiToken(user, { deviceAuthRequestCode: code }); return { status: 'approved',