From e4c351d861c3425fad7e867f0975b823702ef213 Mon Sep 17 00:00:00 2001 From: Samuel Colvin Date: Mon, 20 Oct 2025 21:48:35 -0700 Subject: [PATCH 1/4] double cache api keys --- gateway/src/auth.ts | 64 ++++++++++++++++++--------- gateway/src/gateway.ts | 8 ++-- gateway/src/index.ts | 1 + gateway/test/auth.spec.ts | 93 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 142 insertions(+), 24 deletions(-) create mode 100644 gateway/test/auth.spec.ts diff --git a/gateway/src/auth.ts b/gateway/src/auth.ts index 8cac620..6920dfe 100644 --- a/gateway/src/auth.ts +++ b/gateway/src/auth.ts @@ -1,11 +1,14 @@ import type { GatewayOptions } from '.' import type { ApiKeyInfo } from './types' -import { ResponseError } from './utils' +import { ResponseError, runAfter } from './utils' -const CACHE_VERSION = 1 -const CACHE_TTL = 86400 +const CACHE_TTL = 86400 * 30 -export async function apiKeyAuth(request: Request, options: GatewayOptions): Promise { +export async function apiKeyAuth( + request: Request, + ctx: ExecutionContext, + options: GatewayOptions, +): Promise { const authHeader = request.headers.get('authorization') let key: string @@ -23,26 +26,47 @@ export async function apiKeyAuth(request: Request, options: GatewayOptions): Pro throw new ResponseError(401, 'Unauthorized - Key too long') } - const cacheKey = apiKeyCacheKey(key, options) - const cacheResult = await options.kv.getWithMetadata(cacheKey, { type: 'json' }) + const cacheKey = apiKeyCacheKey(key, options.kvVersion) + const cacheResult = await options.kv.getWithMetadata(cacheKey, { type: 'json' }) - let apiKey: ApiKeyInfo | null - if (cacheResult && cacheResult.metadata === CACHE_VERSION && cacheResult.value) { - apiKey = cacheResult.value - } else { - apiKey = await options.keysDb.getApiKey(key) - if (!apiKey) { - throw new ResponseError(401, 'Unauthorized - Key not found') + if (cacheResult?.value) { + const apiKey = cacheResult.value + const projectState = await options.kv.get(projectStateCacheKey(apiKey.project, options.kvVersion)) + // we only return a cache match if the org state is the same, so updating the org state invalidates the cache + if (projectState === null && projectState === cacheResult.metadata) { + return apiKey } - await options.kv.put(cacheKey, JSON.stringify(apiKey), { metadata: CACHE_VERSION, expirationTtl: CACHE_TTL }) } - // check all key validity in gateway.ts - return apiKey + + const apiKey = await options.keysDb.getApiKey(key) + if (apiKey) { + runAfter(ctx, 'setApiKeyCache', setApiKeyCache(apiKey, options)) + return apiKey + } + throw new ResponseError(401, 'Unauthorized - Key not found') +} + +export async function setApiKeyCache( + apiKey: ApiKeyInfo, + options: Pick, + expirationTtl?: number, +) { + const projectState = await options.kv.get(projectStateCacheKey(apiKey.org, options.kvVersion)) + + await options.kv.put(apiKeyCacheKey(apiKey.key, options.kvVersion), JSON.stringify(apiKey), { + metadata: projectState, + expirationTtl: expirationTtl || CACHE_TTL, + }) +} + +export async function deleteApiKeyCache(apiKey: ApiKeyInfo, options: Pick) { + await options.kv.delete(apiKeyCacheKey(apiKey.key, options.kvVersion)) } -export async function disableApiKeyAuth(apiKey: ApiKeyInfo, options: GatewayOptions, expirationTtl?: number) { - const cacheKey = apiKeyCacheKey(apiKey.key, options) - await options.kv.put(cacheKey, JSON.stringify(apiKey), { metadata: CACHE_VERSION, expirationTtl }) +export async function changeProjectState(project: number, options: Pick) { + const cacheKey = projectStateCacheKey(project, options.kvVersion) + await options.kv.put(cacheKey, crypto.randomUUID(), { expirationTtl: CACHE_TTL }) } -const apiKeyCacheKey = (key: string, options: GatewayOptions) => `apiKeyAuth:${options.kvVersion}:${key}` +const apiKeyCacheKey = (key: string, kvVersion: string) => `apiKeyAuth:${kvVersion}:${key}` +const projectStateCacheKey = (project: number, kvVersion: string) => `projectState:${kvVersion}:${project}` diff --git a/gateway/src/gateway.ts b/gateway/src/gateway.ts index 6ce1173..f73a36a 100644 --- a/gateway/src/gateway.ts +++ b/gateway/src/gateway.ts @@ -1,5 +1,5 @@ import type { GatewayOptions } from '.' -import { apiKeyAuth, disableApiKeyAuth } from './auth' +import { apiKeyAuth, setApiKeyCache } from './auth' import { type ExceededScope, endOfMonth, endOfWeek, type SpendScope, scopeIntervals } from './db' import { OtelTrace } from './otel' import { genAiOtelAttributes } from './otel/attributes' @@ -23,7 +23,7 @@ export async function gateway( return textResponse(400, `Invalid provider '${provider}', should be one of ${providerIdArray.join(', ')}`) } - const apiKeyInfo = await apiKeyAuth(request, options) + const apiKeyInfo = await apiKeyAuth(request, ctx, options) if (apiKeyInfo.status !== 'active') { return textResponse(403, `Unauthorized - Key ${apiKeyInfo.status}`) @@ -72,7 +72,7 @@ export async function gateway( } else if ('error' in result) { const { error, disableKey } = result if (disableKey) { - runAfter(ctx, 'disableApiKey', blockApiKey(apiKeyInfo, options, 'Invalid request')) + runAfter(ctx, 'blockApiKey', blockApiKey(apiKeyInfo, options, 'Invalid request')) response = textResponse(400, `${error}, API key disabled`) } else { response = textResponse(400, error) @@ -97,7 +97,7 @@ export async function disableApiKey( expirationTtl?: number, ): Promise { apiKey.status = newStatus - await disableApiKeyAuth(apiKey, options, expirationTtl) + await setApiKeyCache(apiKey, options, expirationTtl) await options.keysDb.disableKey(apiKey.id, reason, newStatus, expirationTtl) } diff --git a/gateway/src/index.ts b/gateway/src/index.ts index 13c2b26..6798c26 100644 --- a/gateway/src/index.ts +++ b/gateway/src/index.ts @@ -21,6 +21,7 @@ import type { DefaultProviderProxy, Middleware, Next } from './providers/default import type { SubFetch } from './types' import { ctHeader, ResponseError, response405, textResponse } from './utils' +export { changeProjectState as setProjectState, deleteApiKeyCache, setApiKeyCache } from './auth' export type { DefaultProviderProxy, Middleware, Next } export * from './db' export * from './types' diff --git a/gateway/test/auth.spec.ts b/gateway/test/auth.spec.ts new file mode 100644 index 0000000..4ea7228 --- /dev/null +++ b/gateway/test/auth.spec.ts @@ -0,0 +1,93 @@ +/** biome-ignore-all lint/suspicious/useAwait: don't care in tests */ +import { createExecutionContext, env, waitOnExecutionContext } from 'cloudflare:test' +import type { KeysDb } from '@pydantic/ai-gateway' +import { describe, expect } from 'vitest' +import { apiKeyAuth, changeProjectState } from '../src/auth' +import type { ApiKeyInfo } from '../src/types' +import { test } from './setup' +import { buildGatewayEnv, IDS } from './worker' + +class CountingKeysDb implements KeysDb { + callCount = 0 + private wrapped: KeysDb + + constructor(wrapped: KeysDb) { + this.wrapped = wrapped + } + + async getApiKey(key: string): Promise { + this.callCount++ + return this.wrapped.getApiKey(key) + } + + async disableKey(id: number, reason: string, newStatus: string, expirationTtl?: number): Promise { + return this.wrapped.disableKey(id, reason, newStatus, expirationTtl) + } +} + +describe('apiKeyAuth cache invalidation', () => { + test('caches api key and returns cached value', async () => { + const ctx = createExecutionContext() + const baseOptions = buildGatewayEnv(env, [], fetch) + const countingDb = new CountingKeysDb(baseOptions.keysDb) + const options = { ...baseOptions, keysDb: countingDb } + + const request = new Request('https://example.com', { headers: { Authorization: 'healthy' } }) + + // First call should fetch from DB + const apiKey1 = await apiKeyAuth(request, ctx, options) + expect(apiKey1.key).toBe('healthy') + // Wait for cache to be set (it's set asynchronously via runAfter) + await waitOnExecutionContext(ctx) + expect(countingDb.callCount).toBe(1) + + // Verify cache was set + const cached = await env.KV.get('apiKeyAuth:test:healthy') + expect(cached).toBeTypeOf('string') + + // Second call should use cache, not hit DB + const ctx2 = createExecutionContext() + const apiKey2 = await apiKeyAuth(request, ctx2, options) + expect(apiKey2.key).toBe('healthy') + + expect(countingDb.callCount).toBe(1) + }) + + test('invalidates cache when project state changes', async () => { + const ctx = createExecutionContext() + const baseOptions = buildGatewayEnv(env, [], fetch) + const countingDb = new CountingKeysDb(baseOptions.keysDb) + const options = { ...baseOptions, keysDb: countingDb } + + const request = new Request('https://example.com', { headers: { Authorization: 'healthy' } }) + + // First call - fetch from DB and cache + await apiKeyAuth(request, ctx, options) + await waitOnExecutionContext(ctx) + expect(countingDb.callCount).toBe(1) + + const cached1 = await env.KV.getWithMetadata('apiKeyAuth:test:healthy') + expect(cached1.value).not.toBeNull() + expect(cached1.metadata).toBeNull() + + // Second call - should use cache, not hit DB + const ctx2 = createExecutionContext() + await apiKeyAuth(request, ctx2, options) + await waitOnExecutionContext(ctx2) + expect(countingDb.callCount).toBe(1) + + // Change project state - this invalidates the cache + await changeProjectState(IDS.projectDefault, options) + + const projectState = await env.KV.get(`projectState:test:${IDS.projectDefault}`) + expect(projectState).not.toBeNull() + + // Third call - cache is invalidated, should hit DB again + const ctx3 = createExecutionContext() + const apiKey3 = await apiKeyAuth(request, ctx3, options) + expect(apiKey3.key).toBe('healthy') + await waitOnExecutionContext(ctx3) + + expect(countingDb.callCount).toBe(2) + }) +}) From 41cd199cac8ef7c484d43f1baf61faab0f915449 Mon Sep 17 00:00:00 2001 From: Daniel Cruz Date: Tue, 21 Oct 2025 20:55:05 -0600 Subject: [PATCH 2/4] fix: project cache --- gateway/src/auth.ts | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/gateway/src/auth.ts b/gateway/src/auth.ts index 6920dfe..bd1b96f 100644 --- a/gateway/src/auth.ts +++ b/gateway/src/auth.ts @@ -32,8 +32,9 @@ export async function apiKeyAuth( if (cacheResult?.value) { const apiKey = cacheResult.value const projectState = await options.kv.get(projectStateCacheKey(apiKey.project, options.kvVersion)) - // we only return a cache match if the org state is the same, so updating the org state invalidates the cache - if (projectState === null && projectState === cacheResult.metadata) { + // we only return a cache match if the project state is the same, so updating the project state invalidates the cache + // projectState is null if we have never invalidated the cache which will only be true for the first request after a deployment + if (projectState === null || projectState === cacheResult.metadata) { return apiKey } } @@ -51,7 +52,7 @@ export async function setApiKeyCache( options: Pick, expirationTtl?: number, ) { - const projectState = await options.kv.get(projectStateCacheKey(apiKey.org, options.kvVersion)) + const projectState = await options.kv.get(projectStateCacheKey(apiKey.project, options.kvVersion)) await options.kv.put(apiKeyCacheKey(apiKey.key, options.kvVersion), JSON.stringify(apiKey), { metadata: projectState, From c0981107c82c5d884e0044bc2f919ba22fb8d96b Mon Sep 17 00:00:00 2001 From: Daniel Cruz Date: Tue, 21 Oct 2025 21:14:55 -0600 Subject: [PATCH 3/4] fix: fallback --- gateway/src/auth.ts | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/gateway/src/auth.ts b/gateway/src/auth.ts index bd1b96f..6992c2c 100644 --- a/gateway/src/auth.ts +++ b/gateway/src/auth.ts @@ -56,7 +56,10 @@ export async function setApiKeyCache( await options.kv.put(apiKeyCacheKey(apiKey.key, options.kvVersion), JSON.stringify(apiKey), { metadata: projectState, - expirationTtl: expirationTtl || CACHE_TTL, + // Note: 0 is a valid expirationTtl (for immediate cache expiry if, e.g., the user hits a limit at the end of an interval). + // Do not use logical OR (||) for a fallback, as it would treat 0 as false and incorrectly default to CACHE_TTL, + // potentially locking out the user much longer than intended. + expirationTtl: expirationTtl ?? CACHE_TTL, }) } From 7ce75fca1b0d8703af846cdf7065504cad807030 Mon Sep 17 00:00:00 2001 From: Daniel Cruz Date: Tue, 21 Oct 2025 22:44:05 -0600 Subject: [PATCH 4/4] refactor: narrow delete api key types --- gateway/src/auth.ts | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/gateway/src/auth.ts b/gateway/src/auth.ts index 6992c2c..193cf49 100644 --- a/gateway/src/auth.ts +++ b/gateway/src/auth.ts @@ -63,7 +63,10 @@ export async function setApiKeyCache( }) } -export async function deleteApiKeyCache(apiKey: ApiKeyInfo, options: Pick) { +export async function deleteApiKeyCache( + apiKey: Pick, + options: Pick, +) { await options.kv.delete(apiKeyCacheKey(apiKey.key, options.kvVersion)) }