[codex] Fix CLI OAuth login polling

freebuff/web/src/app/api/auth/cli/status/_db.ts

@@ -0,0 +1,44 @@
+import db from '@codebuff/internal/db'
+import * as schema from '@codebuff/internal/db/schema'
+import { and, eq, gt } from 'drizzle-orm'
+
+export interface LoginStatusUser {
+  id: string
+  email: string | null
+  name: string | null
+  authToken: string
+}
+
+export interface LoginStatusDb {
+  getCliSessionForAuth(
+    fingerprintId: string,
+    fingerprintHash: string,
+  ): Promise<LoginStatusUser | null>
+}
+
+export function createLoginStatusDb(): LoginStatusDb {
+  return {
+    getCliSessionForAuth: async (fingerprintId, fingerprintHash) => {
+      const users = await db
+        .select({
+          id: schema.user.id,
+          email: schema.user.email,
+          name: schema.user.name,
+          authToken: schema.session.sessionToken,
+        })
+        .from(schema.session)
+        .innerJoin(schema.user, eq(schema.session.userId, schema.user.id))
+        .where(
+          and(
+            eq(schema.session.fingerprint_id, fingerprintId),
+            eq(schema.session.cli_auth_hash, fingerprintHash),
+            eq(schema.session.type, 'cli'),
+            gt(schema.session.expires, new Date()),
+          ),
+        )
+        .limit(1)
+
+      return users[0] ?? null
+    },
+  }
+}

freebuff/web/src/app/api/auth/cli/status/_get.ts

@@ -0,0 +1,101 @@
+import { genAuthCode } from '@codebuff/common/util/credentials'
+import { NextResponse } from 'next/server'
+import { z } from 'zod/v4'
+
+import type { LoginStatusDb } from './_db'
+import type { Logger } from '@codebuff/common/types/contracts/logger'
+
+export type { LoginStatusDb } from './_db'
+
+interface GetLoginStatusDeps {
+  req: Request
+  db: LoginStatusDb
+  logger: Logger
+  secret: string
+  now?: () => number
+}
+
+const reqSchema = z.object({
+  fingerprintId: z.string(),
+  fingerprintHash: z.string(),
+  expiresAt: z.coerce.number().finite().int().positive(),
+})
+
+export async function getLoginStatus({
+  req,
+  db,
+  logger,
+  secret,
+  now = Date.now,
+}: GetLoginStatusDeps): Promise<NextResponse> {
+  const { searchParams } = new URL(req.url)
+  const result = reqSchema.safeParse({
+    fingerprintId: searchParams.get('fingerprintId'),
+    fingerprintHash: searchParams.get('fingerprintHash'),
+    expiresAt: searchParams.get('expiresAt'),
+  })
+  if (!result.success) {
+    return NextResponse.json(
+      { error: 'Invalid query parameters' },
+      { status: 400 },
+    )
+  }
+
+  const { fingerprintId, fingerprintHash, expiresAt } = result.data
+
+  if (now() > expiresAt) {
+    logger.info(
+      { fingerprintId, fingerprintHash, expiresAt },
+      'Auth code expired',
+    )
+    return NextResponse.json(
+      { error: 'Authentication failed' },
+      { status: 401 },
+    )
+  }
+
+  const expectedHash = genAuthCode(fingerprintId, expiresAt.toString(), secret)
+  if (fingerprintHash !== expectedHash) {
+    logger.info(
+      { fingerprintId, fingerprintHash, expectedHash },
+      'Invalid auth code',
+    )
+    return NextResponse.json(
+      { error: 'Authentication failed' },
+      { status: 401 },
+    )
+  }
+
+  try {
+    const user = await db.getCliSessionForAuth(fingerprintId, fingerprintHash)
+
+    if (!user) {
+      logger.info(
+        { fingerprintId, fingerprintHash },
+        'No active CLI session found for login auth code',
+      )
+      return NextResponse.json(
+        { error: 'Authentication failed' },
+        { status: 401 },
+      )
+    }
+
+    return NextResponse.json({
+      user: {
+        id: user.id,
+        name: user.name,
+        email: user.email,
+        authToken: user.authToken,
+        fingerprintId,
+        fingerprintHash,
+      },
+      message: 'Authentication successful!',
+    })
+  } catch (error) {
+    logger.error({ error }, 'Error checking login status')
+    return NextResponse.json(
+      { error: 'Internal server error' },
+      { status: 500 },
+    )
+  }
+}

