From db08f6c1758822c05d42a2815948ae491f1f88e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Wed, 29 Jul 2026 13:50:14 +0200 Subject: [PATCH] fix(01-05): serialize MCP refresh families - lock and rotate refresh families inside one database transaction - revoke CAS losers and insert successors before commit - sign refresh artifacts only after committed state --- src/app/api/mcp/oauth/token/route.ts | 86 ++++++++++++++++++++-------- src/lib/auth/issue-token.ts | 4 +- src/lib/mcp/oauth/connections.ts | 41 ++++++++++--- 3 files changed, 100 insertions(+), 31 deletions(-) diff --git a/src/app/api/mcp/oauth/token/route.ts b/src/app/api/mcp/oauth/token/route.ts index c61342fdd..e9493db20 100644 --- a/src/app/api/mcp/oauth/token/route.ts +++ b/src/app/api/mcp/oauth/token/route.ts @@ -29,6 +29,8 @@ import { randomUUID } from "node:crypto"; import { NextRequest } from "next/server"; +import type { Prisma } from "@/generated/prisma/client"; +import { prisma } from "@/lib/db"; import { annotate } from "@/lib/logging/context"; import { withBackgroundEvent } from "@/lib/logging/background"; import { auditLog } from "@/lib/auth/audit"; @@ -109,6 +111,16 @@ interface RefreshClaims { resource: string; } +type TokenPairArgs = { + userId: string; + clientId: string; + scope: string; + resource: string; + grant: "authorization_code" | "refresh_token"; + connectionId?: string; + refreshJti?: string; +}; + /** * Mint the access token (+ refresh artifact when `offline_access` is granted). * @@ -119,15 +131,15 @@ interface RefreshClaims { * the connection seeded with it; for `refresh_token` the caller has just * advanced the connection to it. */ -async function mintTokenPair(args: { - userId: string; - clientId: string; - scope: string; - resource: string; - grant: "authorization_code" | "refresh_token"; - connectionId?: string; - refreshJti?: string; -}): Promise { +async function mintTokenPair(args: TokenPairArgs): Promise { + const access = await persistAccessToken(args); + return finalizeTokenPair(args, access); +} + +async function persistAccessToken( + args: TokenPairArgs, + db?: Pick, +) { // Grant EXACTLY the scopes the authorization code carried (read, or // read+write). `health:read` is always present; `health:write` is added // only when the user consented to it at the authorize step. The token is @@ -139,14 +151,20 @@ async function mintTokenPair(args: { permissions.push(SCOPE_HEALTH_WRITE); } - const access = await issueApiToken({ + const options = { userId: args.userId, name: "MCP connector", permissions, expiresInMinutes: ACCESS_TOKEN_TTL_MINUTES, ...(args.connectionId ? { mcpConnectionId: args.connectionId } : {}), - }); + }; + return db ? issueApiToken(options, db) : issueApiToken(options); +} +async function finalizeTokenPair( + args: TokenPairArgs, + access: Awaited>, +): Promise { const refresh = args.connectionId && args.refreshJti ? signArtifact( @@ -367,14 +385,42 @@ export async function POST(request: NextRequest): Promise { ); } const newJti = randomUUID(); - const rotation = await rotateConnection({ + const rotationArgs = { connectionId: claims.cid, presentedJti: claims.jti, newJti, clientId: claims.client_id, userId: claims.sub, - }); - if (!rotation.ok) { + }; + const pairArgs: TokenPairArgs = { + userId: claims.sub, + clientId: claims.client_id, + scope: claims.scope, + resource: claims.resource, + grant: "refresh_token", + connectionId: claims.cid, + refreshJti: newJti, + }; + + // Rotation and successor insertion are one serialized state transition. + // A replay transaction commits family revocation; a successful + // transaction commits both JTI advancement and the successor access row. + // Any insertion error rejects the callback and PostgreSQL rolls every + // earlier state change back. No artifact is signed inside this callback. + const exchange = await prisma.$transaction( + async (tx) => { + const rotation = await rotateConnection(rotationArgs, tx); + if (!rotation.ok) { + return { ok: false as const, rotation }; + } + const access = await persistAccessToken(pairArgs, tx); + return { ok: true as const, access }; + }, + { timeout: 10_000 }, + ); + + if (!exchange.ok) { + const { rotation } = exchange; annotate({ action: { name: @@ -392,15 +438,9 @@ export async function POST(request: NextRequest): Promise { ); } - return mintTokenPair({ - userId: claims.sub, - clientId: claims.client_id, - scope: claims.scope, - resource: claims.resource, - grant: "refresh_token", - connectionId: claims.cid, - refreshJti: newJti, - }); + // The transaction has committed. Only now may the stateless successor + // refresh artifact be signed and returned to the client. + return finalizeTokenPair(pairArgs, exchange.access); } return oauthError( diff --git a/src/lib/auth/issue-token.ts b/src/lib/auth/issue-token.ts index 291b34342..0939aaec2 100644 --- a/src/lib/auth/issue-token.ts +++ b/src/lib/auth/issue-token.ts @@ -5,6 +5,7 @@ * cookie. The raw `hlk_` value is returned exactly once. */ import { randomBytes } from "node:crypto"; +import type { Prisma } from "@/generated/prisma/client"; import { prisma } from "@/lib/db"; import { hashToken } from "@/lib/auth/hmac"; @@ -44,6 +45,7 @@ export interface IssueTokenOptions { */ export async function issueApiToken( opts: IssueTokenOptions, + db: Pick = prisma, ): Promise { const rawToken = `hlk_${randomBytes(32).toString("hex")}`; const tokenHashValue = hashToken(rawToken); @@ -52,7 +54,7 @@ export async function issueApiToken( ? new Date(Date.now() + opts.expiresInMinutes * 60 * 1000) : new Date(Date.now() + (opts.expiresInDays ?? 90) * 24 * 60 * 60 * 1000); - const created = await prisma.apiToken.create({ + const created = await db.apiToken.create({ data: { userId: opts.userId, name: opts.name, diff --git a/src/lib/mcp/oauth/connections.ts b/src/lib/mcp/oauth/connections.ts index 091143e01..f2fafa4a9 100644 --- a/src/lib/mcp/oauth/connections.ts +++ b/src/lib/mcp/oauth/connections.ts @@ -26,8 +26,14 @@ * The access tokens themselves stay ordinary `hlk_` `ApiToken` rows, linked * back through `ApiToken.mcpConnectionId`. */ +import type { Prisma } from "@/generated/prisma/client"; import { prisma } from "@/lib/db"; +type ConnectionDb = Pick< + Prisma.TransactionClient, + "mcpOAuthConnection" | "apiToken" | "$queryRaw" +>; + export interface ConnectionRotation { /** The connection id (`cid`) embedded in the refresh artifact. */ connectionId: string; @@ -82,8 +88,22 @@ export async function createConnection(args: { */ export async function rotateConnection( args: ConnectionRotation, + db: ConnectionDb = prisma, ): Promise { - const connection = await prisma.mcpOAuthConnection.findUnique({ + // PostgreSQL has no Prisma model-level row-lock primitive. Inside the token + // route's interactive transaction this serializes every decision for one + // refresh family before we re-read its current state. The runtime guard keeps + // lightweight unit doubles compatible while production clients always lock. + if (typeof db.$queryRaw === "function") { + await db.$queryRaw` + SELECT id + FROM mcp_oauth_connections + WHERE id = ${args.connectionId} + FOR UPDATE + `; + } + + const connection = await db.mcpOAuthConnection.findUnique({ where: { id: args.connectionId }, select: { id: true, @@ -106,13 +126,13 @@ export async function rotateConnection( if (connection.currentJti !== args.presentedJti) { // Replay of an already-rotated refresh token — treat as theft and revoke // the whole connection (reuse-detection family revocation). - await revokeConnectionRow(connection.id); + await revokeConnectionRow(connection.id, db); return { ok: false, reason: "reuse_detected" }; } // Advance the chain atomically: only the caller presenting the current jti // wins, so two concurrent refreshes cannot both rotate the same row. - const advanced = await prisma.mcpOAuthConnection.updateMany({ + const advanced = await db.mcpOAuthConnection.updateMany({ where: { id: connection.id, currentJti: args.presentedJti, @@ -121,12 +141,16 @@ export async function rotateConnection( data: { currentJti: args.newJti, lastUsedAt: new Date() }, }); if (advanced.count === 0) { + // A contender that loses the guarded advancement is a replay signal just + // like a stale JTI observed after taking the row lock. Revoke the complete + // family before returning; the surrounding transaction commits this deny. + await revokeConnectionRow(connection.id, db); return { ok: false, reason: "reuse_detected" }; } // L4 — revoke the access tokens minted before this rotation so a leaked // 60-minute access token cannot outlive the refresh that spawned it. - await prisma.apiToken.updateMany({ + await db.apiToken.updateMany({ where: { mcpConnectionId: connection.id, revoked: false }, data: { revoked: true }, }); @@ -135,12 +159,15 @@ export async function rotateConnection( } /** Internal — stamp `revokedAt` and revoke every linked access token. */ -async function revokeConnectionRow(connectionId: string): Promise { - await prisma.mcpOAuthConnection.updateMany({ +async function revokeConnectionRow( + connectionId: string, + db: Pick = prisma, +): Promise { + await db.mcpOAuthConnection.updateMany({ where: { id: connectionId, revokedAt: null }, data: { revokedAt: new Date() }, }); - await prisma.apiToken.updateMany({ + await db.apiToken.updateMany({ where: { mcpConnectionId: connectionId, revoked: false }, data: { revoked: true }, });