-
Notifications
You must be signed in to change notification settings - Fork 48
fix(auth): harden device-auth poll (atomic consume + blocked re-check) #4328
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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))) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Low, Security: When this re-check trips, the function returns This is backend observability only; it does not change the response the device receives. Suggest logging before the return, matching the existing 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', | ||
|
|
||
There was a problem hiding this comment.
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, butpollDeviceAuthRequestreturnsdenied(line 211), which the route maps to 403 while a consumed or race-lost code returnsexpired/ 410. Two things worth a look:expired) does not match the response (denied), which can be confusing when debugging from the database.expiredhere (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.