freebuff/web/src/app/api/auth/cli/status/route.ts

@@ -1,114 +1,14 @@
-import { genAuthCode } from '@codebuff/common/util/credentials'
-import db from '@codebuff/internal/db'
-import * as schema from '@codebuff/internal/db/schema'
 import { env } from '@codebuff/internal/env'
-import { and, eq, gt, or, isNull } from 'drizzle-orm'
-import { NextResponse } from 'next/server'
-import { z } from 'zod/v4'
 
+import { createLoginStatusDb } from './_db'
+import { getLoginStatus } from './_get'
 import { logger } from '@/util/logger'
 
 export async function GET(req: Request) {
-  const { searchParams } = new URL(req.url)
-  const reqSchema = z.object({
-    fingerprintId: z.string(),
-    fingerprintHash: z.string(),
-    expiresAt: z.string().transform(Number),
+  return getLoginStatus({
+    req,
+    db: createLoginStatusDb(),
+    logger,
+    secret: env.NEXTAUTH_SECRET,
   })
-  const result = reqSchema.safeParse({
-    fingerprintId: searchParams.get('fingerprintId'),
-    fingerprintHash: searchParams.get('fingerprintHash'),
-    expiresAt: searchParams.get('expiresAt'),
-  })
-  if (!result.success) {
-    return NextResponse.json(
-      { error: 'Invalid query parameters' },
-      { status: 400 },
-    )
-  }
-
-  const { fingerprintId, fingerprintHash, expiresAt } = result.data
-
-  if (Date.now() > expiresAt) {
-    logger.info(
-      { fingerprintId, fingerprintHash, expiresAt },
-      'Auth code expired',
-    )
-    return NextResponse.json(
-      { error: 'Authentication failed' },
-      { status: 401 },
-    )
-  }
-
-  const expectedHash = genAuthCode(
-    fingerprintId,
-    expiresAt.toString(),
-    env.NEXTAUTH_SECRET,
-  )
-  if (fingerprintHash !== expectedHash) {
-    logger.info(
-      { fingerprintId, fingerprintHash, expectedHash },
-      'Invalid auth code',
-    )
-    return NextResponse.json(
-      { error: 'Authentication failed' },
-      { status: 401 },
-    )
-  }
-
-  try {
-    const users = await db
-      .select({
-        id: schema.user.id,
-        email: schema.user.email,
-        name: schema.user.name,
-        authToken: schema.session.sessionToken,
-      })
-      .from(schema.user)
-      .leftJoin(schema.session, eq(schema.user.id, schema.session.userId))
-      .leftJoin(
-        schema.fingerprint,
-        eq(schema.session.fingerprint_id, schema.fingerprint.id),
-      )
-      .where(
-        and(
-          eq(schema.session.fingerprint_id, fingerprintId),
-          or(
-            eq(schema.fingerprint.sig_hash, fingerprintHash),
-            isNull(schema.fingerprint.sig_hash),
-          ),
-          gt(schema.session.expires, new Date()),
-        ),
-      )
-
-    if (users.length === 0) {
-      logger.info(
-        { fingerprintId, fingerprintHash },
-        'No active session found or fingerprint claimed by another user',
-      )
-      return NextResponse.json(
-        { error: 'Authentication failed' },
-        { status: 401 },
-      )
-    }
-
-    const user = users[0]
-    return NextResponse.json({
-      user: {
-        id: user.id,
-        name: user.name,
-        email: user.email,
-        authToken: user.authToken,
-        fingerprintId,
-        fingerprintHash,
-      },
-      message: 'Authentication successful!',
-    })
-  } catch (error) {
-    logger.error({ error }, 'Error checking login status')
-    return NextResponse.json(
-      { error: 'Internal server error' },
-      { status: 500 },
-    )
-  }
 }

freebuff/web/src/app/onboard/_db.ts

@@ -1,7 +1,7 @@
 import { MAX_DATE } from '@codebuff/common/old-constants'
 import { db } from '@codebuff/internal/db'
 import * as schema from '@codebuff/internal/db/schema'
-import { and, eq, gt, isNull } from 'drizzle-orm'
+import { and, eq, gt, isNull, ne } from 'drizzle-orm'
 import { cookies } from 'next/headers'
 
 import { logger } from '@/util/logger'
@@ -12,22 +12,19 @@ type DbTransaction = Parameters<typeof db.transaction>[0] extends (
   ? T
   : never
 
-export async function checkReplayAttack(
+export async function hasCliSessionForAuthHash(
   fingerprintHash: string,
   userId: string,
 ): Promise<boolean> {
   const existing = await db
-    .select({ id: schema.user.id })
-    .from(schema.user)
-    .leftJoin(schema.session, eq(schema.user.id, schema.session.userId))
-    .leftJoin(
-      schema.fingerprint,
-      eq(schema.session.fingerprint_id, schema.fingerprint.id),
-    )
+    .select({ id: schema.session.userId })
+    .from(schema.session)
     .where(
       and(
-        eq(schema.fingerprint.sig_hash, fingerprintHash),
-        eq(schema.user.id, userId),
+        eq(schema.session.cli_auth_hash, fingerprintHash),
+        eq(schema.session.userId, userId),
+        eq(schema.session.type, 'cli'),
+        gt(schema.session.expires, new Date()),
       ),
     )
     .limit(1)
@@ -42,19 +39,19 @@ export async function checkFingerprintConflict(
   const existingSession = await db
     .select({
       userId: schema.session.userId,
-      expires: schema.session.expires,
     })
     .from(schema.session)
     .where(
       and(
         eq(schema.session.fingerprint_id, fingerprintId),
+        ne(schema.session.userId, userId),
         gt(schema.session.expires, new Date()),
       ),
     )
     .limit(1)
 
   const activeSession = existingSession[0]
-  if (activeSession && activeSession.userId !== userId) {
+  if (activeSession) {
     return { hasConflict: true, existingUserId: activeSession.userId }
   }
   return { hasConflict: false }
@@ -80,7 +77,7 @@ export async function createCliSession(
   return db.transaction(async (tx: DbTransaction) => {
     await tx
       .insert(schema.fingerprint)
-      .values({ sig_hash: fingerprintHash, id: fingerprintId })
+      .values({ id: fingerprintId })
       .onConflictDoNothing()
 
     const session = await tx
@@ -90,8 +87,10 @@ export async function createCliSession(
         userId,
         expires: MAX_DATE,
         fingerprint_id: fingerprintId,
+        cli_auth_hash: fingerprintHash,
         type: 'cli',
       })
+      .onConflictDoNothing()
       .returning({ userId: schema.session.userId })
 
     if (sessionToken) {

freebuff/web/src/app/onboard/_helpers.ts

@@ -20,5 +20,6 @@ export function validateAuthCode(
 }
 
 export function isAuthCodeExpired(expiresAt: string): boolean {
-  return expiresAt < Date.now().toString()
+  const expiresAtMs = Number(expiresAt)
+  return !Number.isFinite(expiresAtMs) || expiresAtMs < Date.now()
 }

freebuff/web/src/app/onboard/page.tsx

@@ -6,9 +6,9 @@ import { getServerSession } from 'next-auth'
 
 import {
   checkFingerprintConflict,
-  checkReplayAttack,
   createCliSession,
   getSessionTokenFromCookies,
+  hasCliSessionForAuthHash,
 } from './_db'
 import { isAuthCodeExpired, parseAuthCode, validateAuthCode } from './_helpers'
 import { authOptions } from '../api/auth/[...nextauth]/auth-options'
@@ -119,7 +119,7 @@ const Onboard = async ({ searchParams }: PageProps) => {
     )
   }
 
-  const isReplay = await checkReplayAttack(fingerprintHash, user.id)
+  const isReplay = await hasCliSessionForAuthHash(fingerprintHash, user.id)
   if (isReplay) {
     return (
       <StatusCard

packages/internal/src/db/migrations/0048_wide_blob.sql

@@ -0,0 +1 @@
+ALTER TABLE "session" ADD COLUMN "cli_auth_hash" text;

packages/internal/src/db/migrations/0049_loud_madame_masque.sql

@@ -0,0 +1 @@
+CREATE UNIQUE INDEX "session_cli_auth_code_idx" ON "session" USING btree ("fingerprint_id","cli_auth_hash") WHERE "session"."cli_auth_hash" IS NOT NULL;