Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions apps/web/src/lib/device-auth/device-auth.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand Down
37 changes: 27 additions & 10 deletions apps/web/src/lib/device-auth/device-auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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' })

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit, Functional: For the blocked/blacklisted case this row is set to expired, but pollDeviceAuthRequest returns denied (line 211), which the route maps to 403 while a consumed or race-lost code returns expired / 410. Two things worth a look:

  1. The persisted row state (expired) does not match the response (denied), which can be confusing when debugging from the database.
  2. More substantively, that distinct 403-vs-410 lets a polling client tell "you were blocked" apart from "code is gone." If you'd rather not signal the block to the device, returning expired here (still consumed, still terminal) would make the blocked case indistinguishable from a normal consumed code while the backend log from the comment above preserves the real reason for operators.

Not blocking; the current behavior is intentional per the description.

.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))) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Low, Security: When this re-check trips, the function returns { status: 'denied' } silently. The canonical auth path (validateUserAuthorization / authError in lib/user/server.ts:1077) emits an AUTH-FAIL 403 warn line for blocked and blacklisted users, so those denials are visible in backend logs. This branch has no equivalent, so a blocked or blacklisted user who approved before being blocked and then polls (exactly the abuse case this branch defends against) leaves no server-side trace.

This is backend observability only; it does not change the response the device receives. Suggest logging before the return, matching the existing AUTH-FAIL format and distinguishing the two reasons, e.g.:

if (user.blocked_reason || (await isUserBlacklistedByDomain(user))) {
  const reason = user.blocked_reason ? 'blocked' : 'blacklisted-domain';
  console.warn(`AUTH-FAIL 403 (${kiloUserId}): device-auth poll denied (${reason})`);
  return { status: 'denied' };
}

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',
Expand Down