Harden OAuth metadata fetching and make token rotation transactional

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.
 

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)) {

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<Response>;
+export type MetadataFetch = (
+  url: string,
+  init: RequestInit,
+) => Promise<Pick<Response, 'ok' | 'status' | 'text'>>;
+
+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<ResolvedOAuthClient> {
+  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();
   }
 }
 

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');
   });

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 <T>(
+  callback: (tx: Prisma.TransactionClient) => Promise<T>,
+) => Promise<T>;
+
 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<TokenResponse> {
+async function issueTokens(
+  options: {
+    userId: string;
+    clientId: string;
+    scopes: string[];
+    resource: string;
+    issuer: string;
+    passwordHash?: string;
+  },
+  tx: Prisma.TransactionClient,
+): Promise<TokenResponse> {
   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,
+    );
   });
 }
 

src/queries/prisma/oauth.ts

@@ -1,5 +1,10 @@
+import type { Prisma } from '@/generated/prisma/client';
 import prisma from '@/lib/prisma';
 
+type OAuthDatabase = Pick<Prisma.TransactionClient, 'oauthAuthorizationCode' | 'oauthRefreshToken'>;
+// 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() },
   });