From ab638295536ab0a78a35b4d24e163f91fe1ef95b Mon Sep 17 00:00:00 2001 From: li-xiu-qi Date: Sat, 24 Oct 2026 14:17:00 +0800 Subject: [PATCH] fix(mcp): prevent lost updates during concurrent OAuth token saves --- packages/coding-agent/src/step/mcp-oauth.ts | 120 +++++++-- .../fixtures/mcp-oauth-lock-holder-child.ts | 35 +++ .../fixtures/mcp-oauth-save-tokens-child.ts | 63 +++++ .../coding-agent/test/step-mcp-oauth.test.ts | 232 +++++++++++++++++- 4 files changed, 426 insertions(+), 24 deletions(-) create mode 100644 packages/coding-agent/test/fixtures/mcp-oauth-lock-holder-child.ts create mode 100644 packages/coding-agent/test/fixtures/mcp-oauth-save-tokens-child.ts diff --git a/packages/coding-agent/src/step/mcp-oauth.ts b/packages/coding-agent/src/step/mcp-oauth.ts index 66a4d35e..5a25fbe5 100644 --- a/packages/coding-agent/src/step/mcp-oauth.ts +++ b/packages/coding-agent/src/step/mcp-oauth.ts @@ -1,13 +1,15 @@ -import { randomBytes, timingSafeEqual } from "node:crypto"; -import { existsSync, mkdirSync, readFileSync, renameSync, writeFileSync } from "node:fs"; +import { randomBytes, randomUUID, timingSafeEqual } from "node:crypto"; +import { existsSync, mkdirSync, readFileSync, renameSync, unlinkSync, writeFileSync } from "node:fs"; import { createServer } from "node:http"; import { join } from "node:path"; +import { setTimeout as sleep } from "node:timers/promises"; import { auth } from "@modelcontextprotocol/sdk/client/auth.js"; import type { OAuthClientInformationMixed, OAuthClientMetadata, OAuthTokens, } from "@modelcontextprotocol/sdk/shared/auth.js"; +import lockfile from "proper-lockfile"; import { openBrowser } from "../utils/open-browser.ts"; import { resolveStepConfigRoot } from "./environment.ts"; @@ -38,7 +40,7 @@ export function createStoredMcpOAuthProvider( clientInformation(): OAuthClientInformationMixed | undefined; tokens(): OAuthTokens | undefined; saveClientInformation(value: OAuthClientInformationMixed): void; - saveTokens(value: OAuthTokens): void; + saveTokens(value: OAuthTokens): Promise; redirectToAuthorization(): never; saveCodeVerifier(value: string): void; codeVerifier(): string; @@ -60,10 +62,8 @@ export function createStoredMcpOAuthProvider( saveClientInformation: (value) => { clientInformation = value; }, - saveTokens: (tokens) => { - const store = loadStore(env); - store[key(name, serverUrl)] = { serverName: name, serverUrl, clientInformation, tokens }; - saveStore(env, store); + saveTokens: async (tokens) => { + await saveCredentialTokens(env, name, serverUrl, clientInformation, tokens); }, redirectToAuthorization: () => { throw new Error(`MCP server '${name}' requires login; run step mcp login ${name}`); @@ -94,6 +94,77 @@ function key(name: string, url: string): string { return `${name}|${url}`; } +function errorCode(error: unknown): string | undefined { + return typeof error === "object" && error !== null && "code" in error + ? String((error as { code?: unknown }).code) + : undefined; +} + +function acquireCredentialStoreLock(path: string): () => void { + const maxAttempts = 100; + const delayMs = 20; + for (let attempt = 1; attempt <= maxAttempts; attempt += 1) { + try { + return lockfile.lockSync(path, { realpath: false }); + } catch (error) { + if (errorCode(error) !== "ELOCKED" || attempt === maxAttempts) throw error; + const started = Date.now(); + while (Date.now() - started < delayMs) { + // Keep the synchronous logout API while a peer writer finishes. + } + } + } + throw new Error("Failed to acquire MCP OAuth credential store lock"); +} + +function withCredentialStoreLock(env: NodeJS.ProcessEnv, fn: () => T): T { + mkdirSync(resolveStepConfigRoot(env), { recursive: true, mode: 0o700 }); + const release = acquireCredentialStoreLock(credentialsPath(env)); + try { + return fn(); + } finally { + release(); + } +} + +async function acquireCredentialStoreLockAsync(path: string): Promise<() => Promise> { + const timeoutMs = 30_000; + const delayMs = 20; + const deadline = Date.now() + timeoutMs; + while (true) { + try { + return await lockfile.lock(path, { realpath: false, retries: 0, stale: 60_000 }); + } catch (error) { + if (errorCode(error) !== "ELOCKED" || Date.now() >= deadline) throw error; + await sleep(delayMs); + } + } +} + +async function withCredentialStoreLockAsync(env: NodeJS.ProcessEnv, fn: () => T | Promise): Promise { + mkdirSync(resolveStepConfigRoot(env), { recursive: true, mode: 0o700 }); + const release = await acquireCredentialStoreLockAsync(credentialsPath(env)); + try { + return await fn(); + } finally { + await release(); + } +} + +async function saveCredentialTokens( + env: NodeJS.ProcessEnv, + name: string, + serverUrl: string, + clientInformation: OAuthClientInformationMixed | undefined, + tokens: OAuthTokens, +): Promise { + await withCredentialStoreLockAsync(env, () => { + const store = loadStore(env); + store[key(name, serverUrl)] = { serverName: name, serverUrl, clientInformation, tokens }; + saveStore(env, store); + }); +} + function loadStore(env: NodeJS.ProcessEnv): CredentialStore { const path = credentialsPath(env); if (!existsSync(path)) return {}; @@ -104,9 +175,18 @@ function loadStore(env: NodeJS.ProcessEnv): CredentialStore { function saveStore(env: NodeJS.ProcessEnv, store: CredentialStore): void { const path = credentialsPath(env); mkdirSync(resolveStepConfigRoot(env), { recursive: true, mode: 0o700 }); - const temporary = `${path}.${process.pid}.tmp`; - writeFileSync(temporary, `${JSON.stringify(store, null, 2)}\n`, { encoding: "utf8", mode: 0o600 }); - renameSync(temporary, path); + const temporary = `${path}.${process.pid}.${randomUUID()}.tmp`; + try { + writeFileSync(temporary, `${JSON.stringify(store, null, 2)}\n`, { encoding: "utf8", mode: 0o600 }); + renameSync(temporary, path); + } catch (error) { + try { + unlinkSync(temporary); + } catch { + // Cleanup is best effort and must not mask the original write or rename failure. + } + throw error; + } } export async function loginMcpServer( @@ -181,10 +261,8 @@ export async function loginMcpServer( }, state: () => state, tokens: () => loadStore(env)[key(name, serverUrl)]?.tokens, - saveTokens: (tokens: OAuthTokens) => { - const store = loadStore(env); - store[key(name, serverUrl)] = { serverName: name, serverUrl, clientInformation, tokens }; - saveStore(env, store); + saveTokens: async (tokens: OAuthTokens) => { + await saveCredentialTokens(env, name, serverUrl, clientInformation, tokens); }, redirectToAuthorization: (url: URL) => { process.stderr.write(`Open this URL to authorize ${name}:\n${url}\n`); @@ -278,10 +356,12 @@ function matchesState(received: string | null, expected: string): boolean { } export function logoutMcpServer(name: string, serverUrl: string, env: NodeJS.ProcessEnv = process.env): boolean { - const store = loadStore(env); - const entry = key(name, serverUrl); - if (!store[entry]) return false; - delete store[entry]; - saveStore(env, store); - return true; + return withCredentialStoreLock(env, () => { + const store = loadStore(env); + const entry = key(name, serverUrl); + if (!store[entry]) return false; + delete store[entry]; + saveStore(env, store); + return true; + }); } diff --git a/packages/coding-agent/test/fixtures/mcp-oauth-lock-holder-child.ts b/packages/coding-agent/test/fixtures/mcp-oauth-lock-holder-child.ts new file mode 100644 index 00000000..8ae214bd --- /dev/null +++ b/packages/coding-agent/test/fixtures/mcp-oauth-lock-holder-child.ts @@ -0,0 +1,35 @@ +import { readFileSync, renameSync, writeFileSync } from "node:fs"; +import process from "node:process"; +import lockfile from "proper-lockfile"; + +const path = process.env.MCP_OAUTH_CREDENTIALS_PATH; +const serverName = process.env.MCP_OAUTH_SERVER_NAME; +const serverUrl = process.env.MCP_OAUTH_SERVER_URL; +if (!path || !serverName || !serverUrl) throw new Error("Missing MCP OAuth lock holder environment"); + +const release = await lockfile.lock(path, { realpath: false, stale: 60_000 }); +process.send?.({ type: "ready", pid: process.pid }); + +process.on("message", async (message: { type?: string; holdMs?: number }) => { + if (message.type !== "write" || typeof message.holdMs !== "number") return; + try { + const started = Date.now(); + while (Date.now() - started < message.holdMs) { + // Simulate a slow synchronous read-modify-write that blocks lock heartbeat. + } + const store = JSON.parse(readFileSync(path, "utf8")); + store[`${serverName}|${serverUrl}`] = { + serverName, + serverUrl, + tokens: { access_token: `token-${serverName}`, token_type: "Bearer" }, + }; + const temporary = `${path}.${process.pid}.holder.tmp`; + writeFileSync(temporary, `${JSON.stringify(store, null, 2)}\n`, { encoding: "utf8", mode: 0o600 }); + renameSync(temporary, path); + process.send?.({ type: "released" }); + } catch (error) { + process.send?.({ type: "save-error", error: error instanceof Error ? error.stack : String(error) }); + } finally { + await release(); + } +}); diff --git a/packages/coding-agent/test/fixtures/mcp-oauth-save-tokens-child.ts b/packages/coding-agent/test/fixtures/mcp-oauth-save-tokens-child.ts new file mode 100644 index 00000000..7a3af7e4 --- /dev/null +++ b/packages/coding-agent/test/fixtures/mcp-oauth-save-tokens-child.ts @@ -0,0 +1,63 @@ +import fs, { existsSync, writeFileSync } from "node:fs"; +import { syncBuiltinESMExports } from "node:module"; +import { join } from "node:path"; +import process from "node:process"; +import lockfile from "proper-lockfile"; + +const moduleUrl = process.env.MCP_OAUTH_MODULE; +const credentialsPath = process.env.MCP_OAUTH_CREDENTIALS_PATH; +const barrierDir = process.env.MCP_OAUTH_BARRIER_DIR; +const serverName = process.env.MCP_OAUTH_SERVER_NAME; +const serverUrl = process.env.MCP_OAUTH_SERVER_URL; +if (!moduleUrl || !credentialsPath || !barrierDir || !serverName || !serverUrl) { + throw new Error("Missing MCP OAuth child fixture environment"); +} + +let armedRound: number | undefined; +const originalReadFileSync = fs.readFileSync; +fs.readFileSync = ((...args: unknown[]) => { + const result: unknown = Reflect.apply(originalReadFileSync, fs, args); + if (armedRound !== undefined && String(args[0]) === credentialsPath) { + writeFileSync(join(barrierDir, `read-${serverName}-${armedRound}`), ""); + const releasePath = join(barrierDir, `release-${serverName}-${armedRound}`); + while (!existsSync(releasePath)) { + // The parent releases this read only after the competing writer is gated too. + } + } + return result; +}) as typeof fs.readFileSync; +syncBuiltinESMExports(); + +const originalLock = lockfile.lock; +lockfile.lock = (async (...args: unknown[]) => { + const release: unknown = await Reflect.apply(originalLock, lockfile, args); + if (armedRound !== undefined) { + writeFileSync(join(barrierDir, `lock-${serverName}-${armedRound}`), ""); + } + return release; +}) as typeof lockfile.lock; + +const { createStoredMcpOAuthProvider } = await import(moduleUrl); +const provider = createStoredMcpOAuthProvider(serverName, serverUrl, process.env); + +process.on("message", async (message: { type?: string; round?: number; token?: string }) => { + if (message.type === "exit") { + process.exit(0); + } + if (message.type !== "save" || typeof message.round !== "number" || !message.token) return; + try { + armedRound = message.round; + await provider.saveTokens({ access_token: message.token, token_type: "Bearer" }); + process.send?.({ type: "saved", round: message.round }); + } catch (error) { + process.send?.({ + type: "save-error", + round: message.round, + error: error instanceof Error ? error.stack : String(error), + }); + } finally { + armedRound = undefined; + } +}); + +process.send?.({ type: "ready", pid: process.pid }); diff --git a/packages/coding-agent/test/step-mcp-oauth.test.ts b/packages/coding-agent/test/step-mcp-oauth.test.ts index ac61c34f..812c1ea6 100644 --- a/packages/coding-agent/test/step-mcp-oauth.test.ts +++ b/packages/coding-agent/test/step-mcp-oauth.test.ts @@ -1,7 +1,11 @@ -import { existsSync, mkdirSync, readFileSync, rmSync } from "node:fs"; +import { type ChildProcess, spawn } from "node:child_process"; +import fs, { existsSync, mkdirSync, readdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { createServer } from "node:http"; +import { syncBuiltinESMExports } from "node:module"; import { join } from "node:path"; -import { afterEach, describe, expect, it } from "vitest"; +import { pathToFileURL } from "node:url"; +import lockfile from "proper-lockfile"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { buildOAuthClientMetadata, createStoredMcpOAuthProvider, @@ -11,6 +15,9 @@ import { resolveOAuthCallback, } from "../src/step/mcp-oauth.ts"; +const { authMock } = vi.hoisted(() => ({ authMock: vi.fn() })); +vi.mock("@modelcontextprotocol/sdk/client/auth.js", () => ({ auth: authMock })); + function sink(): { calls: string[]; resolve(code: string): void; reject(error: Error): void } { const calls: string[] = []; return { @@ -24,6 +31,36 @@ function callbackUrl(query: string, path = "/callback"): URL { return new URL(`${path}${query}`, "http://127.0.0.1"); } +function waitForChildMessage(child: ChildProcess, type: string, round?: number, timeoutMs = 10000): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + cleanup(); + reject(new Error(`Timed out waiting for ${type} from child ${child.pid}`)); + }, timeoutMs); + const cleanup = () => { + clearTimeout(timer); + child.off("message", onMessage); + child.off("exit", onExit); + }; + const onMessage = (message: { type?: string; round?: number; error?: string }) => { + if (message.type === "save-error") { + cleanup(); + reject(new Error(message.error)); + return; + } + if (message.type !== type || (round !== undefined && message.round !== round)) return; + cleanup(); + resolve(); + }; + const onExit = (code: number | null, signal: NodeJS.Signals | null) => { + cleanup(); + reject(new Error(`Child ${child.pid} exited early: code=${code}, signal=${signal}`)); + }; + child.on("message", onMessage); + child.on("exit", onExit); + }); +} + describe("MCP OAuth callback", () => { it("ignores any path other than the callback path", () => { const target = sink(); @@ -131,13 +168,15 @@ describe("MCP OAuth credential storage", () => { for (const root of roots.splice(0)) { if (existsSync(root)) rmSync(root, { recursive: true, force: true }); } + authMock.mockReset(); + vi.restoreAllMocks(); }); - it("stores credentials beside the config, not inside an injected agent directory", () => { + it("stores credentials beside the config, not inside an injected agent directory", async () => { const { env, configRoot } = makeEnv(); const provider = createStoredMcpOAuthProvider("demo", "https://example.invalid/mcp", env); - provider.saveTokens({ access_token: "token", token_type: "Bearer" }); + await provider.saveTokens({ access_token: "token", token_type: "Bearer" }); const path = join(configRoot, ".credentials.json"); expect(existsSync(path)).toBe(true); @@ -151,4 +190,189 @@ describe("MCP OAuth credential storage", () => { expect(logoutMcpServer("demo", "https://example.invalid/mcp", env)).toBe(false); expect(hasStoredMcpOAuthCredential("demo", "https://example.invalid/mcp", env)).toBe(false); }); + + it("persists interactive login tokens through the credential store lock", async () => { + const { env, configRoot } = makeEnv(); + const lockSpy = vi.spyOn(lockfile, "lock"); + authMock.mockImplementation( + async (provider: { + saveTokens(tokens: { access_token: string; token_type: string }): void | Promise; + }) => { + await provider.saveTokens({ access_token: "login-token", token_type: "Bearer" }); + return "AUTHORIZED"; + }, + ); + + await loginMcpServer("demo", "https://example.invalid/mcp", {}, env); + + expect(lockSpy).toHaveBeenCalledWith(join(configRoot, ".credentials.json"), { + realpath: false, + retries: 0, + stale: 60_000, + }); + expect(JSON.parse(readFileSync(join(configRoot, ".credentials.json"), "utf8"))).toMatchObject({ + "demo|https://example.invalid/mcp": { tokens: { access_token: "login-token" } }, + }); + }); + + it("fails closed when the credential store lock cannot be acquired", async () => { + const { env, configRoot } = makeEnv(); + const path = join(configRoot, ".credentials.json"); + writeFileSync(path, "{}\n", { encoding: "utf8", mode: 0o600 }); + const lockSpy = vi.spyOn(lockfile, "lock").mockImplementation(async () => { + throw Object.assign(new Error("locked"), { code: "ELOCKED" }); + }); + const startedAt = Date.now(); + let now = startedAt; + vi.spyOn(Date, "now").mockImplementation(() => { + now += 1000; + return now; + }); + const provider = createStoredMcpOAuthProvider("demo", "https://example.invalid/mcp", env); + + await expect(provider.saveTokens({ access_token: "blocked", token_type: "Bearer" })).rejects.toThrow("locked"); + expect(now - startedAt).toBeGreaterThanOrEqual(30_000); + expect(lockSpy.mock.calls.length).toBeGreaterThan(1); + expect(readFileSync(path, "utf8")).toBe("{}\n"); + }); + + it("cleans up a failed publication and releases the credential store lock", async () => { + const { env, configRoot } = makeEnv(); + const path = join(configRoot, ".credentials.json"); + const original = '{"version":1}\n'; + writeFileSync(path, original, { encoding: "utf8", mode: 0o600 }); + const provider = createStoredMcpOAuthProvider("demo", "https://example.invalid/mcp", env); + const renameSpy = vi.spyOn(fs, "renameSync").mockImplementation(() => { + throw new Error("rename blocked"); + }); + syncBuiltinESMExports(); + + try { + await expect(provider.saveTokens({ access_token: "new-token", token_type: "Bearer" })).rejects.toThrow( + "rename blocked", + ); + } finally { + renameSpy.mockRestore(); + syncBuiltinESMExports(); + } + + expect(readFileSync(path, "utf8")).toBe(original); + expect(readdirSync(configRoot).filter((entry) => entry.endsWith(".tmp"))).toEqual([]); + + await provider.saveTokens({ access_token: "after-error", token_type: "Bearer" }); + expect(JSON.parse(readFileSync(path, "utf8"))).toMatchObject({ + "demo|https://example.invalid/mcp": { tokens: { access_token: "after-error" } }, + }); + }); + + it("does not steal a live lock while the holder is blocked in synchronous I/O", async () => { + const { env, configRoot } = makeEnv(); + const path = join(configRoot, ".credentials.json"); + writeFileSync(path, "{}\n", { encoding: "utf8", mode: 0o600 }); + const childScript = join(process.cwd(), "test", "fixtures", "mcp-oauth-lock-holder-child.ts"); + const child = spawn(process.execPath, ["--import", "tsx", childScript], { + cwd: process.cwd(), + env: { + ...process.env, + ...env, + MCP_OAUTH_CREDENTIALS_PATH: path, + MCP_OAUTH_SERVER_NAME: "server-a", + MCP_OAUTH_SERVER_URL: "https://server-a.example.invalid/mcp", + }, + stdio: ["ignore", "ignore", "inherit", "ipc"], + }); + + try { + await waitForChildMessage(child, "ready"); + const provider = createStoredMcpOAuthProvider("server-b", "https://server-b.example.invalid/mcp", env); + const saving = provider.saveTokens({ access_token: "token-server-b", token_type: "Bearer" }); + child.send({ type: "write", holdMs: 11_000 }); + await Promise.all([saving, waitForChildMessage(child, "released", undefined, 15000)]); + + const store = JSON.parse(readFileSync(path, "utf8")); + expect(store["server-a|https://server-a.example.invalid/mcp"]?.tokens?.access_token).toBe("token-server-a"); + expect(store["server-b|https://server-b.example.invalid/mcp"]?.tokens?.access_token).toBe("token-server-b"); + } finally { + const exited = new Promise((resolve) => { + if (child.exitCode !== null || child.signalCode !== null) resolve(); + else child.once("exit", () => resolve()); + }); + child.kill(); + await exited; + } + }, 20000); + + it("preserves concurrent token saves from separate processes", async () => { + const { env, configRoot } = makeEnv(); + const path = join(configRoot, ".credentials.json"); + writeFileSync(path, "{}\n", { encoding: "utf8", mode: 0o600 }); + const moduleUrl = pathToFileURL(join(process.cwd(), "src", "step", "mcp-oauth.ts")).href; + const childScript = join(process.cwd(), "test", "fixtures", "mcp-oauth-save-tokens-child.ts"); + const servers = [ + { name: "server-a", url: "https://server-a.example.invalid/mcp" }, + { name: "server-b", url: "https://server-b.example.invalid/mcp" }, + ]; + const children = servers.map((server) => + spawn(process.execPath, ["--import", "tsx", childScript], { + cwd: process.cwd(), + env: { + ...process.env, + ...env, + MCP_OAUTH_MODULE: moduleUrl, + MCP_OAUTH_CREDENTIALS_PATH: path, + MCP_OAUTH_BARRIER_DIR: configRoot, + MCP_OAUTH_SERVER_NAME: server.name, + MCP_OAUTH_SERVER_URL: server.url, + }, + stdio: ["ignore", "ignore", "inherit", "ipc"], + }), + ); + + try { + await Promise.all(children.map((child) => waitForChildMessage(child, "ready"))); + const round = 1; + const saved = children.map((child) => waitForChildMessage(child, "saved", round)); + children.forEach((child, index) => { + child.send({ type: "save", round, token: `token-${servers[index].name}` }); + }); + + const readMarkers = servers.map((server) => join(configRoot, `read-${server.name}-${round}`)); + const releaseMarkers = servers.map((server) => join(configRoot, `release-${server.name}-${round}`)); + const lockMarkers = servers.map((server) => join(configRoot, `lock-${server.name}-${round}`)); + const released = new Set(); + const deadline = Date.now() + 10000; + while (released.size < servers.length) { + const allReadersGated = readMarkers.every((marker) => existsSync(marker)); + for (let index = 0; index < servers.length; index += 1) { + if (released.has(index) || !existsSync(readMarkers[index])) continue; + // Without a lock, release only after both processes hold the same old snapshot. + // With a lock, this reader owns it and must finish before its peer can read. + if (allReadersGated || existsSync(lockMarkers[index])) { + writeFileSync(releaseMarkers[index], ""); + released.add(index); + } + } + if (released.size < servers.length) { + if (Date.now() > deadline) throw new Error("Timed out gating concurrent credential writes"); + await new Promise((resolve) => setTimeout(resolve, 5)); + } + } + await Promise.all(saved); + + const store = JSON.parse(readFileSync(path, "utf8")); + for (const server of servers) { + expect(store[`${server.name}|${server.url}`]?.tokens?.access_token).toBe(`token-${server.name}`); + } + } finally { + const exited = children.map( + (child) => + new Promise((resolve) => { + if (child.exitCode !== null || child.signalCode !== null) resolve(); + else child.once("exit", () => resolve()); + }), + ); + for (const child of children) child.kill(); + await Promise.all(exited); + } + }, 20000); });