replace checkAuthToken with getUserInfoFromApiKey

charleslien · charleslien · commit 00c97eaccaaa · 2025-10-10T17:47:47.000-07:00
diff --git a/backend/src/api/usage.ts b/backend/src/api/usage.ts
@@ -2,6 +2,7 @@ import { getOrganizationUsageResponse } from '@codebuff/billing'
 import { INVALID_AUTH_TOKEN_MESSAGE } from '@codebuff/common/old-constants'
 import { z } from 'zod/v4'
 
+import { BACKEND_AGENT_RUNTIME_IMPL } from '../impl/agent-runtime'
 import { checkAuth } from '../util/check-auth'
 import { logger } from '../util/logger'
 import { getUserInfoFromApiKey } from '../websockets/auth'
@@ -31,10 +32,9 @@ async function usageHandler(
     const clientSessionId = `api-${fingerprintId}-${Date.now()}`
 
     const authResult = await checkAuth({
-      fingerprintId,
+      ...BACKEND_AGENT_RUNTIME_IMPL,
       authToken,
       clientSessionId,
-      logger,
     })
     if (authResult) {
       const errorMessage =
diff --git a/backend/src/util/check-auth.ts b/backend/src/util/check-auth.ts
@@ -1,29 +1,31 @@
-import db from '@codebuff/common/db'
-import * as schema from '@codebuff/common/db/schema'
 import { utils } from '@codebuff/internal'
-import { eq } from 'drizzle-orm'
 
 import { extractAuthTokenFromHeader } from './auth-helpers'
+import { getUserInfoFromApiKey } from '../websockets/auth'
 
 import type { ServerAction } from '@codebuff/common/actions'
+import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { Request, Response, NextFunction } from 'express'
 
 export const checkAuth = async (params: {
-  fingerprintId?: string
   authToken?: string
   clientSessionId: string
+  getUserInfoFromApiKey: GetUserInfoFromApiKeyFn
   logger: Logger
 }): Promise<void | ServerAction> => {
-  const { fingerprintId, authToken, clientSessionId, logger } = params
+  const { authToken, clientSessionId, getUserInfoFromApiKey, logger } = params
+
   // Use shared auth check functionality
-  const authResult = await utils.checkAuthToken({
-    fingerprintId,
-    authToken,
-  })
+  const authResult = authToken
+    ? await getUserInfoFromApiKey({
+        apiKey: authToken,
+        fields: ['id'],
+      })
+    : null
 
-  if (!authResult.success) {
-    const errorMessage = authResult.error?.message || 'Authentication failed'
+  if (!authResult) {
+    const errorMessage = 'Authentication failed'
     logger.error({ clientSessionId, error: errorMessage }, errorMessage)
     return {
       type: 'action-error',
@@ -43,65 +45,44 @@ export const checkAuth = async (params: {
 }
 
 // Express middleware for checking admin access
-export const checkAdmin = (logger: Logger) => async (
-  req: Request,
-  res: Response,
-  next: NextFunction,
-) => {
-  // Extract auth token from x-codebuff-api-key header
-  const authToken = extractAuthTokenFromHeader(req)
-  if (!authToken) {
-    return res.status(401).json({ error: 'Missing x-codebuff-api-key header' })
-  }
-
-  // Generate a client session ID for this request
-  const clientSessionId = `admin-relabel-${Date.now()}`
-
-  // Check authentication
-  const authResult = await checkAuth({
-    authToken,
-    clientSessionId,
-    logger,
-  })
+export const checkAdmin =
+  (logger: Logger) =>
+  async (req: Request, res: Response, next: NextFunction) => {
+    // Extract auth token from x-codebuff-api-key header
+    const authToken = extractAuthTokenFromHeader(req)
+    if (!authToken) {
+      return res
+        .status(401)
+        .json({ error: 'Missing x-codebuff-api-key header' })
+    }
 
-  if (authResult) {
-    // checkAuth returns an error action if auth fails
-    const errorMessage =
-      authResult.type === 'action-error'
-        ? authResult.message
-        : 'Authentication failed'
-    return res.status(401).json({ error: errorMessage })
-  }
+    // Generate a client session ID for this request
+    const clientSessionId = `admin-relabel-${Date.now()}`
 
-  // Get the user ID associated with this session token
-  const user = await db
-    .select({
-      id: schema.user.id,
-      email: schema.user.email,
+    // Check authentication
+    const user = await getUserInfoFromApiKey({
+      apiKey: authToken,
+      fields: ['id', 'email'],
     })
-    .from(schema.user)
-    .innerJoin(schema.session, eq(schema.user.id, schema.session.userId))
-    .where(eq(schema.session.sessionToken, authToken))
-    .then((users) => users[0])
 
-  if (!user) {
-    return res.status(401).json({ error: 'Invalid session' })
-  }
+    if (!user) {
+      return res.status(401).json({ error: 'Invalid session' })
+    }
 
-  // Check if user has admin access using shared utility
-  const adminUser = await utils.checkUserIsCodebuffAdmin(user.id)
-  if (!adminUser) {
-    logger.warn(
-      { userId: user.id, email: user.email, clientSessionId },
-      'Unauthorized access attempt to admin endpoint',
-    )
-    return res.status(403).json({ error: 'Forbidden' })
-  }
+    // Check if user has admin access using shared utility
+    const adminUser = await utils.checkUserIsCodebuffAdmin(user.id)
+    if (!adminUser) {
+      logger.warn(
+        { userId: user.id, email: user.email, clientSessionId },
+        'Unauthorized access attempt to admin endpoint',
+      )
+      return res.status(403).json({ error: 'Forbidden' })
+    }
 
-  // Store user info in request for handlers to use if needed
-  // req.user = adminUser // TODO: ensure type check passes
+    // Store user info in request for handlers to use if needed
+    // req.user = adminUser // TODO: ensure type check passes
 
-  // Auth passed and user is admin, proceed to next middleware
-  next()
-  return
-}
+    // Auth passed and user is admin, proceed to next middleware
+    next()
+    return
+  }
diff --git a/backend/src/websockets/middleware.ts b/backend/src/websockets/middleware.ts
@@ -60,25 +60,27 @@ export class WebSocketMiddleware {
   }
 
   use<T extends ClientAction['type']>(
-    callback: (params: {
-      action: ClientAction<T>
-      clientSessionId: string
-      ws: WebSocket
-      userInfo: { id: string } | null
-      logger: Logger
-    }) => Promise<void | ServerAction>,
+    callback: (
+      params: {
+        action: ClientAction<T>
+        clientSessionId: string
+        ws: WebSocket
+        userInfo: { id: string } | null
+      } & AgentRuntimeDeps,
+    ) => Promise<void | ServerAction>,
   ) {
     this.middlewares.push(callback as MiddlewareCallback)
   }
 
-  async execute(params: {
-    action: ClientAction
-    clientSessionId: string
-    ws: WebSocket
-    silent?: boolean
-    getUserInfoFromApiKey: GetUserInfoFromApiKeyFn
-    logger: Logger
-  }): Promise<boolean> {
+  async execute(
+    params: {
+      action: ClientAction
+      clientSessionId: string
+      ws: WebSocket
+      silent?: boolean
+      getUserInfoFromApiKey: GetUserInfoFromApiKeyFn
+    } & AgentRuntimeDeps,
+  ): Promise<boolean> {
     const { action, clientSessionId, ws, silent, logger } = params
 
     const userInfo =
@@ -91,11 +93,11 @@ export class WebSocketMiddleware {
 
     for (const middleware of this.middlewares) {
       const actionOrContinue = await middleware({
+        ...params,
         action,
         clientSessionId,
         ws,
         userInfo,
-        logger,
       })
       if (actionOrContinue) {
         logger.warn(
@@ -173,14 +175,13 @@ export class WebSocketMiddleware {
 
 export const protec = new WebSocketMiddleware(BACKEND_AGENT_RUNTIME_IMPL)
 
-protec.use(async ({ action, clientSessionId, logger }) =>
-  checkAuth({
-    fingerprintId: 'fingerprintId' in action ? action.fingerprintId : undefined,
+protec.use(async (params) => {
+  const { action } = params
+  return checkAuth({
+    ...params,
     authToken: 'authToken' in action ? action.authToken : undefined,
-    clientSessionId,
-    logger,
-  }),
-)
+  })
+})
 
 // Organization repository coverage detection middleware
 protec.use(async ({ action, userInfo, logger }) => {
diff --git a/packages/internal/src/knowledge.md b/packages/internal/src/knowledge.md
@@ -25,9 +25,3 @@ All environment variables are defined and validated in `env.ts`:
 - **Purpose**: Transactional emails (invitations, basic messages)
 - **Functions**: `sendOrganizationInvitationEmail`, `sendBasicEmail`, `sendSignupEventToLoops`
 - **Environment**: Requires `LOOPS_API_KEY`
-
-### Auth Utilities
-
-- **Purpose**: Admin user verification and session validation
-- **Functions**: `checkAuthToken`, `checkSessionIsAdmin`, `isCodebuffAdmin`
-- **Usage**: Used by admin routes and protected endpoints
diff --git a/packages/internal/src/utils/auth.ts b/packages/internal/src/utils/auth.ts
@@ -2,8 +2,6 @@ import db from '@codebuff/common/db'
 import * as schema from '@codebuff/common/db/schema'
 import { eq } from 'drizzle-orm'
 
-import { INVALID_AUTH_TOKEN_MESSAGE } from '@codebuff/common/old-constants'
-
 // List of admin user emails - single source of truth
 const CODEBUFF_ADMIN_USER_EMAILS = [
   'venkateshrameshkumar+1@gmail.com',
@@ -38,62 +36,6 @@ export interface AuthResult {
   }
 }
 
-/**
- * Check authentication based on auth token and fingerprint ID
- * Returns structured result instead of throwing or logging directly
- */
-export async function checkAuthToken({
-  fingerprintId,
-  authToken,
-}: {
-  fingerprintId?: string
-  authToken?: string
-}): Promise<AuthResult> {
-  if (!authToken) {
-    if (!fingerprintId) {
-      return {
-        success: false,
-        error: {
-          type: 'missing-credentials',
-          message: 'Auth token and fingerprint ID are missing',
-        },
-      }
-    }
-    return { success: true }
-  }
-
-  const user = await db
-    .select({
-      id: schema.user.id,
-      email: schema.user.email,
-      discord_id: schema.user.discord_id,
-    })
-    .from(schema.user)
-    .innerJoin(schema.session, eq(schema.user.id, schema.session.userId))
-    .where(eq(schema.session.sessionToken, authToken))
-    .then((users) => {
-      if (users.length === 1) {
-        return users[0]
-      }
-      return undefined
-    })
-
-  if (!user) {
-    return {
-      success: false,
-      error: {
-        type: 'invalid-token',
-        message: INVALID_AUTH_TOKEN_MESSAGE,
-      },
-    }
-  }
-
-  return {
-    success: true,
-    user,
-  }
-}
-
 /**
  * Check if the current user session corresponds to a Codebuff admin
  * Returns the admin user if authorized, null if not
diff --git a/web/src/app/api/agents/publish/route.ts b/web/src/app/api/agents/publish/route.ts
@@ -3,7 +3,6 @@ import * as schema from '@codebuff/common/db/schema'
 import { validateAgentsWithSpawnableAgents } from '@codebuff/common/templates/agent-validation'
 import { publishAgentsRequestSchema } from '@codebuff/common/types/api/agents/publish'
 import {
-  checkAuthToken,
   determineNextVersion,
   stringifyVersion,
   versionExists,
@@ -22,6 +21,7 @@ import { authOptions } from '../../auth/[...nextauth]/auth-options'
 import type { Version } from '@codebuff/internal'
 import type { NextRequest } from 'next/server'
 
+import { getUserInfoFromApiKey } from '@/db/user'
 import { logger } from '@/util/logger'
 
 async function getPublishedAgentIds(publisherId: string) {
@@ -96,9 +96,12 @@ export async function POST(request: NextRequest) {
     if (session?.user?.id) {
       userId = session.user.id
     } else if (authToken) {
-      const authResult = await checkAuthToken({ authToken })
-      if (authResult.success && authResult.user) {
-        userId = authResult.user.id
+      const user = await getUserInfoFromApiKey({
+        apiKey: authToken,
+        fields: ['id'],
+      })
+      if (user) {
+        userId = user.id
       }
     }
 
