From 1fa985cd7b5be857fa7d016b4af3ea79bf441d4f Mon Sep 17 00:00:00 2001 From: Matt Carey Date: Thu, 14 May 2026 20:51:59 +0100 Subject: [PATCH] Use KV for API token identity cache --- src/auth/api-token-mode.ts | 43 +++++++++++++++++++++---- src/auth/oauth-handler.ts | 26 +++------------- src/tests/auth/api-token-mode.test.ts | 36 ++++++++++++++++++--- src/tests/auth/oauth-handler.test.ts | 45 --------------------------- 4 files changed, 73 insertions(+), 77 deletions(-) diff --git a/src/auth/api-token-mode.ts b/src/auth/api-token-mode.ts index 639e6d8..804c3e6 100644 --- a/src/auth/api-token-mode.ts +++ b/src/auth/api-token-mode.ts @@ -1,13 +1,48 @@ +import { env as cloudflareEnv } from 'cloudflare:workers' + import { getUserAndAccounts } from './oauth-handler' import { OAuthError } from './workers-oauth-utils' -import type { AuthProps } from './types' +import type { AccountSchema, AuthProps, UserSchema } from './types' + +const env = cloudflareEnv as Env +const API_TOKEN_IDENTITY_CACHE_TTL_SECONDS = 2_592_000 + +type ApiTokenIdentity = { + user: UserSchema | null + accounts: AccountSchema[] +} async function hashApiToken(token: string): Promise { const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(token)) return [...new Uint8Array(digest)].map((byte) => byte.toString(16).padStart(2, '0')).join('') } +async function getCachedApiTokenIdentity(token: string): Promise { + const tokenHash = await hashApiToken(token) + const cacheKey = `api-token-identity:${tokenHash}` + try { + const cached = await env.OAUTH_KV.get(cacheKey, 'json') + if (cached) { + return cached + } + } catch (error) { + console.warn('api_token_identity_probe kv-cache read failed', error) + } + + const identity = await getUserAndAccounts(token, 'api_token_identity_probe') + + try { + await env.OAUTH_KV.put(cacheKey, JSON.stringify(identity), { + expirationTtl: API_TOKEN_IDENTITY_CACHE_TTL_SECONDS + }) + } catch (error) { + console.warn('api_token_identity_probe kv-cache write failed', error) + } + + return identity +} + /** * Check if the request contains a direct Cloudflare API token * (as opposed to an OAuth token issued by workers-oauth-provider) @@ -56,11 +91,7 @@ export async function handleApiTokenRequest( } try { - const { user, accounts } = await getUserAndAccounts( - token, - 'api_token_identity_probe', - await hashApiToken(token) - ) + const { user, accounts } = await getCachedApiTokenIdentity(token) // Account-scoped token if (!user) { diff --git a/src/auth/oauth-handler.ts b/src/auth/oauth-handler.ts index dc9410a..8598b18 100644 --- a/src/auth/oauth-handler.ts +++ b/src/auth/oauth-handler.ts @@ -213,25 +213,14 @@ function throwCombinedCloudflareApiError(userResp: Response, accountsResp: Respo async function fetchCloudflareProbes( accessToken: string, - caller = 'oauth_callback_identity_probe', - apiTokenCacheKeyHash?: string + caller = 'oauth_callback_identity_probe' ): Promise<[Response, Response]> { const headers = { Authorization: `Bearer ${accessToken}` } - const userUrl = `${env.CLOUDFLARE_API_BASE}/user` - const accountsUrl = `${env.CLOUDFLARE_API_BASE}/accounts` - const cacheCf = (path: 'user' | 'accounts') => - apiTokenCacheKeyHash - ? { - cacheEverything: true, - cacheKey: `${env.CLOUDFLARE_API_BASE}/mcp/api-token-identity/${apiTokenCacheKeyHash}/${path}`, - cacheTtl: 2_592_000 - } - : undefined try { return await Promise.all([ - fetchWithRetry(userUrl, { headers, cf: cacheCf('user') }, { caller }), - fetchWithRetry(accountsUrl, { headers, cf: cacheCf('accounts') }, { caller }) + fetchWithRetry(`${env.CLOUDFLARE_API_BASE}/user`, { headers }, { caller }), + fetchWithRetry(`${env.CLOUDFLARE_API_BASE}/accounts`, { headers }, { caller }) ]) } catch (error) { console.error('Cloudflare API request failed', error) @@ -244,17 +233,12 @@ async function fetchCloudflareProbes( */ export async function getUserAndAccounts( accessToken: string, - caller = 'oauth_callback_identity_probe', - apiTokenCacheKeyHash?: string + caller = 'oauth_callback_identity_probe' ): Promise<{ user: UserSchema | null accounts: AccountSchema[] }> { - const [userResp, accountsResp] = await fetchCloudflareProbes( - accessToken, - caller, - apiTokenCacheKeyHash - ) + const [userResp, accountsResp] = await fetchCloudflareProbes(accessToken, caller) // Check for upstream errors before parsing if (!userResp.ok && !accountsResp.ok) { diff --git a/src/tests/auth/api-token-mode.test.ts b/src/tests/auth/api-token-mode.test.ts index 71d2d43..f7a8ac4 100644 --- a/src/tests/auth/api-token-mode.test.ts +++ b/src/tests/auth/api-token-mode.test.ts @@ -1,4 +1,5 @@ -import { describe, it, expect, vi, beforeEach } from 'vitest' +import { env } from 'cloudflare:workers' +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' import { getUserAndAccounts } from '../../auth/oauth-handler' import { isDirectApiToken, @@ -28,6 +29,10 @@ beforeEach(() => { getUserAndAccountsMock.mockReset() }) +afterEach(() => { + vi.restoreAllMocks() +}) + describe('isDirectApiToken', () => { it('should return false for requests without Authorization header', () => { const request = mockRequest() @@ -151,21 +156,42 @@ describe('buildAuthProps', () => { describe('handleApiTokenRequest identity probe caching', () => { const token = 'api-token-123' + const tokenHash = '9bdb81d121b42d1c7819c816fa3cfbb6ee109726f9ed2475edb169374881d7b3' + const cacheKey = `api-token-identity:${tokenHash}` const user = { id: 'user-1', email: 'test@example.com' } const accounts = [{ id: 'acc-1', name: 'Account One' }] - it('passes a stable token hash for Cloudflare fetch cache keys', async () => { + it('stores API token identity lookups in KV by token hash', async () => { + const getSpy = vi.spyOn(env.OAUTH_KV, 'get').mockResolvedValue(null) + const putSpy = vi.spyOn(env.OAUTH_KV, 'put').mockResolvedValue(undefined) getUserAndAccountsMock.mockResolvedValue({ user, accounts }) const createMcpResponse = vi.fn().mockResolvedValue(new Response('ok')) const request = mockRequest(`Bearer ${token}`) await handleApiTokenRequest(request, createMcpResponse) + expect(getSpy).toHaveBeenCalledWith(cacheKey, 'json') expect(getUserAndAccountsMock).toHaveBeenCalledTimes(1) - expect(getUserAndAccountsMock).toHaveBeenCalledWith( + expect(getUserAndAccountsMock).toHaveBeenCalledWith(token, 'api_token_identity_probe') + expect(putSpy).toHaveBeenCalledWith(cacheKey, JSON.stringify({ user, accounts }), { + expirationTtl: 2_592_000 + }) + }) + + it('uses cached API token identity from KV', async () => { + vi.spyOn(env.OAUTH_KV, 'get').mockResolvedValue({ user, accounts }) + const putSpy = vi.spyOn(env.OAUTH_KV, 'put').mockResolvedValue(undefined) + const createMcpResponse = vi.fn().mockResolvedValue(new Response('ok')) + const request = mockRequest(`Bearer ${token}`) + + await handleApiTokenRequest(request, createMcpResponse) + + expect(getUserAndAccountsMock).not.toHaveBeenCalled() + expect(putSpy).not.toHaveBeenCalled() + expect(createMcpResponse).toHaveBeenCalledWith( token, - 'api_token_identity_probe', - '9bdb81d121b42d1c7819c816fa3cfbb6ee109726f9ed2475edb169374881d7b3' + undefined, + buildAuthProps(token, user, accounts) ) }) }) diff --git a/src/tests/auth/oauth-handler.test.ts b/src/tests/auth/oauth-handler.test.ts index bfa6e2f..b49acd0 100644 --- a/src/tests/auth/oauth-handler.test.ts +++ b/src/tests/auth/oauth-handler.test.ts @@ -128,51 +128,6 @@ describe('getUserAndAccounts', () => { }) }) - it('sets Cloudflare fetch cache options for API token identity probes', async () => { - const fetchMock = vi - .fn() - .mockResolvedValueOnce(new Response('Forbidden', { status: 403 })) - .mockResolvedValueOnce( - jsonResponse({ - success: true, - result: [{ id: 'acc-1', name: 'Primary Account' }] - }) - ) - - vi.stubGlobal('fetch', fetchMock) - - await expect( - getUserAndAccounts('test-token', 'api_token_identity_probe', 'token-hash') - ).resolves.toEqual({ - user: null, - accounts: [{ id: 'acc-1', name: 'Primary Account' }] - }) - - expect(fetchMock).toHaveBeenNthCalledWith( - 1, - 'https://api.cloudflare.com/client/v4/user', - expect.objectContaining({ - cf: { - cacheEverything: true, - cacheKey: 'https://api.cloudflare.com/client/v4/mcp/api-token-identity/token-hash/user', - cacheTtl: 2_592_000 - } - }) - ) - expect(fetchMock).toHaveBeenNthCalledWith( - 2, - 'https://api.cloudflare.com/client/v4/accounts', - expect.objectContaining({ - cf: { - cacheEverything: true, - cacheKey: - 'https://api.cloudflare.com/client/v4/mcp/api-token-identity/token-hash/accounts', - cacheTtl: 2_592_000 - } - }) - ) - }) - it('accepts user tokens when /accounts fails but /user succeeds', async () => { const fetchMock = vi .fn()