From 67449d76efbf8ac7d0038f47abf2eac1c8eb98a6 Mon Sep 17 00:00:00 2001 From: Mike Cao Date: Sat, 5 Sep 2026 14:09:33 -0700 Subject: [PATCH] Harden OAuth metadata fetching and make token rotation transactional --- packages/mcp/README.md | 3 + src/app/api/oauth/register/route.ts | 39 +++++-- src/lib/oauth/client.ts | 65 ++++++++++- src/lib/oauth/token.test.ts | 17 ++- src/lib/oauth/token.ts | 172 +++++++++++++++++----------- src/queries/prisma/oauth.ts | 53 ++++++--- 6 files changed, 245 insertions(+), 104 deletions(-) diff --git a/packages/mcp/README.md b/packages/mcp/README.md index 01a46d37f..f053170dd 100644 --- a/packages/mcp/README.md +++ b/packages/mcp/README.md @@ -54,6 +54,9 @@ The client discovers `/.well-known/oauth-protected-resource/mcp`, sends you to U and approve access (`websites:read`, `analytics:read`), and receives an OAuth token. No API key is pasted anywhere. Clients register via Client ID Metadata Documents; dynamic client registration is available for compatibility (`OAUTH_DISABLE_DCR=1` turns it off). +Registration is limited to 20 attempts per hour per IP when Redis is configured. Without Redis, +the limit is shared by all clients in each server process and resets when that process restarts. +Use Redis for a shared limit across multiple processes or serverless instances. Self-hosted instances can alternatively pass an API key (`umami_…`) as the bearer token. diff --git a/src/app/api/oauth/register/route.ts b/src/app/api/oauth/register/route.ts index b66f9d85e..e14297397 100644 --- a/src/app/api/oauth/register/route.ts +++ b/src/app/api/oauth/register/route.ts @@ -11,6 +11,24 @@ import { createOauthClient } from '@/queries/prisma/oauth'; import { clientRegistrationRequestSchema } from '../schema'; const REGISTRATIONS_PER_HOUR = 20; +const REGISTRATION_WINDOW_MS = 60 * 60 * 1000; +let registrationWindow = { count: 0, expiresAt: 0 }; + +function isLocalRegistrationLimited() { + const now = Date.now(); + + if (now >= registrationWindow.expiresAt) { + registrationWindow = { count: 0, expiresAt: now + REGISTRATION_WINDOW_MS }; + } + + // A process-wide cap keeps both registrations and limiter memory bounded without Redis. + if (registrationWindow.count >= REGISTRATIONS_PER_HOUR) { + return true; + } + + registrationWindow.count++; + return false; +} export function OPTIONS() { return corsPreflight(); @@ -37,17 +55,18 @@ export async function POST(request: Request) { } try { - if (redis.enabled) { - const ip = getIpAddress(request.headers) ?? 'unknown'; - const allowed = await redis.client.rateLimit( - `oauth:register:${ip}`, - REGISTRATIONS_PER_HOUR, - 60 * 60, - ); + const limited = redis.enabled + ? await redis.client.rateLimit( + `oauth:register:${getIpAddress(request.headers) ?? 'unknown'}`, + REGISTRATIONS_PER_HOUR + 1, + REGISTRATION_WINDOW_MS / 1000, + ) + : isLocalRegistrationLimited(); - if (!allowed) { - throw new OAuthError('temporarily_unavailable', 'Too many registrations. Try again later.'); - } + if (limited) { + throw new OAuthError('temporarily_unavailable', 'Too many registrations. Try again later.', { + status: 429, + }); } if (!body.redirect_uris.every(isAcceptableRedirectUri)) { diff --git a/src/lib/oauth/client.ts b/src/lib/oauth/client.ts index 0d08813b8..e4a6b62c3 100644 --- a/src/lib/oauth/client.ts +++ b/src/lib/oauth/client.ts @@ -1,4 +1,8 @@ +import { lookup } from 'node:dns'; +import { Agent } from 'node:https'; import debug from 'debug'; +import ipaddr from 'ipaddr.js'; +import nodeFetch from 'node-fetch'; import { CLIENT_METADATA_CACHE_SECONDS } from '@/lib/oauth/config'; import { OAuthError } from '@/lib/oauth/errors'; import { isAcceptableRedirectUri, parseUrl } from '@/lib/oauth/redirect'; @@ -11,6 +15,40 @@ const METADATA_FETCH_TIMEOUT_MS = 5000; const METADATA_MAX_BYTES = 64 * 1024; const MAX_REDIRECT_URIS = 20; +const metadataAgent = new Agent({ + keepAlive: false, + lookup(hostname, options, callback) { + // Validate the actual socket lookup, then connect only to those addresses. A separate + // preflight lookup would let a rebinding hostname return a private IP on the next lookup. + lookup(hostname, { all: true }, (error, addresses) => { + if (error) { + callback(error, '', 0); + return; + } + + if ( + !addresses.length || + addresses.some(({ address }) => ipaddr.parse(address).range() !== 'unicast') + ) { + callback(new Error('Client metadata must resolve only to public addresses.'), '', 0); + return; + } + + const destinations = options.family + ? addresses.filter(({ family }) => family === options.family) + : addresses; + + if (!destinations.length) { + callback(new Error('No public address found for the requested address family.'), '', 0); + } else if (options.all) { + callback(null, destinations); + } else { + callback(null, destinations[0].address, destinations[0].family); + } + }); + }, +}); + export interface ResolvedOAuthClient { clientId: string; clientName: string; @@ -162,12 +200,32 @@ async function writeCache(clientId: string, value: ResolvedOAuthClient) { }); } -export type MetadataFetch = (url: string, init: RequestInit) => Promise; +export type MetadataFetch = ( + url: string, + init: RequestInit, +) => Promise>; + +const fetchPublicMetadata: MetadataFetch = (url, init) => + nodeFetch(url, { + method: 'GET', + headers: { accept: 'application/json' }, + redirect: 'error', + signal: init.signal, + agent: metadataAgent, + size: METADATA_MAX_BYTES, + }); export async function fetchClientMetadata( clientId: string, - fetchImpl: MetadataFetch = fetch, + fetchImpl: MetadataFetch = fetchPublicMetadata, ): Promise { + if (!isClientMetadataUrl(clientId)) { + throw new OAuthError( + 'invalid_client', + 'Client metadata must use a public https URL with a path.', + ); + } + const controller = new AbortController(); const timer = setTimeout(() => controller.abort(), METADATA_FETCH_TIMEOUT_MS); @@ -188,7 +246,7 @@ export async function fetchClientMetadata( const text = await response.text(); - if (text.length > METADATA_MAX_BYTES) { + if (Buffer.byteLength(text, 'utf8') > METADATA_MAX_BYTES) { throw new OAuthError('invalid_client_metadata', 'Client metadata document is too large.'); } @@ -215,6 +273,7 @@ export async function fetchClientMetadata( throw new OAuthError('invalid_client', 'Could not fetch client metadata document.'); } finally { clearTimeout(timer); + controller.abort(); } } diff --git a/src/lib/oauth/token.test.ts b/src/lib/oauth/token.test.ts index b7aa25040..b326614ac 100644 --- a/src/lib/oauth/token.test.ts +++ b/src/lib/oauth/token.test.ts @@ -12,6 +12,14 @@ import { computeCodeChallenge } from './pkce'; import { handleRevocationRequest, handleTokenRequest } from './token'; import { generateRefreshToken, hashToken, verifyAccessToken } from './tokens'; +const { transactionClient } = vi.hoisted(() => ({ transactionClient: {} })); + +vi.mock('@/lib/prisma', () => ({ + default: { + transaction: vi.fn(async callback => callback(transactionClient)), + }, +})); + vi.mock('@/queries/prisma/oauth', () => ({ consumeAuthorizationCode: vi.fn(), createRefreshToken: vi.fn(), @@ -107,6 +115,7 @@ describe('authorization_code grant', () => { clientId: CLIENT_ID, tokenHash: hashToken(tokens.refresh_token), }), + transactionClient, ); }); @@ -146,7 +155,11 @@ describe('authorization_code grant', () => { await expect(handleTokenRequest(grant(), ISSUER)).rejects.toMatchObject({ code: 'invalid_grant', }); - expect(revokeRefreshTokensForClientMock).toHaveBeenCalledWith(USER.id, CLIENT_ID); + expect(revokeRefreshTokensForClientMock).toHaveBeenCalledWith( + USER.id, + CLIENT_ID, + transactionClient, + ); }); test('rejects a mismatched resource indicator', async () => { @@ -194,7 +207,7 @@ describe('refresh_token grant', () => { ); expect(getRefreshTokenByHashMock).toHaveBeenCalledWith(hashToken(presented)); - expect(revokeRefreshTokenMock).toHaveBeenCalledWith('rt-1'); + expect(revokeRefreshTokenMock).toHaveBeenCalledWith('rt-1', transactionClient); expect(tokens.refresh_token).not.toBe(presented); expect(tokens.scope).toBe('analytics:read websites:read'); }); diff --git a/src/lib/oauth/token.ts b/src/lib/oauth/token.ts index 33edf0c91..98d172b81 100644 --- a/src/lib/oauth/token.ts +++ b/src/lib/oauth/token.ts @@ -1,3 +1,4 @@ +import type { Prisma } from '@/generated/prisma/client'; import { uuid } from '@/lib/crypto'; import { getMcpResource, normalizeResource, REFRESH_TOKEN_TTL_SECONDS } from '@/lib/oauth/config'; import { OAuthError } from '@/lib/oauth/errors'; @@ -9,6 +10,7 @@ import { hashToken, isRefreshToken, } from '@/lib/oauth/tokens'; +import prisma from '@/lib/prisma'; import { consumeAuthorizationCode, createRefreshToken, @@ -19,6 +21,10 @@ import { } from '@/queries/prisma/oauth'; import { getUser } from '@/queries/prisma/user'; +const transaction = prisma.transaction as ( + callback: (tx: Prisma.TransactionClient) => Promise, +) => Promise; + export interface TokenRequestParams { grant_type?: unknown; code?: unknown; @@ -60,14 +66,17 @@ async function loadUser(userId: string) { return user; } -async function issueTokens(options: { - userId: string; - clientId: string; - scopes: string[]; - resource: string; - issuer: string; - passwordHash?: string; -}): Promise { +async function issueTokens( + options: { + userId: string; + clientId: string; + scopes: string[]; + resource: string; + issuer: string; + passwordHash?: string; + }, + tx: Prisma.TransactionClient, +): Promise { const scopes = normalizeScopes(options.scopes); const access = createAccessToken({ userId: options.userId, @@ -79,15 +88,18 @@ async function issueTokens(options: { }); const refreshToken = generateRefreshToken(); - await createRefreshToken({ - id: uuid(), - tokenHash: hashToken(refreshToken), - userId: options.userId, - clientId: options.clientId, - scope: scopes.join(' '), - resource: options.resource, - expiresAt: new Date(Date.now() + REFRESH_TOKEN_TTL_SECONDS * 1000), - }); + await createRefreshToken( + { + id: uuid(), + tokenHash: hashToken(refreshToken), + userId: options.userId, + clientId: options.clientId, + scope: scopes.join(' '), + resource: options.resource, + expiresAt: new Date(Date.now() + REFRESH_TOKEN_TTL_SECONDS * 1000), + }, + tx, + ); return { access_token: access.token, @@ -105,52 +117,67 @@ export async function handleAuthorizationCodeGrant( const code = requireString(params.code, 'code', 512); const clientId = requireString(params.client_id, 'client_id', 2048); const codeVerifier = requireString(params.code_verifier, 'code_verifier', 128); - const { code: record, replayed } = await consumeAuthorizationCode(hashToken(code)); + const result = await transaction(async tx => { + const { code: record, replayed } = await consumeAuthorizationCode(hashToken(code), tx); - if (!record) { - if (replayed) { - // Replay of a redeemed code: revoke everything that code produced (OAuth 2.1 §4.1.2). - await revokeRefreshTokensForClient(replayed.userId, replayed.clientId); + if (!record) { + if (replayed) { + // Replay of a redeemed code: revoke everything that code produced (OAuth 2.1 §4.1.2). + await revokeRefreshTokensForClient(replayed.userId, replayed.clientId, tx); + } + + // Return the error so replay revocations commit before the request is rejected. + return new OAuthError( + 'invalid_grant', + 'Authorization code is invalid, expired or already used.', + ); } - throw new OAuthError( - 'invalid_grant', - 'Authorization code is invalid, expired or already used.', + if (record.clientId !== clientId) { + throw new OAuthError('invalid_grant', 'client_id does not match the authorization code.'); + } + + const redirectUri = typeof params.redirect_uri === 'string' ? params.redirect_uri : undefined; + + if (redirectUri !== undefined && redirectUri !== record.redirectUri) { + throw new OAuthError( + 'invalid_grant', + 'redirect_uri does not match the authorization request.', + ); + } + + if (!verifyPkce(codeVerifier, record.codeChallenge, record.codeChallengeMethod)) { + throw new OAuthError('invalid_grant', 'PKCE verification failed.'); + } + + const expectedResource = record.resource ?? getMcpResource(issuer); + const resourceParam = + typeof params.resource === 'string' ? normalizeResource(params.resource) : null; + + if (params.resource !== undefined && resourceParam !== expectedResource) { + throw new OAuthError('invalid_target', `resource must be ${expectedResource}.`); + } + + const user = await loadUser(record.userId); + + return issueTokens( + { + userId: user.id, + clientId: record.clientId, + scopes: parseScopes(record.scope), + resource: expectedResource, + issuer, + passwordHash: user.password, + }, + tx, ); - } - - if (record.clientId !== clientId) { - throw new OAuthError('invalid_grant', 'client_id does not match the authorization code.'); - } - - const redirectUri = typeof params.redirect_uri === 'string' ? params.redirect_uri : undefined; - - if (redirectUri !== undefined && redirectUri !== record.redirectUri) { - throw new OAuthError('invalid_grant', 'redirect_uri does not match the authorization request.'); - } - - if (!verifyPkce(codeVerifier, record.codeChallenge, record.codeChallengeMethod)) { - throw new OAuthError('invalid_grant', 'PKCE verification failed.'); - } - - const expectedResource = record.resource ?? getMcpResource(issuer); - const resourceParam = - typeof params.resource === 'string' ? normalizeResource(params.resource) : null; - - if (params.resource !== undefined && resourceParam !== expectedResource) { - throw new OAuthError('invalid_target', `resource must be ${expectedResource}.`); - } - - const user = await loadUser(record.userId); - - return issueTokens({ - userId: user.id, - clientId: record.clientId, - scopes: parseScopes(record.scope), - resource: expectedResource, - issuer, - passwordHash: user.password, }); + + if (result instanceof OAuthError) { + throw result; + } + + return result; } export async function handleRefreshTokenGrant( @@ -195,21 +222,26 @@ export async function handleRefreshTokenGrant( const user = await loadUser(record.userId); // Rotate: the presented refresh token is retired and a new one issued (OAuth 2.1 §4.3.1). - const revoked = await revokeRefreshToken(record.id); + return transaction(async tx => { + const revoked = await revokeRefreshToken(record.id, tx); - if (!revoked) { - throw new OAuthError('invalid_grant', 'Refresh token has already been used.'); - } + if (!revoked) { + throw new OAuthError('invalid_grant', 'Refresh token has already been used.'); + } - await touchRefreshToken(record.id); + await touchRefreshToken(record.id, tx); - return issueTokens({ - userId: user.id, - clientId: record.clientId, - scopes: narrowed, - resource: expectedResource, - issuer, - passwordHash: user.password, + return issueTokens( + { + userId: user.id, + clientId: record.clientId, + scopes: narrowed, + resource: expectedResource, + issuer, + passwordHash: user.password, + }, + tx, + ); }); } diff --git a/src/queries/prisma/oauth.ts b/src/queries/prisma/oauth.ts index fcee957d0..6f379c37c 100644 --- a/src/queries/prisma/oauth.ts +++ b/src/queries/prisma/oauth.ts @@ -1,5 +1,10 @@ +import type { Prisma } from '@/generated/prisma/client'; import prisma from '@/lib/prisma'; +type OAuthDatabase = Pick; +// The read-replica extension preserves these delegates, but changes their generic signatures. +const oauthDatabase = prisma.client as OAuthDatabase; + export async function createOauthClient(data: { id: string; name: string; @@ -39,14 +44,17 @@ export async function createAuthorizationCode(data: { * Atomically marks a code as used and returns it, or `null` when the code is unknown, expired or * has already been redeemed. Reuse of a code is a replay signal, handled by the caller. */ -export async function consumeAuthorizationCode(codeHash: string) { +export async function consumeAuthorizationCode( + codeHash: string, + db: OAuthDatabase = oauthDatabase, +) { const now = new Date(); - const { count } = await prisma.client.oauthAuthorizationCode.updateMany({ + const { count } = await db.oauthAuthorizationCode.updateMany({ where: { codeHash, usedAt: null, expiresAt: { gt: now } }, data: { usedAt: now }, }); - const code = await prisma.client.oauthAuthorizationCode.findUnique({ where: { codeHash } }); + const code = await db.oauthAuthorizationCode.findUnique({ where: { codeHash } }); if (!code) { return { code: null, replayed: null }; @@ -60,24 +68,27 @@ export async function consumeAuthorizationCode(codeHash: string) { return { code: null, replayed: code.usedAt && code.usedAt < new Date() ? code : null }; } -export async function createRefreshToken(data: { - id: string; - tokenHash: string; - userId: string; - clientId: string; - scope: string; - resource?: string | null; - expiresAt: Date; -}) { - return prisma.client.oauthRefreshToken.create({ data }); +export async function createRefreshToken( + data: { + id: string; + tokenHash: string; + userId: string; + clientId: string; + scope: string; + resource?: string | null; + expiresAt: Date; + }, + db: OAuthDatabase = oauthDatabase, +) { + return db.oauthRefreshToken.create({ data }); } export async function getRefreshTokenByHash(tokenHash: string) { return prisma.client.oauthRefreshToken.findUnique({ where: { tokenHash } }); } -export async function revokeRefreshToken(id: string) { - const { count } = await prisma.client.oauthRefreshToken.updateMany({ +export async function revokeRefreshToken(id: string, db: OAuthDatabase = oauthDatabase) { + const { count } = await db.oauthRefreshToken.updateMany({ where: { id, revokedAt: null }, data: { revokedAt: new Date() }, }); @@ -85,8 +96,12 @@ export async function revokeRefreshToken(id: string) { return count; } -export async function revokeRefreshTokensForClient(userId: string, clientId: string) { - const { count } = await prisma.client.oauthRefreshToken.updateMany({ +export async function revokeRefreshTokensForClient( + userId: string, + clientId: string, + db: OAuthDatabase = oauthDatabase, +) { + const { count } = await db.oauthRefreshToken.updateMany({ where: { userId, clientId, revokedAt: null }, data: { revokedAt: new Date() }, }); @@ -94,8 +109,8 @@ export async function revokeRefreshTokensForClient(userId: string, clientId: str return count; } -export async function touchRefreshToken(id: string) { - return prisma.client.oauthRefreshToken.update({ +export async function touchRefreshToken(id: string, db: OAuthDatabase = oauthDatabase) { + return db.oauthRefreshToken.update({ where: { id }, data: { lastUsedAt: new Date() }, });