diff --git a/.changeset/calm-oauth-custody.md b/.changeset/calm-oauth-custody.md new file mode 100644 index 0000000000..c810c78103 --- /dev/null +++ b/.changeset/calm-oauth-custody.md @@ -0,0 +1,5 @@ +--- +"@agent-native/core": patch +--- + +Add reusable owner- and resource-bound OAuth credential lifecycle primitives with concurrency-safe refresh, revocation, and explicit connection states. diff --git a/packages/core/src/extensions/url-safety.spec.ts b/packages/core/src/extensions/url-safety.spec.ts index ccec286d23..86b9e11af3 100644 --- a/packages/core/src/extensions/url-safety.spec.ts +++ b/packages/core/src/extensions/url-safety.spec.ts @@ -130,6 +130,21 @@ describe("ssrfSafeFetch per-hop policies", () => { expect(redirectResponse.bodyUsed).toBe(true); }); + it("can return a validated redirect for a caller with its own redirect policy", async () => { + const redirectResponse = new Response("moved", { + status: 302, + headers: { location: httpOrigin }, + }); + const fetchMock = vi.fn(async () => redirectResponse); + vi.stubGlobal("fetch", fetchMock); + + await expect( + ssrfSafeFetch(httpsOrigin, {}, { followRedirects: false }), + ).resolves.toBe(redirectResponse); + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(redirectResponse.bodyUsed).toBe(false); + }); + it("allows configured loopback aliases without allowing an unconfigured port", async () => { const fetchMock = vi.fn(async () => new Response("ok", { status: 200 })); vi.stubGlobal("fetch", fetchMock); diff --git a/packages/core/src/extensions/url-safety.ts b/packages/core/src/extensions/url-safety.ts index 563dc712e5..8f0ad4c17c 100644 --- a/packages/core/src/extensions/url-safety.ts +++ b/packages/core/src/extensions/url-safety.ts @@ -331,6 +331,7 @@ export async function ssrfSafeFetch( init: RequestInit = {}, options: { maxRedirects?: number; + followRedirects?: boolean; httpsOnly?: boolean; assertUrlAllowed?: (url: string) => void | Promise; /** @@ -387,6 +388,7 @@ export async function ssrfSafeFetch( const response = await fetch(currentUrl, fetchOpts); if (response.status >= 300 && response.status < 400) { + if (options.followRedirects === false) return response; const location = response.headers.get("location"); if (!location) return response; // Drain the redirect body so the hop's connection is released instead diff --git a/packages/core/src/mcp-client/oauth-client.spec.ts b/packages/core/src/mcp-client/oauth-client.spec.ts index d5ff4efca6..d204cbc409 100644 --- a/packages/core/src/mcp-client/oauth-client.spec.ts +++ b/packages/core/src/mcp-client/oauth-client.spec.ts @@ -6,6 +6,9 @@ const validateAuthorizationResponseIssuerMock = vi.hoisted(() => vi.fn()); const deleteOAuthTokensMock = vi.hoisted(() => vi.fn()); const getOAuthTokensMock = vi.hoisted(() => vi.fn()); const saveOAuthTokensMock = vi.hoisted(() => vi.fn()); +const replaceOAuthTokensIfRevisionMock = vi.hoisted(() => vi.fn()); +const deleteOAuthTokensIfRevisionMock = vi.hoisted(() => vi.fn()); +const ssrfSafeFetchMock = vi.hoisted(() => vi.fn()); vi.mock("@modelcontextprotocol/client", () => ({ auth: authMock, @@ -17,6 +20,29 @@ vi.mock("../oauth-tokens/store.js", () => ({ deleteOAuthTokens: deleteOAuthTokensMock, getOAuthTokens: getOAuthTokensMock, saveOAuthTokens: saveOAuthTokensMock, + getOAuthTokenSnapshot: vi.fn( + async (provider: string, accountId: string, owner: string) => { + const tokens = await getOAuthTokensMock(provider, accountId, owner); + return tokens ? { tokens, owner, revision: 1 } : null; + }, + ), + replaceOAuthTokensIfRevision: replaceOAuthTokensIfRevisionMock, + deleteOAuthTokensIfRevision: deleteOAuthTokensIfRevisionMock, +})); + +vi.mock("../settings/store.js", () => ({ + mutateSetting: vi.fn( + async ( + _key: string, + updater: ( + current: Record | null, + ) => Record | Promise>, + ) => updater(null), + ), +})); + +vi.mock("../extensions/url-safety.js", () => ({ + ssrfSafeFetch: ssrfSafeFetchMock, })); import { @@ -25,6 +51,7 @@ import { getMcpOAuthAccessToken, McpOAuthClientProvider, readMcpOAuthCredentials, + revokeMcpOAuthCredentials, saveMcpOAuthCredentials, startMcpOAuthAuthorization, tokenExpiresAt, @@ -60,11 +87,37 @@ const credentials = { }; beforeEach(() => { + vi.restoreAllMocks(); authMock.mockReset(); refreshAuthorizationMock.mockReset(); deleteOAuthTokensMock.mockReset(); getOAuthTokensMock.mockReset(); saveOAuthTokensMock.mockReset(); + replaceOAuthTokensIfRevisionMock.mockReset(); + replaceOAuthTokensIfRevisionMock.mockImplementation( + async ( + _provider: string, + _accountId: string, + _owner: string, + _revision: number, + tokens: Record, + ) => { + getOAuthTokensMock.mockResolvedValue(tokens); + saveOAuthTokensMock(_provider, _accountId, tokens, _owner); + return true; + }, + ); + deleteOAuthTokensIfRevisionMock.mockReset(); + deleteOAuthTokensIfRevisionMock.mockImplementation( + async (provider: string, accountId: string, owner: string) => { + deleteOAuthTokensMock(provider, accountId, owner); + getOAuthTokensMock.mockResolvedValue(null); + return true; + }, + ); + ssrfSafeFetchMock.mockImplementation((url: string, init?: RequestInit) => + fetch(url, init), + ); validateAuthorizationResponseIssuerMock.mockReset(); }); @@ -142,6 +195,7 @@ describe("MCP OAuth client", () => { await expect( fetchFn!("https://127.0.0.1/.well-known/oauth-authorization-server"), ).rejects.toThrow(/private\/internal address/); + expect(ssrfSafeFetchMock).not.toHaveBeenCalled(); const provider = new McpOAuthClientProvider({ serverUrl: "https://mcp.example.com/mcp", redirectUrl: "https://app.example.com/callback", @@ -154,6 +208,44 @@ describe("MCP OAuth client", () => { ).toThrow(/private\/internal address/); }); + it("routes OAuth discovery through the DNS-aware SSRF guard", async () => { + let fetchFn: + | ((url: string | URL, init?: RequestInit) => Promise) + | undefined; + authMock.mockImplementationOnce( + async ( + provider: McpOAuthClientProvider, + options: { fetchFn?: typeof fetchFn }, + ) => { + fetchFn = options.fetchFn; + provider.saveClientInformation(clientInformation as any); + provider.saveCodeVerifier(""); + provider.redirectToAuthorization( + new URL("https://auth.example.com/authorize"), + ); + return "REDIRECT"; + }, + ); + ssrfSafeFetchMock.mockResolvedValueOnce(new Response("ok")); + + await startMcpOAuthAuthorization({ + serverUrl: "https://mcp.example.com/mcp", + redirectUrl: "https://app.example.com/callback", + state: "", + }); + await fetchFn!("https://auth.example.com/discovery"); + + expect(ssrfSafeFetchMock).toHaveBeenCalledWith( + "https://auth.example.com/discovery", + expect.objectContaining({ redirect: "manual" }), + expect.objectContaining({ + maxRedirects: 0, + followRedirects: false, + allowedPrivateOrigins: [], + }), + ); + }); + it("validates every OAuth redirect hop and strips credentials across origins", async () => { let fetchFn: | ((url: string | URL, init?: RequestInit) => Promise) @@ -302,7 +394,15 @@ describe("MCP OAuth client", () => { expect(saveOAuthTokensMock).toHaveBeenCalledWith( "mcp", "mcp_oauth:test", - credentials, + expect.objectContaining({ + ...credentials, + oauthLifecycle: { + version: 1, + provider: "mcp", + resource: "https://mcp.example.com/mcp", + owner: "user:alice@example.com", + }, + }), "user:alice@example.com", ); }); @@ -312,7 +412,7 @@ describe("MCP OAuth client", () => { ...credentials, tokenExpiresAt: Date.now() - 1, }; - getOAuthTokensMock.mockResolvedValueOnce(expiring); + getOAuthTokensMock.mockResolvedValue(expiring); refreshAuthorizationMock.mockResolvedValueOnce({ access_token: "", token_type: "bearer", @@ -341,7 +441,7 @@ describe("MCP OAuth client", () => { }); it("requires reauthorization for expiring legacy credentials without issuer binding", async () => { - getOAuthTokensMock.mockResolvedValueOnce({ + getOAuthTokensMock.mockResolvedValue({ ...credentials, clientInformation: { client_id: "legacy-client", @@ -361,7 +461,7 @@ describe("MCP OAuth client", () => { }); it("does not return an expired token when refresh fails", async () => { - getOAuthTokensMock.mockResolvedValueOnce({ + getOAuthTokensMock.mockResolvedValue({ ...credentials, tokenExpiresAt: Date.now() - 1, }); @@ -380,7 +480,7 @@ describe("MCP OAuth client", () => { }); it("keeps a still-valid token when an early refresh fails", async () => { - getOAuthTokensMock.mockResolvedValueOnce({ + getOAuthTokensMock.mockResolvedValue({ ...credentials, tokenExpiresAt: Date.now() + 30_000, }); @@ -398,6 +498,93 @@ describe("MCP OAuth client", () => { ).resolves.toBe(""); }); + it("does not return a legacy MCP token for a different resource", async () => { + getOAuthTokensMock.mockResolvedValue(credentials); + + await expect( + getMcpOAuthAccessToken({ + key: "mcp_oauth:test", + scope: "user", + scopeId: "alice@example.com", + serverUrl: "https://other.example.com/mcp", + }), + ).resolves.toBeNull(); + expect(refreshAuthorizationMock).not.toHaveBeenCalled(); + }); + + it("revokes the refresh token before deleting local MCP custody", async () => { + getOAuthTokensMock.mockResolvedValue({ + ...credentials, + discoveryState: { + ...credentials.discoveryState, + authorizationServerMetadata: { + ...credentials.discoveryState.authorizationServerMetadata, + revocation_endpoint: "https://auth.example.com/revoke", + }, + }, + }); + const fetchMock = vi + .spyOn(globalThis, "fetch") + .mockResolvedValueOnce(new Response(null, { status: 200 })); + + await expect( + revokeMcpOAuthCredentials({ + key: "mcp_oauth:test", + scope: "user", + scopeId: "alice@example.com", + serverUrl: "https://mcp.example.com/mcp", + }), + ).resolves.toEqual({ remote: "succeeded", local: "deleted" }); + + expect(String(fetchMock.mock.calls[0]?.[0])).toBe( + "https://auth.example.com/revoke", + ); + const request = fetchMock.mock.calls[0]?.[1] as RequestInit; + const body = new URLSearchParams(String(request.body)); + expect(body.get("token")).toBe(""); + expect(body.get("token_type_hint")).toBe("refresh_token"); + expect(body.get("client_id")).toBe("mcp-client-test"); + expect(deleteOAuthTokensIfRevisionMock).toHaveBeenCalledTimes(1); + expect(ssrfSafeFetchMock).toHaveBeenCalledWith( + "https://auth.example.com/revoke", + expect.any(Object), + expect.objectContaining({ maxRedirects: 0, httpsOnly: true }), + ); + }); + + it("fails closed instead of posting a token to a loopback revocation endpoint", async () => { + getOAuthTokensMock.mockResolvedValue({ + ...credentials, + discoveryState: { + ...credentials.discoveryState, + authorizationServerMetadata: { + ...credentials.discoveryState.authorizationServerMetadata, + revocation_endpoint: "http://127.0.0.1:9000/revoke", + }, + }, + }); + ssrfSafeFetchMock.mockRejectedValueOnce( + new Error("SSRF blocked: refusing to fetch private/internal address"), + ); + const fetchMock = vi.spyOn(globalThis, "fetch"); + + await expect( + revokeMcpOAuthCredentials({ + key: "mcp_oauth:test", + scope: "user", + scopeId: "alice@example.com", + serverUrl: "https://mcp.example.com/mcp", + }), + ).resolves.toEqual({ remote: "failed", local: "deleted" }); + + expect(ssrfSafeFetchMock).toHaveBeenCalledWith( + "http://127.0.0.1:9000/revoke", + expect.any(Object), + expect.objectContaining({ maxRedirects: 0, httpsOnly: true }), + ); + expect(fetchMock).not.toHaveBeenCalled(); + }); + it("rejects malformed stored bundles", async () => { getOAuthTokensMock.mockResolvedValueOnce({ access_token: "" }); @@ -406,13 +593,14 @@ describe("MCP OAuth client", () => { key: "mcp_oauth:test", scope: "user", scopeId: "alice@example.com", + serverUrl: "https://mcp.example.com/mcp", }), ).resolves.toBeNull(); }); - it("binds reads and deletes to the credential owner", async () => { - getOAuthTokensMock.mockResolvedValueOnce(null); - deleteOAuthTokensMock.mockResolvedValueOnce(1); + it("preserves owner-bound legacy reads and deletes without a serverUrl argument", async () => { + getOAuthTokensMock.mockResolvedValue(credentials); + deleteOAuthTokensMock.mockResolvedValue(1); await readMcpOAuthCredentials({ key: "mcp_oauth:test", diff --git a/packages/core/src/mcp-client/oauth-client.ts b/packages/core/src/mcp-client/oauth-client.ts index 269bf78cae..6d2248d327 100644 --- a/packages/core/src/mcp-client/oauth-client.ts +++ b/packages/core/src/mcp-client/oauth-client.ts @@ -22,15 +22,23 @@ import { OAuthTokens, } from "@modelcontextprotocol/client"; +import { ssrfSafeFetch } from "../extensions/url-safety.js"; import { - deleteOAuthTokens, - getOAuthTokens, - saveOAuthTokens, -} from "../oauth-tokens/store.js"; + readOAuthCredentialState, + resolveOAuthCredentialAccess, + revokeOAuthCredential, + saveOAuthCredential, + type OAuthCredential, + type OAuthCredentialIdentity, + type OAuthCredentialState, + type OAuthRevocationResult, +} from "../oauth-tokens/lifecycle.js"; +import { deleteOAuthTokens, getOAuthTokens } from "../oauth-tokens/store.js"; import { validateRemoteUrl } from "./remote-url.js"; const TOKEN_EXPIRY_SKEW_MS = 60_000; const MAX_OAUTH_REDIRECTS = 5; +const MCP_OAUTH_PRIVATE_ORIGINS_ENV = "AGENT_NATIVE_MCP_OAUTH_PRIVATE_ORIGINS"; type GuardedFetch = ( url: string | URL, @@ -48,12 +56,22 @@ function checkedRemoteUrl(value: string | URL, label: string): URL { } function guardedOAuthFetch(): GuardedFetch { + const allowedPrivateOrigins = ( + process.env[MCP_OAUTH_PRIVATE_ORIGINS_ENV] ?? "" + ) + .split(",") + .map((origin) => origin.trim()) + .filter(Boolean); return async (url, init) => { let currentUrl = checkedRemoteUrl(url, "request"); let currentInit: RequestInit = { ...init, redirect: "manual" }; for (let redirectCount = 0; ; redirectCount += 1) { - const response = await fetch(currentUrl, currentInit); + const response = await ssrfSafeFetch(currentUrl.href, currentInit, { + maxRedirects: 0, + followRedirects: false, + allowedPrivateOrigins, + }); if (![301, 302, 303, 307, 308].includes(response.status)) { return response; } @@ -100,6 +118,23 @@ function guardedOAuthFetch(): GuardedFetch { }; } +async function guardedRevocationFetch( + url: string, + init: RequestInit, +): Promise { + const endpoint = checkedRemoteUrl(url, "revocation endpoint"); + return ssrfSafeFetch(endpoint.href, init, { + maxRedirects: 0, + httpsOnly: true, + assertUrlAllowed: (candidate) => { + const checked = checkedRemoteUrl(candidate, "revocation redirect"); + if (checked.origin !== endpoint.origin) { + throw new Error("MCP OAuth revocation redirect changed origin."); + } + }, + }); +} + function validateDiscoveryUrls(state: { authorizationServerUrl: string; authorizationServerMetadata?: AuthorizationServerMetadata; @@ -129,13 +164,9 @@ function validateDiscoveryUrls(state: { } } -function credentialOwner(options: { scope: "user" | "org"; scopeId: string }) { - return `${options.scope}:${options.scopeId}`; -} - export type McpOAuthDiscoveryState = OAuthDiscoveryState; -export interface McpOAuthCredentialBundle { +export interface McpOAuthCredentialBundle extends OAuthCredential { serverUrl: string; clientInformation: StoredOAuthClientInformation; discoveryState?: McpOAuthDiscoveryState; @@ -171,6 +202,20 @@ function issuerForDiscovery( return typeof issuer === "string" && issuer ? issuer : undefined; } +function credentialIdentity(options: { + key: string; + scope: "user" | "org"; + scopeId: string; + serverUrl: string; +}): OAuthCredentialIdentity { + return { + provider: "mcp", + accountId: options.key, + resource: options.serverUrl, + owner: { scope: options.scope, id: options.scopeId }, + }; +} + function withIssuer( value: T, context: OAuthClientInformationContext | undefined, @@ -421,11 +466,13 @@ export async function saveMcpOAuthCredentials(options: { if (options.credentials.discoveryState) { validateDiscoveryUrls(options.credentials.discoveryState); } - await saveOAuthTokens( - "mcp", - options.key, - options.credentials as unknown as Record, - `${options.scope}:${options.scopeId}`, + await saveOAuthCredential( + credentialIdentity({ + ...options, + serverUrl: options.credentials.serverUrl, + }), + options.credentials, + { legacyAccountKey: true }, ); } @@ -433,22 +480,33 @@ export async function readMcpOAuthCredentials(options: { key: string; scope: "user" | "org"; scopeId: string; + serverUrl?: string; }): Promise { - const stored = await getOAuthTokens( - "mcp", - options.key, - credentialOwner(options), - ); - if (!stored) return null; - const parsed = stored as Partial; - if ( - typeof parsed.serverUrl !== "string" || - !parsed.clientInformation || - !parsed.tokens || - typeof parsed.tokens.access_token !== "string" - ) { - return null; + if (!options.serverUrl) { + const stored = await getOAuthTokens( + "mcp", + options.key, + `${options.scope}:${options.scope === "user" ? options.scopeId.toLowerCase() : options.scopeId}`, + ); + if (!stored) return null; + const parsed = stored as Partial; + if ( + typeof parsed.serverUrl !== "string" || + !parsed.clientInformation || + !parsed.tokens || + typeof parsed.tokens.access_token !== "string" + ) { + return null; + } + return parsed as McpOAuthCredentialBundle; } + const state = await getMcpOAuthConnectionState({ + ...options, + serverUrl: options.serverUrl, + }); + if (state.kind === "missing" || state.kind === "malformed") return null; + const parsed = state.credential; + if (parsed.serverUrl !== options.serverUrl) return null; if (!validateRemoteUrl(parsed.serverUrl).ok) return null; if (parsed.discoveryState) { try { @@ -457,17 +515,101 @@ export async function readMcpOAuthCredentials(options: { return null; } } - return parsed as McpOAuthCredentialBundle; + return parsed; +} + +export async function getMcpOAuthConnectionState(options: { + key: string; + scope: "user" | "org"; + scopeId: string; + serverUrl: string; +}): Promise> { + return readOAuthCredentialState( + credentialIdentity(options), + { + allowLegacy: true, + legacyAccountKey: true, + validateCredential: (credential) => + credential.serverUrl === options.serverUrl, + }, + ); } export async function deleteMcpOAuthCredentials(options: { key: string; scope: "user" | "org"; scopeId: string; + serverUrl?: string; }): Promise { - return ( - (await deleteOAuthTokens("mcp", options.key, credentialOwner(options))) > 0 - ); + if (!options.serverUrl) { + return ( + (await deleteOAuthTokens( + "mcp", + options.key, + `${options.scope}:${options.scope === "user" ? options.scopeId.toLowerCase() : options.scopeId}`, + )) > 0 + ); + } + const current = await readMcpOAuthCredentials(options); + if (!current) return false; + const identity = credentialIdentity({ + ...options, + serverUrl: options.serverUrl, + }); + const result = await revokeOAuthCredential(identity, { + allowLegacy: true, + legacyAccountKey: true, + validateCredential: (credential: McpOAuthCredentialBundle) => + credential.serverUrl === options.serverUrl, + }); + return result.local === "deleted"; +} + +export async function revokeMcpOAuthCredentials(options: { + key: string; + scope: "user" | "org"; + scopeId: string; + serverUrl: string; +}): Promise { + const identity = credentialIdentity(options); + const current = await readMcpOAuthCredentials(options); + if (!current) { + return { remote: "not_attempted", local: "missing" }; + } + return revokeOAuthCredential(identity, { + allowLegacy: true, + legacyAccountKey: true, + validateCredential: (credential) => + credential.serverUrl === options.serverUrl, + revoke: async ({ credential }) => { + const endpoint = ( + credential.discoveryState?.authorizationServerMetadata as + | (AuthorizationServerMetadata & { revocation_endpoint?: string }) + | undefined + )?.revocation_endpoint; + if (!endpoint) return "unsupported"; + const token = + credential.tokens.refresh_token ?? credential.tokens.access_token; + const body = new URLSearchParams({ + token, + token_type_hint: credential.tokens.refresh_token + ? "refresh_token" + : "access_token", + client_id: credential.clientInformation.client_id, + }); + const response = await guardedRevocationFetch(endpoint, { + method: "POST", + headers: { "content-type": "application/x-www-form-urlencoded" }, + body, + }); + if (!response.ok) { + await response.body?.cancel().catch(() => undefined); + throw new Error("MCP OAuth revocation failed."); + } + await response.body?.cancel().catch(() => undefined); + return "succeeded"; + }, + }); } /** @@ -482,72 +624,60 @@ export async function getMcpOAuthAccessToken(options: { serverUrl: string; }): Promise { if (!validateRemoteUrl(options.serverUrl).ok) return null; - const credentials = await readMcpOAuthCredentials(options); - if (!credentials || credentials.serverUrl !== options.serverUrl) return null; - - const accessToken = credentials.tokens.access_token; - const now = Date.now(); - if ( - typeof credentials.tokenExpiresAt !== "number" || - credentials.tokenExpiresAt - now > TOKEN_EXPIRY_SKEW_MS - ) { - return accessToken; - } - const tokenIsExpired = credentials.tokenExpiresAt <= now; - - const refreshToken = credentials.tokens.refresh_token; - const discovery = credentials.discoveryState; - if (!refreshToken || !discovery?.authorizationServerUrl) { - return tokenIsExpired ? null : accessToken; - } - - try { - const expectedIssuer = issuerForDiscovery(discovery); - if ( - !expectedIssuer || - credentials.clientInformation.issuer !== expectedIssuer || - credentials.tokens.issuer !== expectedIssuer - ) { - return null; - } - const resource = discovery.resourceMetadata?.resource - ? checkedRemoteUrl(discovery.resourceMetadata.resource, "resource") - : undefined; - const authorizationServerUrl = checkedRemoteUrl( - discovery.authorizationServerUrl, - "authorization server", - ); - const refreshed = await refreshAuthorization(authorizationServerUrl, { - metadata: discovery.authorizationServerMetadata, - clientInformation: credentials.clientInformation, - refreshToken, - resource, - fetchFn: guardedOAuthFetch(), - }); - const nextTokens: StoredOAuthTokens = { - ...credentials.tokens, - ...refreshed, - issuer: expectedIssuer, - ...(refreshed.refresh_token - ? { refresh_token: refreshed.refresh_token } - : credentials.tokens.refresh_token - ? { refresh_token: credentials.tokens.refresh_token } - : {}), - }; - const next: McpOAuthCredentialBundle = { - ...credentials, - tokens: nextTokens, - tokenExpiresAt: tokenExpiresAt(nextTokens), - }; - await saveMcpOAuthCredentials({ - ...options, - credentials: next, - }); - return nextTokens.access_token; - } catch { - // A still-valid token can survive a transient refresh failure. Once it has - // expired, omit it so callers surface reauthorization instead of retrying - // a credential that can never authenticate. - return tokenIsExpired ? null : accessToken; - } + const result = await resolveOAuthCredentialAccess( + credentialIdentity(options), + { + allowLegacy: true, + legacyAccountKey: true, + validateCredential: (credential) => + credential.serverUrl === options.serverUrl, + expirySkewMs: TOKEN_EXPIRY_SKEW_MS, + refresh: async ({ credential: credentials }) => { + const refreshToken = credentials.tokens.refresh_token; + const discovery = credentials.discoveryState; + if (!refreshToken || !discovery?.authorizationServerUrl) { + throw new Error("MCP OAuth refresh is unavailable."); + } + const expectedIssuer = issuerForDiscovery(discovery); + if ( + !expectedIssuer || + credentials.clientInformation.issuer !== expectedIssuer || + credentials.tokens.issuer !== expectedIssuer + ) { + throw new Error("MCP OAuth refresh issuer binding is invalid."); + } + const resource = discovery.resourceMetadata?.resource + ? checkedRemoteUrl(discovery.resourceMetadata.resource, "resource") + : undefined; + const authorizationServerUrl = checkedRemoteUrl( + discovery.authorizationServerUrl, + "authorization server", + ); + const refreshed = await refreshAuthorization(authorizationServerUrl, { + metadata: discovery.authorizationServerMetadata, + clientInformation: credentials.clientInformation, + refreshToken, + resource, + fetchFn: guardedOAuthFetch(), + }); + const nextTokens: StoredOAuthTokens = { + ...credentials.tokens, + ...refreshed, + issuer: expectedIssuer, + ...(refreshed.refresh_token + ? { refresh_token: refreshed.refresh_token } + : credentials.tokens.refresh_token + ? { refresh_token: credentials.tokens.refresh_token } + : {}), + }; + const next: McpOAuthCredentialBundle = { + ...credentials, + tokens: nextTokens, + tokenExpiresAt: tokenExpiresAt(nextTokens), + }; + return next; + }, + }, + ); + return result.accessToken; } diff --git a/packages/core/src/mcp-client/remote-store.ts b/packages/core/src/mcp-client/remote-store.ts index 68a119c78c..4c8dca85ab 100644 --- a/packages/core/src/mcp-client/remote-store.ts +++ b/packages/core/src/mcp-client/remote-store.ts @@ -42,6 +42,7 @@ import type { McpHttpServerConfig } from "./config.js"; import { deleteMcpOAuthCredentials, getMcpOAuthAccessToken, + revokeMcpOAuthCredentials, saveMcpOAuthCredentials, type McpOAuthCredentialBundle, } from "./oauth-client.js"; @@ -270,6 +271,7 @@ export async function addOAuthRemoteServer( key: oauthSecretKey, scope, scopeId, + serverUrl: input.credentials.serverUrl, }); } return result; @@ -278,6 +280,7 @@ export async function addOAuthRemoteServer( key: oauthSecretKey, scope, scopeId, + serverUrl: input.credentials.serverUrl, }).catch(() => {}); return { ok: false, @@ -496,11 +499,17 @@ export async function removeRemoteServer( } if (removed?.oauthSecretKey) { try { - await deleteMcpOAuthCredentials({ + const result = await revokeMcpOAuthCredentials({ key: removed.oauthSecretKey, scope, scopeId, + serverUrl: removed.url, }); + if (result.remote === "failed") { + console.warn( + `[mcp-client] MCP OAuth revocation failed for ${removed.name}; local credentials were removed.`, + ); + } } catch (err: any) { // eslint-disable-next-line no-console console.warn( diff --git a/packages/core/src/oauth-tokens/index.ts b/packages/core/src/oauth-tokens/index.ts index 99804f3041..b19cf52613 100644 --- a/packages/core/src/oauth-tokens/index.ts +++ b/packages/core/src/oauth-tokens/index.ts @@ -1,14 +1,29 @@ export { getOAuthTokens, + getOAuthTokenSnapshot, OAuthAccountOwnedByOtherUserError, saveOAuthTokens, + replaceOAuthTokensIfRevision, deleteOAuthTokens, + deleteOAuthTokensIfRevision, listOAuthAccounts, listOAuthAccountsByOwner, hasOAuthTokens, setOAuthDisplayName, } from "./store.js"; +export { + readOAuthCredentialState, + resolveOAuthCredentialAccess, + revokeOAuthCredential, + saveOAuthCredential, + type OAuthCredential, + type OAuthCredentialIdentity, + type OAuthCredentialOwner, + type OAuthCredentialState, + type OAuthRevocationResult, +} from "./lifecycle.js"; + export { refreshExpiringGoogleTokens, startGoogleTokenRefreshLoop, diff --git a/packages/core/src/oauth-tokens/lifecycle.spec.ts b/packages/core/src/oauth-tokens/lifecycle.spec.ts new file mode 100644 index 0000000000..8e1ef31dbc --- /dev/null +++ b/packages/core/src/oauth-tokens/lifecycle.spec.ts @@ -0,0 +1,454 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const state = vi.hoisted(() => ({ + rows: new Map< + string, + { + tokens: Record; + owner: string; + revision: number; + } + >(), + settings: new Map>(), + revision: 1, +})); + +function rowKey(provider: string, accountId: string): string { + return `${provider}:${accountId}`; +} + +vi.mock("./store.js", () => ({ + saveOAuthTokens: vi.fn( + async ( + provider: string, + accountId: string, + tokens: Record, + owner: string, + ) => { + state.rows.set(rowKey(provider, accountId), { + tokens, + owner, + revision: state.revision++, + }); + }, + ), + getOAuthTokenSnapshot: vi.fn( + async (provider: string, accountId: string, owner: string) => { + const row = state.rows.get(rowKey(provider, accountId)); + return row?.owner === owner ? structuredClone(row) : null; + }, + ), + replaceOAuthTokensIfRevision: vi.fn( + async ( + provider: string, + accountId: string, + owner: string, + expectedRevision: number, + tokens: Record, + ) => { + const key = rowKey(provider, accountId); + const row = state.rows.get(key); + if (row?.owner !== owner || row.revision !== expectedRevision) { + return false; + } + state.rows.set(key, { + tokens, + owner, + revision: state.revision++, + }); + return true; + }, + ), + deleteOAuthTokensIfRevision: vi.fn( + async ( + provider: string, + accountId: string, + owner: string, + expectedRevision: number, + ) => { + const key = rowKey(provider, accountId); + const row = state.rows.get(key); + if (row?.owner !== owner || row.revision !== expectedRevision) { + return false; + } + state.rows.delete(key); + return true; + }, + ), +})); + +vi.mock("../settings/store.js", () => ({ + mutateSetting: vi.fn( + async ( + key: string, + updater: ( + current: Record | null, + ) => Record | Promise>, + ) => { + const next = await updater(state.settings.get(key) ?? null); + state.settings.set(key, structuredClone(next)); + return structuredClone(next); + }, + ), +})); + +import { + readOAuthCredentialState, + resolveOAuthCredentialAccess, + revokeOAuthCredential, + saveOAuthCredential, + type OAuthCredential, + type OAuthCredentialIdentity, +} from "./lifecycle.js"; + +const identity: OAuthCredentialIdentity = { + provider: "builder", + accountId: "managed-ai", + resource: "https://api.builder.io", + owner: { scope: "user", id: "Alice@Example.com" }, +}; + +function credential( + options: { + access?: string; + refresh?: string; + expiresAt?: number; + } = {}, +): OAuthCredential { + return { + tokens: { + access_token: options.access ?? "", + ...(options.refresh === undefined + ? { refresh_token: "" } + : options.refresh + ? { refresh_token: options.refresh } + : {}), + }, + tokenExpiresAt: options.expiresAt ?? Date.now() + 3_600_000, + }; +} + +beforeEach(() => { + state.rows.clear(); + state.settings.clear(); + state.revision = 1; + vi.clearAllMocks(); +}); + +describe("OAuth credential lifecycle", () => { + it("binds custody to provider, resource, and normalized owner", async () => { + await saveOAuthCredential(identity, credential()); + + await expect(readOAuthCredentialState(identity)).resolves.toMatchObject({ + kind: "connected", + credential: { + oauthLifecycle: { + version: 1, + provider: "builder", + resource: "https://api.builder.io", + owner: "user:alice@example.com", + }, + }, + }); + await expect( + readOAuthCredentialState({ + ...identity, + resource: "https://mcp.builder.io/mcp/fusion", + }), + ).resolves.toEqual({ kind: "missing" }); + await expect( + readOAuthCredentialState({ + ...identity, + owner: { scope: "user", id: "bob@example.com" }, + }), + ).resolves.toEqual({ kind: "missing" }); + }); + + it("keeps credentials for two resources with the same provider and account independently retrievable", async () => { + const fusionIdentity = { + ...identity, + resource: "https://api.builder.io/mcp/fusion", + }; + await saveOAuthCredential( + identity, + credential({ access: "" }), + ); + await saveOAuthCredential( + fusionIdentity, + credential({ access: "" }), + ); + + await expect(readOAuthCredentialState(identity)).resolves.toMatchObject({ + kind: "connected", + credential: { tokens: { access_token: "" } }, + }); + await expect( + readOAuthCredentialState(fusionIdentity), + ).resolves.toMatchObject({ + kind: "connected", + credential: { tokens: { access_token: "" } }, + }); + expect(state.rows.size).toBe(2); + }); + + it("distinguishes missing, malformed, expired, and reconnect-required custody", async () => { + await expect(readOAuthCredentialState(identity)).resolves.toEqual({ + kind: "missing", + }); + + await saveOAuthCredential(identity, credential()); + const malformedRow = [...state.rows.values()][0]; + malformedRow.tokens = { tokens: {} }; + await expect(readOAuthCredentialState(identity)).resolves.toMatchObject({ + kind: "malformed", + }); + + await saveOAuthCredential( + identity, + credential({ expiresAt: Date.now() - 1 }), + ); + await expect(readOAuthCredentialState(identity)).resolves.toMatchObject({ + kind: "expired", + }); + + const row = [...state.rows.values()][0]; + row.tokens = { + ...row.tokens, + oauthLifecycle: { + ...(row.tokens.oauthLifecycle as Record), + reconnectReason: "refresh_failed", + }, + }; + await expect(readOAuthCredentialState(identity)).resolves.toMatchObject({ + kind: "reconnect_required", + }); + }); + + it("redeems one rotating refresh token and makes concurrent waiters reload the winner", async () => { + await saveOAuthCredential( + identity, + credential({ expiresAt: Date.now() - 1 }), + ); + let finishRefresh!: () => void; + const refreshGate = new Promise((resolve) => { + finishRefresh = resolve; + }); + const refresh = vi.fn(async ({ credential: current }) => { + await refreshGate; + return { + ...current, + tokens: { + ...current.tokens, + access_token: "", + refresh_token: "", + }, + tokenExpiresAt: Date.now() + 3_600_000, + }; + }); + const options = { + refresh, + waitMs: 1, + maxWaitMs: 1_000, + dependencies: { + sleep: (ms: number) => new Promise((r) => setTimeout(r, ms)), + }, + }; + + const first = resolveOAuthCredentialAccess(identity, options); + await Promise.resolve(); + const second = resolveOAuthCredentialAccess(identity, options); + await new Promise((resolve) => setTimeout(resolve, 5)); + finishRefresh(); + + await expect(Promise.all([first, second])).resolves.toEqual([ + expect.objectContaining({ accessToken: "" }), + expect.objectContaining({ accessToken: "" }), + ]); + expect(refresh).toHaveBeenCalledTimes(1); + }); + + it("renews the lease while a slow rotating refresh is in flight", async () => { + await saveOAuthCredential( + identity, + credential({ expiresAt: Date.now() - 1 }), + ); + const refresh = vi.fn(async ({ credential: current }) => { + await new Promise((resolve) => setTimeout(resolve, 40)); + return { + ...current, + tokens: { + ...current.tokens, + access_token: "", + refresh_token: "", + }, + tokenExpiresAt: Date.now() + 3_600_000, + }; + }); + const options = { + refresh, + leaseMs: 12, + waitMs: 1, + maxWaitMs: 1_000, + }; + + const first = resolveOAuthCredentialAccess(identity, options); + await new Promise((resolve) => setTimeout(resolve, 20)); + const second = resolveOAuthCredentialAccess(identity, options); + + await expect(Promise.all([first, second])).resolves.toEqual([ + expect.objectContaining({ accessToken: "" }), + expect.objectContaining({ accessToken: "" }), + ]); + expect(refresh).toHaveBeenCalledTimes(1); + }); + + it("preserves a successful rotated token when the refresh lease is stolen", async () => { + await saveOAuthCredential( + identity, + credential({ expiresAt: Date.now() - 1 }), + ); + let finishRefresh!: () => void; + const refreshGate = new Promise((resolve) => { + finishRefresh = resolve; + }); + const refresh = vi.fn(async ({ credential: current }) => { + await refreshGate; + return { + ...current, + tokens: { + ...current.tokens, + access_token: "", + refresh_token: "", + }, + tokenExpiresAt: Date.now() + 3_600_000, + }; + }); + + const pending = resolveOAuthCredentialAccess(identity, { + refresh, + leaseMs: 12, + waitMs: 1, + maxWaitMs: 1_000, + }); + await new Promise((resolve) => setTimeout(resolve, 5)); + const key = [...state.settings.keys()][0]; + state.settings.set(key, { + holder: "competing-process", + expiresAt: Date.now() + 10_000, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + finishRefresh(); + + await expect(pending).resolves.toMatchObject({ + accessToken: "", + state: { kind: "connected" }, + }); + await expect(readOAuthCredentialState(identity)).resolves.toMatchObject({ + kind: "connected", + credential: { + tokens: { refresh_token: "" }, + }, + }); + }); + + it("reloads the winning rotation instead of marking reconnect after lease loss", async () => { + await saveOAuthCredential( + identity, + credential({ expiresAt: Date.now() - 1 }), + ); + let rejectRefresh!: () => void; + const refreshGate = new Promise((_resolve, reject) => { + rejectRefresh = () => reject(new Error("rotating token already used")); + }); + + const pending = resolveOAuthCredentialAccess(identity, { + refresh: async () => refreshGate, + leaseMs: 12, + waitMs: 1, + maxWaitMs: 1_000, + }); + await new Promise((resolve) => setTimeout(resolve, 5)); + const key = [...state.settings.keys()][0]; + state.settings.set(key, { + holder: "competing-process", + expiresAt: Date.now() + 10_000, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + rejectRefresh(); + await new Promise((resolve) => setTimeout(resolve, 2)); + await saveOAuthCredential( + identity, + credential({ + access: "", + refresh: "", + }), + ); + + await expect(pending).resolves.toMatchObject({ + accessToken: "", + state: { kind: "connected" }, + }); + }); + + it("marks an expired credential for reconnect after refresh fails", async () => { + await saveOAuthCredential( + identity, + credential({ expiresAt: Date.now() - 1 }), + ); + + await expect( + resolveOAuthCredentialAccess(identity, { + refresh: async () => { + throw new Error("refresh token rejected"); + }, + }), + ).resolves.toMatchObject({ + accessToken: null, + state: { kind: "reconnect_required" }, + }); + }); + + it("attempts remote revocation, deletes local custody, and reports failure honestly", async () => { + await saveOAuthCredential(identity, credential()); + + await expect( + revokeOAuthCredential(identity, { + revoke: async () => { + throw new Error("provider unavailable"); + }, + }), + ).resolves.toEqual({ remote: "failed", local: "deleted" }); + await expect(readOAuthCredentialState(identity)).resolves.toEqual({ + kind: "missing", + }); + }); + + it("does not delete a newer authorization that lands during revocation", async () => { + await saveOAuthCredential(identity, credential()); + let finishRevocation!: () => void; + const revocationGate = new Promise((resolve) => { + finishRevocation = resolve; + }); + const revocation = revokeOAuthCredential(identity, { + revoke: async () => { + await revocationGate; + return "succeeded"; + }, + }); + await Promise.resolve(); + await saveOAuthCredential( + identity, + credential({ access: "" }), + ); + finishRevocation(); + + await expect(revocation).resolves.toEqual({ + remote: "succeeded", + local: "replaced", + }); + await expect(readOAuthCredentialState(identity)).resolves.toMatchObject({ + kind: "connected", + credential: { tokens: { access_token: "" } }, + }); + }); +}); diff --git a/packages/core/src/oauth-tokens/lifecycle.ts b/packages/core/src/oauth-tokens/lifecycle.ts new file mode 100644 index 0000000000..bbcc55ba5e --- /dev/null +++ b/packages/core/src/oauth-tokens/lifecycle.ts @@ -0,0 +1,548 @@ +import { createHash, randomUUID } from "node:crypto"; + +import { mutateSetting } from "../settings/store.js"; +import { + deleteOAuthTokensIfRevision, + getOAuthTokenSnapshot, + replaceOAuthTokensIfRevision, + saveOAuthTokens, +} from "./store.js"; + +const DEFAULT_EXPIRY_SKEW_MS = 60_000; +const DEFAULT_LEASE_MS = 15_000; +const DEFAULT_WAIT_MS = 50; +const DEFAULT_MAX_WAIT_MS = 20_000; +const LIFECYCLE_VERSION = 1; + +export interface OAuthCredentialOwner { + scope: "user" | "org"; + id: string; +} + +export interface OAuthCredentialIdentity { + provider: string; + accountId: string; + resource: string; + owner: OAuthCredentialOwner; +} + +export interface OAuthCredentialTokens { + access_token: string; + refresh_token?: string; + [key: string]: unknown; +} + +interface OAuthLifecycleMetadata { + version: 1; + provider: string; + resource: string; + owner: string; + reconnectReason?: "refresh_failed"; +} + +export interface OAuthCredential { + tokens: OAuthCredentialTokens; + tokenExpiresAt?: number; + oauthLifecycle?: OAuthLifecycleMetadata; + [key: string]: unknown; +} + +interface CredentialSnapshot { + credential: T; + revision: number; +} + +export type OAuthCredentialState = + | { kind: "missing" } + | { kind: "malformed"; revision: number } + | ({ + kind: "connected" | "expired" | "reconnect_required"; + } & CredentialSnapshot); + +export interface OAuthCredentialAccessResult< + T extends OAuthCredential = OAuthCredential, +> { + state: OAuthCredentialState; + accessToken: string | null; +} + +export interface OAuthRefreshContext { + identity: OAuthCredentialIdentity; + credential: T; +} + +export interface OAuthRevocationResult { + remote: "succeeded" | "failed" | "unsupported" | "not_attempted"; + local: "deleted" | "missing" | "replaced"; +} + +interface LifecycleDependencies { + now: () => number; + sleep: (ms: number) => Promise; + holderId: () => string; +} + +const defaultDependencies: LifecycleDependencies = { + now: () => Date.now(), + sleep: (ms) => new Promise((resolve) => setTimeout(resolve, ms)), + holderId: () => randomUUID(), +}; + +function ownerKey(owner: OAuthCredentialOwner): string { + const id = owner.id.trim(); + if (!id) throw new Error("OAuth credential owner is required."); + return `${owner.scope}:${owner.scope === "user" ? id.toLowerCase() : id}`; +} + +function assertIdentity(identity: OAuthCredentialIdentity): void { + if (!identity.provider.trim()) throw new Error("OAuth provider is required."); + if (!identity.accountId.trim()) { + throw new Error("OAuth account id is required."); + } + if (!identity.resource.trim()) throw new Error("OAuth resource is required."); + ownerKey(identity.owner); +} + +function lifecycleMetadata( + identity: OAuthCredentialIdentity, + reconnectReason?: OAuthLifecycleMetadata["reconnectReason"], +): OAuthLifecycleMetadata { + return { + version: LIFECYCLE_VERSION, + provider: identity.provider, + resource: identity.resource, + owner: ownerKey(identity.owner), + ...(reconnectReason ? { reconnectReason } : {}), + }; +} + +function withLifecycle( + identity: OAuthCredentialIdentity, + credential: T, + reconnectReason?: OAuthLifecycleMetadata["reconnectReason"], +): T { + return { + ...credential, + oauthLifecycle: lifecycleMetadata(identity, reconnectReason), + }; +} + +function metadataMatches( + identity: OAuthCredentialIdentity, + metadata: OAuthLifecycleMetadata | undefined, +): boolean { + return Boolean( + metadata && + metadata.version === LIFECYCLE_VERSION && + metadata.provider === identity.provider && + metadata.resource === identity.resource && + metadata.owner === ownerKey(identity.owner), + ); +} + +function leaseKey(identity: OAuthCredentialIdentity): string { + const digest = createHash("sha256") + .update( + JSON.stringify([ + identity.provider, + identity.accountId, + identity.resource, + ownerKey(identity.owner), + ]), + ) + .digest("hex"); + return `oauth-refresh-lease:${digest}`; +} + +function storageAccountId( + identity: OAuthCredentialIdentity, + legacyAccountKey = false, +): string { + if (legacyAccountKey) return identity.accountId; + const resourceHash = createHash("sha256") + .update(identity.resource) + .digest("hex"); + return `${identity.accountId}:resource:${resourceHash}`; +} + +async function acquireLease( + identity: OAuthCredentialIdentity, + holder: string, + leaseMs: number, + now: number, +): Promise { + const next = await mutateSetting(leaseKey(identity), (current) => { + const currentHolder = + typeof current?.holder === "string" ? current.holder : ""; + const expiresAt = + typeof current?.expiresAt === "number" ? current.expiresAt : 0; + if (currentHolder && currentHolder !== holder && expiresAt > now) { + return current!; + } + return { holder, expiresAt: now + leaseMs }; + }); + return next.holder === holder; +} + +async function releaseLease( + identity: OAuthCredentialIdentity, + holder: string, +): Promise { + await mutateSetting(leaseKey(identity), (current) => + current?.holder === holder ? { holder: "", expiresAt: 0 } : (current ?? {}), + ).catch(() => undefined); +} + +function startLeaseHeartbeat( + identity: OAuthCredentialIdentity, + holder: string, + leaseMs: number, + dependencies: LifecycleDependencies, +): () => Promise { + let renewal: Promise | undefined; + const timer = setInterval( + () => { + if (renewal) return; + renewal = acquireLease(identity, holder, leaseMs, dependencies.now()) + .then(() => undefined) + .catch(() => undefined) + .finally(() => { + renewal = undefined; + }); + }, + Math.max(1, Math.floor(leaseMs / 3)), + ); + timer.unref?.(); + return async () => { + clearInterval(timer); + await renewal; + }; +} + +export async function saveOAuthCredential( + identity: OAuthCredentialIdentity, + credential: T, + options: { legacyAccountKey?: boolean } = {}, +): Promise { + assertIdentity(identity); + await saveOAuthTokens( + identity.provider, + storageAccountId(identity, options.legacyAccountKey), + withLifecycle(identity, credential), + ownerKey(identity.owner), + ); +} + +export async function readOAuthCredentialState< + T extends OAuthCredential = OAuthCredential, +>( + identity: OAuthCredentialIdentity, + options: { + allowLegacy?: boolean; + legacyAccountKey?: boolean; + now?: number; + validateCredential?: (credential: T) => boolean; + } = {}, +): Promise> { + assertIdentity(identity); + const stored = await getOAuthTokenSnapshot( + identity.provider, + storageAccountId(identity, options.legacyAccountKey), + ownerKey(identity.owner), + ); + if (!stored) return { kind: "missing" }; + const parsed = stored.tokens as Partial; + if ( + !parsed.tokens || + typeof parsed.tokens.access_token !== "string" || + (!options.allowLegacy && + !metadataMatches(identity, parsed.oauthLifecycle)) || + (parsed.oauthLifecycle && !metadataMatches(identity, parsed.oauthLifecycle)) + ) { + return { kind: "malformed", revision: stored.revision }; + } + const credential = parsed as T; + if (options.validateCredential && !options.validateCredential(credential)) { + return { kind: "malformed", revision: stored.revision }; + } + if (credential.oauthLifecycle?.reconnectReason) { + return { + kind: "reconnect_required", + credential, + revision: stored.revision, + }; + } + const now = options.now ?? Date.now(); + if ( + typeof credential.tokenExpiresAt === "number" && + credential.tokenExpiresAt <= now + ) { + return { kind: "expired", credential, revision: stored.revision }; + } + return { kind: "connected", credential, revision: stored.revision }; +} + +async function markReconnectRequired( + identity: OAuthCredentialIdentity, + snapshot: CredentialSnapshot, + legacyAccountKey: boolean, +): Promise { + await replaceOAuthTokensIfRevision( + identity.provider, + storageAccountId(identity, legacyAccountKey), + ownerKey(identity.owner), + snapshot.revision, + withLifecycle(identity, snapshot.credential, "refresh_failed"), + ); +} + +export async function resolveOAuthCredentialAccess< + T extends OAuthCredential = OAuthCredential, +>( + identity: OAuthCredentialIdentity, + options: { + refresh: (context: OAuthRefreshContext) => Promise; + allowLegacy?: boolean; + legacyAccountKey?: boolean; + validateCredential?: (credential: T) => boolean; + expirySkewMs?: number; + leaseMs?: number; + waitMs?: number; + maxWaitMs?: number; + dependencies?: Partial; + }, +): Promise> { + assertIdentity(identity); + const dependencies = { ...defaultDependencies, ...options.dependencies }; + const expirySkewMs = options.expirySkewMs ?? DEFAULT_EXPIRY_SKEW_MS; + const leaseMs = options.leaseMs ?? DEFAULT_LEASE_MS; + const waitMs = options.waitMs ?? DEFAULT_WAIT_MS; + const maxWaitMs = options.maxWaitMs ?? DEFAULT_MAX_WAIT_MS; + const startedAt = dependencies.now(); + let state = await readOAuthCredentialState(identity, { + allowLegacy: options.allowLegacy, + legacyAccountKey: options.legacyAccountKey, + now: startedAt, + validateCredential: options.validateCredential, + }); + if ( + state.kind === "connected" && + (typeof state.credential.tokenExpiresAt !== "number" || + state.credential.tokenExpiresAt - startedAt > expirySkewMs) + ) { + return { state, accessToken: state.credential.tokens.access_token }; + } + if ( + state.kind === "missing" || + state.kind === "malformed" || + state.kind === "reconnect_required" + ) { + return { state, accessToken: null }; + } + if (!state.credential.tokens.refresh_token) { + return { + state, + accessToken: + state.kind === "connected" + ? state.credential.tokens.access_token + : null, + }; + } + + const baselineRevision = state.revision; + const holder = dependencies.holderId(); + while (dependencies.now() - startedAt <= maxWaitMs) { + const acquired = await acquireLease( + identity, + holder, + leaseMs, + dependencies.now(), + ); + if (!acquired) { + await dependencies.sleep(waitMs); + state = await readOAuthCredentialState(identity, { + allowLegacy: options.allowLegacy, + legacyAccountKey: options.legacyAccountKey, + now: dependencies.now(), + validateCredential: options.validateCredential, + }); + if ( + state.kind === "connected" && + (state.revision !== baselineRevision || + typeof state.credential.tokenExpiresAt !== "number" || + state.credential.tokenExpiresAt - dependencies.now() > expirySkewMs) + ) { + return { state, accessToken: state.credential.tokens.access_token }; + } + if ( + state.kind === "missing" || + state.kind === "malformed" || + state.kind === "reconnect_required" + ) { + return { state, accessToken: null }; + } + continue; + } + + try { + state = await readOAuthCredentialState(identity, { + allowLegacy: options.allowLegacy, + legacyAccountKey: options.legacyAccountKey, + now: dependencies.now(), + validateCredential: options.validateCredential, + }); + if ( + state.kind === "missing" || + state.kind === "malformed" || + state.kind === "reconnect_required" + ) { + return { state, accessToken: null }; + } + if ( + state.kind === "connected" && + (typeof state.credential.tokenExpiresAt !== "number" || + state.credential.tokenExpiresAt - dependencies.now() > expirySkewMs) + ) { + return { + state, + accessToken: state.credential.tokens.access_token, + }; + } + try { + const stopHeartbeat = startLeaseHeartbeat( + identity, + holder, + leaseMs, + dependencies, + ); + const refreshed = withLifecycle( + identity, + await options + .refresh({ identity, credential: state.credential }) + .finally(stopHeartbeat), + ); + const saved = await replaceOAuthTokensIfRevision( + identity.provider, + storageAccountId(identity, options.legacyAccountKey), + ownerKey(identity.owner), + state.revision, + refreshed, + ); + const latest = await readOAuthCredentialState(identity, { + allowLegacy: options.allowLegacy, + legacyAccountKey: options.legacyAccountKey, + now: dependencies.now(), + validateCredential: options.validateCredential, + }); + if (!saved && latest.kind === "connected") { + return { + state: latest, + accessToken: latest.credential.tokens.access_token, + }; + } + return { + state: latest, + accessToken: + latest.kind === "connected" + ? latest.credential.tokens.access_token + : null, + }; + } catch { + const stillOwnsLease = await acquireLease( + identity, + holder, + leaseMs, + dependencies.now(), + ); + if (!stillOwnsLease) { + await dependencies.sleep(waitMs); + continue; + } + const latest = await readOAuthCredentialState(identity, { + allowLegacy: options.allowLegacy, + legacyAccountKey: options.legacyAccountKey, + now: dependencies.now(), + validateCredential: options.validateCredential, + }); + if (latest.kind === "connected") { + return { + state: latest, + accessToken: latest.credential.tokens.access_token, + }; + } + if (latest.kind === "expired") { + await markReconnectRequired( + identity, + latest, + options.legacyAccountKey === true, + ); + const reconnect = await readOAuthCredentialState(identity, { + allowLegacy: options.allowLegacy, + legacyAccountKey: options.legacyAccountKey, + now: dependencies.now(), + validateCredential: options.validateCredential, + }); + return { state: reconnect, accessToken: null }; + } + return { state: latest, accessToken: null }; + } + } finally { + await releaseLease(identity, holder); + } + } + + state = await readOAuthCredentialState(identity, { + allowLegacy: options.allowLegacy, + legacyAccountKey: options.legacyAccountKey, + now: dependencies.now(), + validateCredential: options.validateCredential, + }); + return { + state, + accessToken: + state.kind === "connected" ? state.credential.tokens.access_token : null, + }; +} + +export async function revokeOAuthCredential( + identity: OAuthCredentialIdentity, + options: { + revoke?: ( + context: OAuthRefreshContext, + ) => Promise<"succeeded" | "unsupported">; + allowLegacy?: boolean; + legacyAccountKey?: boolean; + validateCredential?: (credential: T) => boolean; + } = {}, +): Promise { + const state = await readOAuthCredentialState(identity, { + allowLegacy: options.allowLegacy, + legacyAccountKey: options.legacyAccountKey, + validateCredential: options.validateCredential, + }); + if (state.kind === "missing") { + return { remote: "not_attempted", local: "missing" }; + } + let remote: OAuthRevocationResult["remote"] = options.revoke + ? "failed" + : "unsupported"; + if (options.revoke && state.kind !== "malformed") { + try { + remote = await options.revoke({ + identity, + credential: state.credential, + }); + } catch { + remote = "failed"; + } + } else if (state.kind === "malformed") { + remote = "not_attempted"; + } + const deleted = await deleteOAuthTokensIfRevision( + identity.provider, + storageAccountId(identity, options.legacyAccountKey), + ownerKey(identity.owner), + state.revision, + ); + return { remote, local: deleted ? "deleted" : "replaced" }; +} diff --git a/packages/core/src/oauth-tokens/store.spec.ts b/packages/core/src/oauth-tokens/store.spec.ts index bb53a785c8..b67267cd58 100644 --- a/packages/core/src/oauth-tokens/store.spec.ts +++ b/packages/core/src/oauth-tokens/store.spec.ts @@ -15,7 +15,10 @@ interface ExecCall { const execCalls: ExecCall[] = []; let existingOwner: string | null = null; let existingTokens: Record | null = null; +let existingRevision = 100; let mockPostgres = false; +let conflictOwnerAfterUpsert: string | null = null; +let upsertAttempted = false; const mockDb = { execute: vi.fn(async (input: string | { sql: string; args?: unknown[] }) => { @@ -23,6 +26,42 @@ const mockDb = { const args = typeof input === "string" ? [] : (input.args ?? []); execCalls.push({ sql, args }); + if (/^\s*INSERT\s+INTO\s+(?:public\.)?oauth_tokens/i.test(sql)) { + upsertAttempted = true; + return { + rows: [], + rowsAffected: conflictOwnerAfterUpsert ? 0 : 1, + }; + } + + if (/SELECT owner FROM (?:public\.)?oauth_tokens/i.test(sql)) { + const owner = + existingOwner ?? (upsertAttempted ? conflictOwnerAfterUpsert : null); + return { + rows: owner ? [{ owner }] : [], + rowsAffected: 0, + }; + } + + if ( + /SELECT owner, tokens, updated_at FROM (?:public\.)?oauth_tokens/i.test( + sql, + ) + ) { + return { + rows: existingOwner + ? [ + { + owner: existingOwner, + tokens: JSON.stringify(existingTokens ?? {}), + updated_at: existingRevision, + }, + ] + : [], + rowsAffected: 0, + }; + } + if ( /SELECT owner, display_name, tokens FROM (?:public\.)?oauth_tokens/i.test( sql, @@ -42,6 +81,10 @@ const mockDb = { }; } + if (/^(?:UPDATE\s+|DELETE\s+FROM\s+)(?:public\.)?oauth_tokens/i.test(sql)) { + return { rows: [], rowsAffected: existingOwner ? 1 : 0 }; + } + return { rows: [], rowsAffected: 0 }; }), }; @@ -52,8 +95,14 @@ vi.mock("../db/client.js", () => ({ isPostgres: () => mockPostgres, })); -const { deleteOAuthTokens, getOAuthTokens, saveOAuthTokens } = - await import("./store.js"); +const { + deleteOAuthTokens, + deleteOAuthTokensIfRevision, + getOAuthTokenSnapshot, + getOAuthTokens, + replaceOAuthTokensIfRevision, + saveOAuthTokens, +} = await import("./store.js"); function lastInsert(): ExecCall { const inserts = execCalls.filter((c) => /^\s*INSERT\b/i.test(c.sql)); @@ -66,7 +115,10 @@ describe("oauth token store", () => { execCalls.length = 0; existingOwner = null; existingTokens = null; + existingRevision = 100; mockPostgres = false; + conflictOwnerAfterUpsert = null; + upsertAttempted = false; vi.clearAllMocks(); }); @@ -87,6 +139,27 @@ describe("oauth token store", () => { }); }); + it("refuses a different owner that wins the row between the pre-read and upsert", async () => { + conflictOwnerAfterUpsert = "other@example.com"; + + await expect( + saveOAuthTokens( + "google", + "steve@builder.io", + { access_token: "new-token" }, + "steve@builder.io", + ), + ).rejects.toMatchObject({ + name: "OAuthAccountOwnedByOtherUserError", + existingOwner: "other@example.com", + attemptedOwner: "steve@builder.io", + }); + + expect(lastInsert().sql).toContain( + "WHERE oauth_tokens.owner = excluded.owner", + ); + }); + it("supports owner-scoped reads and deletes for tenant-bound OAuth credentials", async () => { await getOAuthTokens("mcp", "mcp_oauth:test", "org:org-test"); await deleteOAuthTokens("mcp", "mcp_oauth:test", "org:org-test"); @@ -108,6 +181,70 @@ describe("oauth token store", () => { ]); }); + it("reads and conditionally replaces one exact owner-bound revision", async () => { + existingOwner = "user:alice@example.com"; + existingTokens = { access_token: "old-access" }; + + await expect( + getOAuthTokenSnapshot("builder", "managed-ai", "user:alice@example.com"), + ).resolves.toMatchObject({ + owner: "user:alice@example.com", + revision: 100, + tokens: { access_token: "old-access" }, + }); + + await expect( + replaceOAuthTokensIfRevision( + "builder", + "managed-ai", + "user:alice@example.com", + 100, + { access_token: "new-access", refresh_token: "new-refresh" }, + ), + ).resolves.toBe(true); + + const conditionalUpdate = execCalls.find((call) => + /^UPDATE\s+(?:public\.)?oauth_tokens/i.test(call.sql), + ); + expect(conditionalUpdate?.sql).toContain("AND updated_at = ?"); + expect(conditionalUpdate?.args.slice(-4)).toEqual([ + "builder", + "managed-ai", + "user:alice@example.com", + 100, + ]); + const encrypted = conditionalUpdate?.args[0] as string; + expect(isEncryptedSecretValue(encrypted)).toBe(true); + expect(JSON.parse(decryptSecretValue(encrypted))).toMatchObject({ + access_token: "new-access", + refresh_token: "new-refresh", + }); + }); + + it("conditionally deletes only the revision and owner that were inspected", async () => { + existingOwner = "user:alice@example.com"; + + await expect( + deleteOAuthTokensIfRevision( + "builder", + "managed-ai", + "user:alice@example.com", + 100, + ), + ).resolves.toBe(true); + + const conditionalDelete = execCalls.find((call) => + /^DELETE\s+FROM\s+(?:public\.)?oauth_tokens/i.test(call.sql), + ); + expect(conditionalDelete?.sql).toContain("AND updated_at = ?"); + expect(conditionalDelete?.args).toEqual([ + "builder", + "managed-ai", + "user:alice@example.com", + 100, + ]); + }); + it("qualifies the real oauth_tokens table on Postgres so temp scoped views cannot shadow OAuth callbacks", async () => { mockPostgres = true; @@ -154,6 +291,21 @@ describe("oauth token store", () => { }); }); + it("advances the revision atomically when a replacement lands in the same clock tick", async () => { + existingOwner = "user:alice@example.com"; + + await saveOAuthTokens( + "builder", + "managed-ai", + { access_token: "new-access" }, + "user:alice@example.com", + ); + + expect(lastInsert().sql).toContain( + "updated_at=MAX(oauth_tokens.updated_at + 1, excluded.updated_at)", + ); + }); + it("encrypts the token bundle at rest (no plaintext refresh token in the column)", async () => { existingOwner = null; diff --git a/packages/core/src/oauth-tokens/store.ts b/packages/core/src/oauth-tokens/store.ts index 81695e5847..5fa027ed7d 100644 --- a/packages/core/src/oauth-tokens/store.ts +++ b/packages/core/src/oauth-tokens/store.ts @@ -157,6 +157,83 @@ export async function getOAuthTokens( return parseStoredTokens(rows[0].tokens as string); } +export interface OAuthTokenSnapshot { + tokens: Record; + owner: string | null; + revision: number; +} + +/** + * Read one credential bundle together with the row revision used by + * refresh/revocation compare-and-swap writes. + */ +export async function getOAuthTokenSnapshot( + provider: string, + accountId: string, + owner: string, +): Promise { + await ensureTable(); + const client = getDbExec(); + const table = oauthTokensTable(); + const { rows } = await client.execute({ + sql: `SELECT owner, tokens, updated_at FROM ${table} WHERE provider = ? AND account_id = ? AND owner = ?`, + args: [provider, accountId, owner], + }); + if (rows.length === 0) return null; + return { + tokens: parseStoredTokens(rows[0].tokens as string), + owner: (rows[0].owner as string) ?? null, + revision: Number(rows[0].updated_at), + }; +} + +/** + * Replace an existing credential bundle only when the caller still owns the + * revision it read. This prevents a slow refresh/revoke flow from overwriting + * a newer authorization completed in another process. + */ +export async function replaceOAuthTokensIfRevision( + provider: string, + accountId: string, + owner: string, + expectedRevision: number, + tokens: Record, +): Promise { + await ensureTable(); + const client = getDbExec(); + const table = oauthTokensTable(); + const nextRevision = Math.max(Date.now(), expectedRevision + 1); + const result = await client.execute({ + sql: `UPDATE ${table} SET tokens = ?, updated_at = ? WHERE provider = ? AND account_id = ? AND owner = ? AND updated_at = ?`, + args: [ + serializeTokens(tokens), + nextRevision, + provider, + accountId, + owner, + expectedRevision, + ], + }); + return result.rowsAffected === 1; +} + +/** Delete only the exact credential revision the caller inspected. */ +export async function deleteOAuthTokensIfRevision( + provider: string, + accountId: string, + owner: string, + expectedRevision: number, +): Promise { + await ensureTable(); + const client = getDbExec(); + const table = oauthTokensTable(); + const result = await client.execute({ + sql: `DELETE FROM ${table} WHERE provider = ? AND account_id = ? AND owner = ? AND updated_at = ?`, + args: [provider, accountId, owner, expectedRevision], + }); + return result.rowsAffected === 1; +} + /** * Thrown when an OAuth save would re-bind an `(provider, account_id)` row * to a different owner than already holds it. Callers should catch this and @@ -256,10 +333,10 @@ export async function saveOAuthTokens( ...cleanedIncomingTokens, }; - await client.execute({ + const result = await client.execute({ sql: isPostgres() - ? `INSERT INTO ${table} (provider, account_id, owner, display_name, tokens, updated_at) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT (provider, account_id) DO UPDATE SET owner=EXCLUDED.owner, display_name=COALESCE(EXCLUDED.display_name, ${table}.display_name), tokens=EXCLUDED.tokens, updated_at=EXCLUDED.updated_at` - : `INSERT OR REPLACE INTO ${table} (provider, account_id, owner, display_name, tokens, updated_at) VALUES (?, ?, ?, ?, ?, ?)`, + ? `INSERT INTO ${table} (provider, account_id, owner, display_name, tokens, updated_at) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT (provider, account_id) DO UPDATE SET display_name=COALESCE(EXCLUDED.display_name, ${table}.display_name), tokens=EXCLUDED.tokens, updated_at=GREATEST(${table}.updated_at + 1, EXCLUDED.updated_at) WHERE ${table}.owner = EXCLUDED.owner` + : `INSERT INTO ${table} (provider, account_id, owner, display_name, tokens, updated_at) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT (provider, account_id) DO UPDATE SET display_name=COALESCE(excluded.display_name, ${table}.display_name), tokens=excluded.tokens, updated_at=MAX(${table}.updated_at + 1, excluded.updated_at) WHERE ${table}.owner = excluded.owner`, args: [ provider, accountId, @@ -269,6 +346,22 @@ export async function saveOAuthTokens( Date.now(), ], }); + if (result.rowsAffected === 1) return; + + const { rows: conflict } = await client.execute({ + sql: `SELECT owner FROM ${table} WHERE provider = ? AND account_id = ?`, + args: [provider, accountId], + }); + const conflictOwner = (conflict[0]?.owner as string | undefined) ?? ""; + if (conflictOwner && conflictOwner !== resolvedOwner) { + throw new OAuthAccountOwnedByOtherUserError({ + provider, + accountId, + existingOwner: conflictOwner, + attemptedOwner: resolvedOwner, + }); + } + throw new Error(`OAuth account ${provider}:${accountId} was not saved.`); } export async function deleteOAuthTokens(