From 5078b4e02d01840f1211854094822706072e853d Mon Sep 17 00:00:00 2001 From: Jakob Wennberg <149234542+jakobwennberg@users.noreply.github.com> Date: Fri, 12 Jun 2026 10:05:10 +0200 Subject: [PATCH] fix(mcp): grace window + idempotent refresh-token replay for OAuth (#710) (#714) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(mcp): grace window + idempotent refresh-token replay for OAuth (#710) OAuth refresh rotated BOTH the refresh token and the access key in one zero-grace CAS. Claude Code's MCP OAuth client fails to persist the rotated refresh token (or fires concurrent refreshes), re-presents the stale one, the CAS matches 0 rows, and the grant dies with invalid_grant — forcing a full re-authorization roughly every 60s in a loop. Regression from #392. Keep rotation (RFC 9700 §4.14.2 requires it for public clients) but add a bounded grace window with idempotent replay, atomic in one SECURITY DEFINER RPC: - Migration adds previous_key_hash / previous_refresh_token_hash (+ *_expires_at) shadow columns. validate_and_increment_api_key accepts the current OR an unexpired previous key_hash, with the rate-limit increment keyed off the resolved row id. - New rotate_mcp_refresh_token RPC: rotated | replayed | reuse_revoked | revoked | invalid. In-grace replay re-issues a fresh pair and slides the window so an actively-refreshing client that cannot persist the rotated token keeps working; reuse after the window revokes the grant family (RFC 9700 4.14.2 reuse detection preserved). - The refresh grant now calls the one RPC, closing the old SELECT-then-CAS TOCTOU gap. All previous_* columns default NULL, so existing keys are unaffected and the RPC return shape is unchanged (callers untouched). Tests: rewired the token-route unit tests to the RPC and replaced the test that codified the bug with a #710 regression (in-grace replay returns 200, not 400); added tests/pg/mcp-oauth-rotation-grace.pg.test.ts (grace accept/expire, revoke-never-graced, rotate->demote, idempotent replay, reuse-after-grace->revoke). Co-Authored-By: Claude Opus 4.8 (1M context) * ci: retrigger checks for #714 No code change — re-running CI. The Supabase Preview check fails on a pre-existing main-branch migration-history drift ("Remote migration versions not found in local migrations directory"), not this PR; pg-real (full migration replay) passes. Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- .../mcp-oauth/token/__tests__/route.test.ts | 96 +++---- app/api/mcp-oauth/token/route.ts | 94 +++--- ...20260621130000_api_keys_rotation_grace.sql | 233 +++++++++++++++ tests/pg/mcp-oauth-rotation-grace.pg.test.ts | 270 ++++++++++++++++++ 4 files changed, 587 insertions(+), 106 deletions(-) create mode 100644 supabase/migrations/20260621130000_api_keys_rotation_grace.sql create mode 100644 tests/pg/mcp-oauth-rotation-grace.pg.test.ts diff --git a/app/api/mcp-oauth/token/__tests__/route.test.ts b/app/api/mcp-oauth/token/__tests__/route.test.ts index f312ec6c..ffc1db65 100644 --- a/app/api/mcp-oauth/token/__tests__/route.test.ts +++ b/app/api/mcp-oauth/token/__tests__/route.test.ts @@ -147,12 +147,10 @@ describe('POST /api/mcp-oauth/token', () => { it('rotates both tokens and returns a fresh access_token', async () => { const { token: refreshToken } = generateRefreshToken() - const { supabase, enqueueMany } = createQueuedMockSupabase() + const { supabase, enqueue } = createQueuedMockSupabase() mocks.supabaseFactory.mockReturnValue(supabase) - enqueueMany([ - { data: { id: 'key-1', revoked_at: null }, error: null }, // SELECT - { data: [{ id: 'key-1' }], error: null }, // UPDATE ... RETURNING - ]) + // rotate_mcp_refresh_token RPC → normal rotation + enqueue({ data: [{ outcome: 'rotated', scopes: null }], error: null }) const res = await POST( formRequest({ @@ -179,7 +177,7 @@ describe('POST /api/mcp-oauth/token', () => { it('returns 400 when refresh_token is unknown', async () => { const { supabase, enqueue } = createQueuedMockSupabase() mocks.supabaseFactory.mockReturnValue(supabase) - enqueue({ data: null, error: null }) // SELECT — no row + enqueue({ data: [{ outcome: 'invalid', scopes: null }], error: null }) const res = await POST( formRequest({ @@ -195,10 +193,7 @@ describe('POST /api/mcp-oauth/token', () => { it('returns 400 when the api_key is revoked', async () => { const { supabase, enqueue } = createQueuedMockSupabase() mocks.supabaseFactory.mockReturnValue(supabase) - enqueue({ - data: { id: 'key-1', revoked_at: '2026-05-01T00:00:00Z' }, - error: null, - }) + enqueue({ data: [{ outcome: 'revoked', scopes: null }], error: null }) const res = await POST( formRequest({ @@ -212,7 +207,7 @@ describe('POST /api/mcp-oauth/token', () => { expect(body.error_description).toContain('revoked') }) - it('returns 500 when the lookup fails with a DB error', async () => { + it('returns 500 when the rotation RPC fails with a DB error', async () => { const { supabase, enqueue } = createQueuedMockSupabase() mocks.supabaseFactory.mockReturnValue(supabase) enqueue({ data: null, error: { message: 'connection reset' } }) @@ -228,32 +223,12 @@ describe('POST /api/mcp-oauth/token', () => { expect(body.error).toBe('server_error') }) - it('returns 500 when the rotation update fails with a DB error', async () => { - const { supabase, enqueueMany } = createQueuedMockSupabase() + it('returns 400 invalid_grant when a refresh token is reused after its grace window (reuse_revoked)', async () => { + const { supabase, enqueue } = createQueuedMockSupabase() mocks.supabaseFactory.mockReturnValue(supabase) - enqueueMany([ - { data: { id: 'key-1', revoked_at: null }, error: null }, // SELECT - { data: null, error: { message: 'deadlock detected' } }, // UPDATE — DB error - ]) - - const res = await POST( - formRequest({ - grant_type: 'refresh_token', - refresh_token: 'gnubok_rt_anything', - }) - ) - expect(res.status).toBe(500) - const body = await res.json() - expect(body.error).toBe('server_error') - }) - - it('returns 400 when the CAS update affects 0 rows (concurrent reuse)', async () => { - const { supabase, enqueueMany } = createQueuedMockSupabase() - mocks.supabaseFactory.mockReturnValue(supabase) - enqueueMany([ - { data: { id: 'key-1', revoked_at: null }, error: null }, // SELECT - { data: [], error: null }, // UPDATE — 0 rows (lost the CAS race) - ]) + // The RPC detected reuse of a previous refresh token past its grace + // window and already revoked the grant family (RFC 9700 §4.14.2). + enqueue({ data: [{ outcome: 'reuse_revoked', scopes: null }], error: null }) const res = await POST( formRequest({ @@ -264,7 +239,29 @@ describe('POST /api/mcp-oauth/token', () => { expect(res.status).toBe(400) const body = await res.json() expect(body.error).toBe('invalid_grant') - expect(body.error_description).toContain('already used') + }) + + it('returns a fresh pair on idempotent in-grace replay instead of 400 (issue #710 regression)', async () => { + // A retried / mis-persisted / concurrent refresh presents the previous + // refresh token within the grace window. The old code returned 400 + // "already used", stranding Claude Code into a re-auth loop; the RPC now + // replays idempotently and the client gets a working pair. + const { supabase, enqueue } = createQueuedMockSupabase() + mocks.supabaseFactory.mockReturnValue(supabase) + enqueue({ data: [{ outcome: 'replayed', scopes: ['transactions:read'] }], error: null }) + + const res = await POST( + formRequest({ + grant_type: 'refresh_token', + refresh_token: 'gnubok_rt_anything', + }) + ) + expect(res.status).toBe(200) + const body = await res.json() + expect(body.access_token).toMatch(/^gnubok_sk_/) + expect(body.refresh_token).toMatch(/^gnubok_rt_/) + expect(body.expires_in).toBe(3600) + expect(body.scope).toBe('transactions:read') }) }) @@ -387,19 +384,17 @@ describe('POST /api/mcp-oauth/token', () => { it('returns the granular scopes the api_key was minted with', async () => { // Greptile P1 — refresh response previously hardcoded scope:'mcp', // causing OAuth 2.1 clients to think they had lost their grant. - const { supabase, enqueueMany } = createQueuedMockSupabase() + const { supabase, enqueue } = createQueuedMockSupabase() mocks.supabaseFactory.mockReturnValue(supabase) - enqueueMany([ - { - data: { - id: 'key-1', - revoked_at: null, + enqueue({ + data: [ + { + outcome: 'rotated', scopes: ['transactions:read', 'invoices:read', 'invoices:write'], }, - error: null, - }, // SELECT - { data: [{ id: 'key-1' }], error: null }, // UPDATE - ]) + ], + error: null, + }) const res = await POST( formRequest({ @@ -415,12 +410,9 @@ describe('POST /api/mcp-oauth/token', () => { }) it('falls back to read-only DEFAULT_OAUTH_SCOPES for legacy keys with null scopes', async () => { - const { supabase, enqueueMany } = createQueuedMockSupabase() + const { supabase, enqueue } = createQueuedMockSupabase() mocks.supabaseFactory.mockReturnValue(supabase) - enqueueMany([ - { data: { id: 'key-1', revoked_at: null, scopes: null }, error: null }, - { data: [{ id: 'key-1' }], error: null }, - ]) + enqueue({ data: [{ outcome: 'rotated', scopes: null }], error: null }) const res = await POST( formRequest({ diff --git a/app/api/mcp-oauth/token/route.ts b/app/api/mcp-oauth/token/route.ts index c8da9bc3..c0b7e559 100644 --- a/app/api/mcp-oauth/token/route.ts +++ b/app/api/mcp-oauth/token/route.ts @@ -13,6 +13,13 @@ import { requireCompanyId } from '@/lib/company/context' const ACCESS_TOKEN_TTL_SECONDS = 3600 +// Grace window (seconds) during which a just-superseded access token and refresh +// token stay valid after a rotation. Lets a client that cannot reliably persist +// the rotated refresh token — or that fires concurrent refreshes — recover via +// idempotent replay instead of being forced into a re-auth loop (issue #710). +// Reuse of a previous refresh token AFTER this window revokes the grant. +const REFRESH_GRACE_SECONDS = 120 + /** * OAuth 2.0 Token Endpoint. * @@ -184,75 +191,54 @@ async function handleRefreshTokenGrant(params: URLSearchParams) { const supabase = createServiceClientNoCookies() const presentedHash = hashRefreshToken(refreshToken) - // Look up the api_key row by refresh_token_hash. The hash is unique among - // non-null values, so there's at most one match. We pull `scopes` so the - // rotated token response advertises the same granular grant the key - // already carries (otherwise OAuth 2.1 clients would re-authorize on - // every refresh). - const { data: row, error: lookupError } = await supabase - .from('api_keys') - .select('id, revoked_at, scopes') - .eq('refresh_token_hash', presentedHash) - .maybeSingle() - - if (lookupError) { - return NextResponse.json( - { error: 'server_error', error_description: 'Failed to look up refresh token' }, - { status: 500 } - ) - } - - if (!row) { - return NextResponse.json( - { error: 'invalid_grant', error_description: 'Invalid refresh token' }, - { status: 400 } - ) - } - - if (row.revoked_at) { - return NextResponse.json( - { error: 'invalid_grant', error_description: 'Refresh token revoked' }, - { status: 400 } - ) - } - - // Rotate both tokens atomically. OAuth 2.1 §6.1 recommends rotating the - // refresh token; we also rotate the api_key because key_hash is one-way - // and we cannot recover the original plaintext to return to the client. - // The .eq('refresh_token_hash', presentedHash) guard makes this a CAS: - // a concurrent refresh with the same token will affect 0 rows. + // Pre-generate the candidate credentials; the RPC decides whether to use them + // (rotate / idempotent replay) or ignore them (reuse / revoked / invalid). + // Doing lookup + rotate + demote in ONE SECURITY DEFINER RPC closes the + // TOCTOU gap the old SELECT-then-CAS had, and lets a just-superseded refresh + // token stay valid for a grace window so a client that cannot persist the + // rotated token — or fires concurrent refreshes — recovers instead of being + // forced into a re-auth loop (issue #710). Reuse AFTER the grace window + // revokes the grant (RFC 9700 §4.14.2 reuse detection). const rotated = generateRefreshToken() const { key: newKey, hash: newKeyHash, prefix: newKeyPrefix } = generateApiKey() - const { data: updated, error: updateError } = await supabase - .from('api_keys') - .update({ - refresh_token_hash: rotated.hash, - key_hash: newKeyHash, - key_prefix: newKeyPrefix, - }) - .eq('id', row.id) - .eq('refresh_token_hash', presentedHash) - .select('id') + const { data, error } = await supabase.rpc('rotate_mcp_refresh_token', { + p_presented_hash: presentedHash, + p_new_refresh_hash: rotated.hash, + p_new_key_hash: newKeyHash, + p_new_key_prefix: newKeyPrefix, + p_grace_seconds: REFRESH_GRACE_SECONDS, + }) - if (updateError) { + if (error) { return NextResponse.json( { error: 'server_error', error_description: 'Failed to rotate refresh token' }, { status: 500 } ) } - if (!updated || updated.length === 0) { + const result = (Array.isArray(data) ? data[0] : data) as + | { outcome: string; scopes: string[] | null } + | undefined + const outcome = result?.outcome + + // 'rotated' (normal) and 'replayed' (idempotent in-grace retry/concurrent) + // both succeed and return the freshly minted pair. Everything else maps to + // invalid_grant; for 'reuse_revoked' the grant was already revoked in the RPC. + if (outcome !== 'rotated' && outcome !== 'replayed') { return NextResponse.json( - { error: 'invalid_grant', error_description: 'Refresh token already used' }, + { + error: 'invalid_grant', + error_description: + outcome === 'revoked' ? 'Refresh token revoked' : 'Invalid or expired refresh token', + }, { status: 400 } ) } - // Return the granular scopes the key was originally minted with. Falling - // back to the read-only OAuth defaults preserves the pre-scope-plumbing - // behaviour for legacy keys whose scopes column is null. - const persistedScopes = validateScopes(row.scopes) ?? DEFAULT_OAUTH_SCOPES + // Carry the granular scopes the key was minted with (read-only OAuth defaults + // for legacy keys with null scopes) so clients don't re-authorize on refresh. + const persistedScopes = validateScopes(result?.scopes ?? null) ?? DEFAULT_OAUTH_SCOPES return NextResponse.json({ access_token: newKey, diff --git a/supabase/migrations/20260621130000_api_keys_rotation_grace.sql b/supabase/migrations/20260621130000_api_keys_rotation_grace.sql new file mode 100644 index 00000000..12c88868 --- /dev/null +++ b/supabase/migrations/20260621130000_api_keys_rotation_grace.sql @@ -0,0 +1,233 @@ +-- Rotation grace + idempotent refresh for MCP OAuth (issue #710). +-- +-- Problem: handleRefreshTokenGrant rotated BOTH the refresh token and the +-- access key in one CAS-guarded UPDATE with zero grace. A client (Claude Code +-- CLI) that fails to persist the rotated refresh token — or fires concurrent +-- refreshes — presents the just-superseded token, the CAS matches 0 rows, and +-- the grant dies with `invalid_grant`, forcing a full re-auth roughly every +-- ~60s in an endless loop. +-- +-- Fix: keep rotation (RFC 9700 §4.14.2 requires public-client refresh tokens to +-- be rotated or sender-constrained) but add a bounded GRACE WINDOW with +-- idempotent replay: +-- * The just-superseded access key (`previous_key_hash`) and refresh token +-- (`previous_refresh_token_hash`) stay valid for a grace window. +-- * Presenting the previous refresh token WITHIN the window re-issues a fresh +-- pair (idempotent replay) and slides the window, so an actively-refreshing +-- client that cannot persist the rotated token is never stranded. +-- * Presenting it AFTER the window is a reuse/breach signal → the grant is +-- revoked (reuse detection preserved). +-- All four shadow columns default NULL → existing keys are unaffected. Because +-- the shadow columns live on the SAME row gated by `revoked_at IS NULL`, a +-- single `revoked_at` UPDATE atomically kills the current AND both previous +-- credentials (RFC 9700 grant-family revocation, no fan-out). + +ALTER TABLE public.api_keys + ADD COLUMN IF NOT EXISTS previous_key_hash text, + ADD COLUMN IF NOT EXISTS previous_key_expires_at timestamptz, + ADD COLUMN IF NOT EXISTS previous_refresh_token_hash text, + ADD COLUMN IF NOT EXISTS previous_refresh_expires_at timestamptz; + +-- Partial-UNIQUE on each previous hash: lookups are point reads and two rows can +-- never claim the same previous hash. Separate from the current-hash indexes so +-- current + previous coexist on one row without collision. +CREATE UNIQUE INDEX IF NOT EXISTS idx_api_keys_previous_key_hash + ON public.api_keys (previous_key_hash) + WHERE previous_key_hash IS NOT NULL; + +CREATE UNIQUE INDEX IF NOT EXISTS idx_api_keys_previous_refresh_token_hash + ON public.api_keys (previous_refresh_token_hash) + WHERE previous_refresh_token_hash IS NOT NULL; + +-- --------------------------------------------------------------------------- +-- validate_and_increment_api_key: accept the current key_hash OR an unexpired +-- previous_key_hash (access-token grace). Identical return shape, so callers in +-- lib/auth/api-keys.ts are unaffected. The rate-limit UPDATE is now keyed off +-- the resolved row id (not p_key_hash) so a grace-hash hit still increments. +-- (CREATE OR REPLACE cannot keep an identical signature across a body change +-- cleanly here; DROP+CREATE mirrors the pattern in 20260512162506.) +-- --------------------------------------------------------------------------- +DROP FUNCTION IF EXISTS public.validate_and_increment_api_key(text); + +CREATE FUNCTION public.validate_and_increment_api_key(p_key_hash text) +RETURNS TABLE( + user_id uuid, + company_id uuid, + api_key_id uuid, + api_key_name text, + rate_limited boolean, + scopes text[], + mode text +) +LANGUAGE plpgsql SECURITY DEFINER AS $$ +DECLARE + v_id uuid; + v_user_id uuid; + v_company_id uuid; + v_api_key_name text; + v_rate_limit_rpm integer; + v_request_count integer; + v_window_start timestamptz; + v_scopes text[]; + v_mode text; +BEGIN + -- Match the live key_hash, OR a previous (just-rotated) key_hash that is still + -- inside its grace window. Both gated by revoked_at IS NULL. + SELECT ak.id, ak.user_id, ak.company_id, ak.name, + ak.rate_limit_rpm, ak.request_count, ak.rate_limit_window_start, ak.scopes, ak.mode + INTO v_id, v_user_id, v_company_id, v_api_key_name, + v_rate_limit_rpm, v_request_count, v_window_start, v_scopes, v_mode + FROM public.api_keys ak + WHERE ak.revoked_at IS NULL + AND ( + ak.key_hash = p_key_hash + OR ( + ak.previous_key_hash = p_key_hash + AND ak.previous_key_expires_at IS NOT NULL + AND ak.previous_key_expires_at > now() + ) + ) + FOR UPDATE; + + IF v_id IS NULL THEN + RETURN; -- no live match (incl. expired grace) → caller returns 401, as before + END IF; + + -- Reset the rate-limit window if it is unset or older than one minute. + IF v_window_start IS NULL OR v_window_start < now() - interval '1 minute' THEN + UPDATE public.api_keys + SET request_count = 1, + rate_limit_window_start = now(), + last_used_at = now() + WHERE id = v_id; + RETURN QUERY SELECT v_user_id, v_company_id, v_id, v_api_key_name, false, v_scopes, v_mode; + RETURN; + END IF; + + IF v_request_count >= v_rate_limit_rpm THEN + RETURN QUERY SELECT v_user_id, v_company_id, v_id, v_api_key_name, true, v_scopes, v_mode; + RETURN; + END IF; + + UPDATE public.api_keys + SET request_count = request_count + 1, + last_used_at = now() + WHERE id = v_id; + + RETURN QUERY SELECT v_user_id, v_company_id, v_id, v_api_key_name, false, v_scopes, v_mode; +END; +$$; + +-- --------------------------------------------------------------------------- +-- rotate_mcp_refresh_token: atomic lookup + rotate + demote + idempotent replay +-- for the OAuth refresh_token grant. Replaces the JS SELECT-then-CAS (which had +-- a TOCTOU gap). The caller pre-generates the candidate new credentials; this +-- function decides whether to use them. +-- +-- Outcomes: +-- 'rotated' presented hash = current refresh token → normal rotation; +-- current creds demoted to previous_* with a grace window. +-- 'replayed' presented hash = an unexpired previous refresh token → the +-- client retried / mis-persisted the rotated token (or a +-- concurrent sibling already rotated). Re-issue a fresh pair +-- and SLIDE the grace window so an actively-refreshing client +-- keeps working. The previous hash is preserved (not chained) +-- so a truly idle gap longer than the window still trips reuse. +-- 'reuse_revoked' presented hash = an EXPIRED previous refresh token → reuse +-- after grace = breach (RFC 9700 §4.14.2) → revoke the grant. +-- 'revoked' the matched grant is already revoked. +-- 'invalid' no matching grant. +-- --------------------------------------------------------------------------- +CREATE OR REPLACE FUNCTION public.rotate_mcp_refresh_token( + p_presented_hash text, + p_new_refresh_hash text, + p_new_key_hash text, + p_new_key_prefix text, + p_grace_seconds integer +) +RETURNS TABLE(outcome text, scopes text[]) +LANGUAGE plpgsql SECURITY DEFINER AS $$ +DECLARE + v_id uuid; + v_revoked timestamptz; + v_scopes text[]; + v_cur_key_hash text; + v_prev_refresh_expires timestamptz; + v_match text; +BEGIN + -- Current refresh token? + SELECT ak.id, ak.revoked_at, ak.scopes, ak.key_hash + INTO v_id, v_revoked, v_scopes, v_cur_key_hash + FROM public.api_keys ak + WHERE ak.refresh_token_hash = p_presented_hash + FOR UPDATE; + + IF v_id IS NOT NULL THEN + v_match := 'current'; + ELSE + -- Previous (just-rotated) refresh token? + SELECT ak.id, ak.revoked_at, ak.scopes, ak.key_hash, ak.previous_refresh_expires_at + INTO v_id, v_revoked, v_scopes, v_cur_key_hash, v_prev_refresh_expires + FROM public.api_keys ak + WHERE ak.previous_refresh_token_hash = p_presented_hash + FOR UPDATE; + IF v_id IS NOT NULL THEN + v_match := 'previous'; + END IF; + END IF; + + IF v_id IS NULL THEN + RETURN QUERY SELECT 'invalid'::text, NULL::text[]; + RETURN; + END IF; + + IF v_revoked IS NOT NULL THEN + RETURN QUERY SELECT 'revoked'::text, NULL::text[]; + RETURN; + END IF; + + IF v_match = 'current' THEN + -- Normal rotation. Demote the consumed refresh token and the superseded + -- access key to previous_* with a grace window, then promote the new pair. + UPDATE public.api_keys + SET previous_refresh_token_hash = p_presented_hash, + previous_refresh_expires_at = now() + make_interval(secs => p_grace_seconds), + previous_key_hash = v_cur_key_hash, + previous_key_expires_at = now() + make_interval(secs => p_grace_seconds), + refresh_token_hash = p_new_refresh_hash, + key_hash = p_new_key_hash, + key_prefix = p_new_key_prefix + WHERE id = v_id; + RETURN QUERY SELECT 'rotated'::text, v_scopes; + RETURN; + END IF; + + -- v_match = 'previous' + IF v_prev_refresh_expires IS NOT NULL AND v_prev_refresh_expires > now() THEN + -- Idempotent in-grace replay: a retried / mis-persisted / concurrent + -- refresh. Re-issue a fresh current pair and SLIDE the grace window. The + -- previous refresh hash stays the presented one (no chaining), so reuse of + -- the original token after an idle gap longer than the window still trips + -- reuse detection below. + UPDATE public.api_keys + SET previous_refresh_expires_at = now() + make_interval(secs => p_grace_seconds), + previous_key_hash = v_cur_key_hash, + previous_key_expires_at = now() + make_interval(secs => p_grace_seconds), + refresh_token_hash = p_new_refresh_hash, + key_hash = p_new_key_hash, + key_prefix = p_new_key_prefix + WHERE id = v_id; + RETURN QUERY SELECT 'replayed'::text, v_scopes; + RETURN; + END IF; + + -- Reuse of a previous refresh token AFTER its grace window = breach signal. + -- Revoke the grant family (single row → kills current + both previous creds). + UPDATE public.api_keys + SET revoked_at = now() + WHERE id = v_id; + RETURN QUERY SELECT 'reuse_revoked'::text, NULL::text[]; +END; +$$; + +NOTIFY pgrst, 'reload schema'; diff --git a/tests/pg/mcp-oauth-rotation-grace.pg.test.ts b/tests/pg/mcp-oauth-rotation-grace.pg.test.ts new file mode 100644 index 00000000..6ebdd8ed --- /dev/null +++ b/tests/pg/mcp-oauth-rotation-grace.pg.test.ts @@ -0,0 +1,270 @@ +import { randomUUID } from 'crypto' +import { describe, expect, it } from 'vitest' +import { seedCompany } from '@/tests/pg/fixtures' +import { getPool } from '@/tests/pg/setup' + +/** + * pg-real coverage for migration 20260621130000_api_keys_rotation_grace. + * + * Locks in the issue #710 fix: + * - The four previous_* shadow columns exist. + * - validate_and_increment_api_key accepts the current key_hash OR an + * unexpired previous_key_hash (access-token grace), rejects an expired one, + * never matches a revoked key, and increments the rate-limit counter on the + * resolved row even when matched via the grace hash. + * - rotate_mcp_refresh_token: 'rotated' demotes the consumed pair to previous_*; + * presenting the previous refresh token within grace 'replayed's (re-issues + * + slides the window); reuse AFTER grace is 'reuse_revoked' and sets + * revoked_at (RFC 9700 §4.14.2); unknown → 'invalid'; revoked grant → 'revoked'. + * + * Inserts go through the pool (superuser, RLS-bypassing) — this is an RPC / + * schema behaviour test, not an RLS test. Hashes are opaque text the RPCs + * compare by equality, so any unique strings work. + */ + +type KeyRow = { + key_hash: string + refresh_token_hash: string | null + previous_key_hash?: string | null + previous_key_expires_at?: string | null + previous_refresh_token_hash?: string | null + previous_refresh_expires_at?: string | null + scopes?: string[] | null + revoked?: boolean +} + +function h(label: string): string { + // Opaque 64-char unique hash-shaped string. + return `${label}_${randomUUID().replace(/-/g, '')}`.padEnd(64, '0').slice(0, 64) +} + +async function insertKey( + userId: string, + companyId: string, + row: KeyRow, +): Promise { + const id = randomUUID() + await getPool().query( + `INSERT INTO public.api_keys + (id, user_id, company_id, key_hash, key_prefix, name, scopes, + refresh_token_hash, previous_key_hash, previous_key_expires_at, + previous_refresh_token_hash, previous_refresh_expires_at, revoked_at) + VALUES ($1,$2,$3,$4,'gnubok_sk_test','pg-real oauth key',$5,$6,$7,$8,$9,$10,$11)`, + [ + id, + userId, + companyId, + row.key_hash, + row.scopes ?? null, + row.refresh_token_hash ?? null, + row.previous_key_hash ?? null, + row.previous_key_expires_at ?? null, + row.previous_refresh_token_hash ?? null, + row.previous_refresh_expires_at ?? null, + row.revoked ? new Date() : null, + ], + ) + return id +} + +async function validate(keyHash: string) { + const { rows } = await getPool().query( + `SELECT * FROM public.validate_and_increment_api_key($1)`, + [keyHash], + ) + return rows +} + +async function rotate( + presented: string, + newRefresh: string, + newKey: string, + graceSeconds = 120, +) { + const { rows } = await getPool().query<{ outcome: string; scopes: string[] | null }>( + `SELECT * FROM public.rotate_mcp_refresh_token($1,$2,$3,'gnubok_sk_new',$4)`, + [presented, newRefresh, newKey, graceSeconds], + ) + return rows[0] +} + +async function rowById(id: string) { + const { rows } = await getPool().query( + `SELECT key_hash, refresh_token_hash, previous_key_hash, previous_key_expires_at, + previous_refresh_token_hash, previous_refresh_expires_at, revoked_at, request_count + FROM public.api_keys WHERE id = $1`, + [id], + ) + return rows[0] +} + +describe('api_keys rotation grace (issue #710)', () => { + it('adds the four previous_* shadow columns', async () => { + const { rows } = await getPool().query<{ column_name: string; data_type: string }>( + `SELECT column_name, data_type FROM information_schema.columns + WHERE table_schema='public' AND table_name='api_keys' + AND column_name IN ('previous_key_hash','previous_key_expires_at', + 'previous_refresh_token_hash','previous_refresh_expires_at') + ORDER BY column_name`, + ) + const byName = Object.fromEntries(rows.map((r) => [r.column_name, r.data_type])) + expect(byName['previous_key_hash']).toBe('text') + expect(byName['previous_key_expires_at']).toBe('timestamp with time zone') + expect(byName['previous_refresh_token_hash']).toBe('text') + expect(byName['previous_refresh_expires_at']).toBe('timestamp with time zone') + }) + + describe('validate_and_increment_api_key access-token grace', () => { + it('accepts the current key_hash', async () => { + const { userId, companyId } = await seedCompany() + const cur = h('cur') + await insertKey(userId, companyId, { key_hash: cur, refresh_token_hash: h('rt') }) + const rows = await validate(cur) + expect(rows).toHaveLength(1) + expect(rows[0].user_id).toBe(userId) + expect(rows[0].rate_limited).toBe(false) + }) + + it('accepts an unexpired previous_key_hash and increments the resolved row', async () => { + const { userId, companyId } = await seedCompany() + const cur = h('cur') + const prev = h('prev') + const id = await insertKey(userId, companyId, { + key_hash: cur, + refresh_token_hash: h('rt'), + previous_key_hash: prev, + previous_key_expires_at: new Date(Date.now() + 60_000).toISOString(), + }) + // Two grace-hash validations: window resets to 1, then increments to 2 — + // proving the increment is keyed off the resolved row id, not p_key_hash. + expect(await validate(prev)).toHaveLength(1) + expect(await validate(prev)).toHaveLength(1) + expect((await rowById(id)).request_count).toBe(2) + }) + + it('rejects an expired previous_key_hash but still accepts the current one', async () => { + const { userId, companyId } = await seedCompany() + const cur = h('cur') + const prev = h('prev') + await insertKey(userId, companyId, { + key_hash: cur, + refresh_token_hash: h('rt'), + previous_key_hash: prev, + previous_key_expires_at: new Date(Date.now() - 1_000).toISOString(), + }) + expect(await validate(prev)).toHaveLength(0) + expect(await validate(cur)).toHaveLength(1) + }) + + it('never matches a revoked key (current or previous)', async () => { + const { userId, companyId } = await seedCompany() + const cur = h('cur') + const prev = h('prev') + await insertKey(userId, companyId, { + key_hash: cur, + refresh_token_hash: h('rt'), + previous_key_hash: prev, + previous_key_expires_at: new Date(Date.now() + 60_000).toISOString(), + revoked: true, + }) + expect(await validate(cur)).toHaveLength(0) + expect(await validate(prev)).toHaveLength(0) + }) + }) + + describe('rotate_mcp_refresh_token', () => { + it('rotates a current refresh token and demotes the consumed pair to previous_*', async () => { + const { userId, companyId } = await seedCompany() + const r1 = h('r1') + const k1 = h('k1') + const id = await insertKey(userId, companyId, { + key_hash: k1, + refresh_token_hash: r1, + scopes: ['transactions:read', 'invoices:read'], + }) + const r2 = h('r2') + const k2 = h('k2') + const res = await rotate(r1, r2, k2) + expect(res.outcome).toBe('rotated') + expect(res.scopes).toEqual(['transactions:read', 'invoices:read']) + + const row = await rowById(id) + expect(row.refresh_token_hash).toBe(r2) + expect(row.key_hash).toBe(k2) + expect(row.previous_refresh_token_hash).toBe(r1) + expect(row.previous_key_hash).toBe(k1) + expect(new Date(row.previous_refresh_expires_at).getTime()).toBeGreaterThan(Date.now()) + expect(new Date(row.previous_key_expires_at).getTime()).toBeGreaterThan(Date.now()) + + // The just-superseded access key still validates within grace. + expect(await validate(k1)).toHaveLength(1) + }) + + it("replays idempotently when the PREVIOUS refresh token is presented within grace (issue #710)", async () => { + const { userId, companyId } = await seedCompany() + const r1 = h('r1') + const k1 = h('k1') + const id = await insertKey(userId, companyId, { + key_hash: k1, + refresh_token_hash: r1, + scopes: ['reports:read'], + }) + // First, a normal rotation r1 -> r2 (client receives r2 but mis-persists, + // keeping r1). + expect((await rotate(r1, h('r2'), h('k2'))).outcome).toBe('rotated') + + // Client retries with the stale r1 — must NOT 400; it replays and gets a + // fresh current pair, and the grace window slides. + const r3 = h('r3') + const k3 = h('k3') + const replay = await rotate(r1, r3, k3) + expect(replay.outcome).toBe('replayed') + expect(replay.scopes).toEqual(['reports:read']) + + const row = await rowById(id) + expect(row.refresh_token_hash).toBe(r3) + expect(row.key_hash).toBe(k3) + expect(row.previous_refresh_token_hash).toBe(r1) // preserved, not chained + expect(row.revoked_at).toBeNull() + }) + + it('treats reuse of a previous refresh token AFTER the grace window as a breach and revokes the grant', async () => { + const { userId, companyId } = await seedCompany() + const r1 = h('r1') + const k1 = h('k1') + const id = await insertKey(userId, companyId, { key_hash: k1, refresh_token_hash: r1 }) + expect((await rotate(r1, h('r2'), h('k2'))).outcome).toBe('rotated') + + // Force the previous-refresh grace to have expired. + await getPool().query( + `UPDATE public.api_keys SET previous_refresh_expires_at = now() - interval '1 second' WHERE id = $1`, + [id], + ) + + const res = await rotate(r1, h('r3'), h('k3')) + expect(res.outcome).toBe('reuse_revoked') + + const row = await rowById(id) + expect(row.revoked_at).not.toBeNull() + // The whole grant family is dead — neither current nor previous validates. + expect(await validate(row.key_hash)).toHaveLength(0) + }) + + it("returns 'invalid' for an unknown refresh token", async () => { + const res = await rotate(h('nope'), h('r2'), h('k2')) + expect(res.outcome).toBe('invalid') + }) + + it("returns 'revoked' when the matched grant is already revoked", async () => { + const { userId, companyId } = await seedCompany() + const r1 = h('r1') + await insertKey(userId, companyId, { + key_hash: h('k1'), + refresh_token_hash: r1, + revoked: true, + }) + const res = await rotate(r1, h('r2'), h('k2')) + expect(res.outcome).toBe('revoked') + }) + }) +})