diff --git a/apps/sim/lib/core/security/input-validation.server.ts b/apps/sim/lib/core/security/input-validation.server.ts index dc9764990ba..cfef31a08d2 100644 --- a/apps/sim/lib/core/security/input-validation.server.ts +++ b/apps/sim/lib/core/security/input-validation.server.ts @@ -452,14 +452,20 @@ export function createPinnedFetch( * caller with a defined connection lifetime (e.g. a long-lived MCP transport) can * tear the Agent down on close instead of waiting for its idle timeout. Closing * the Agent is what releases any pooled keep-alive / HTTP/2 sockets it holds. + * + * `maxResponseSize` caps the (decoded) response body in bytes and makes undici reject + * with `UND_ERR_RES_EXCEEDED_MAX_SIZE` once exceeded — a DoS backstop for one-shot + * callers reading from a URL taken from untrusted metadata. Omit it (the default) to + * leave the response unbounded, which streaming consumers like the MCP transport need. */ export function createPinnedFetchWithDispatcher( resolvedIP: string, - options?: { allowH2?: boolean } + options?: { allowH2?: boolean; maxResponseSize?: number } ): { fetch: typeof fetch; dispatcher: Agent } { const dispatcher = new Agent({ allowH2: options?.allowH2 ?? false, connect: { lookup: createPinnedLookup(resolvedIP) }, + ...(options?.maxResponseSize !== undefined ? { maxResponseSize: options.maxResponseSize } : {}), }) const pinned = async (input: RequestInfo | URL, init?: RequestInit): Promise => { diff --git a/apps/sim/lib/mcp/oauth/probe.test.ts b/apps/sim/lib/mcp/oauth/probe.test.ts index d691f1178c7..23551694cc2 100644 --- a/apps/sim/lib/mcp/oauth/probe.test.ts +++ b/apps/sim/lib/mcp/oauth/probe.test.ts @@ -3,20 +3,30 @@ */ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { mockCreatePinnedFetch, mockCreateSsrfGuardedMcpFetch, mockPinnedFetch, mockGuardedFetch } = - vi.hoisted(() => { - const mockPinnedFetch = vi.fn() - const mockGuardedFetch = vi.fn() - return { - mockPinnedFetch, - mockGuardedFetch, - mockCreatePinnedFetch: vi.fn(() => mockPinnedFetch), - mockCreateSsrfGuardedMcpFetch: vi.fn(() => mockGuardedFetch), - } - }) +const { + mockCreatePinnedFetchWithDispatcher, + mockCreateSsrfGuardedMcpFetch, + mockPinnedFetch, + mockGuardedFetch, + mockDestroy, +} = vi.hoisted(() => { + const mockPinnedFetch = vi.fn() + const mockGuardedFetch = vi.fn() + const mockDestroy = vi.fn(() => Promise.resolve()) + return { + mockPinnedFetch, + mockGuardedFetch, + mockDestroy, + mockCreatePinnedFetchWithDispatcher: vi.fn(() => ({ + fetch: mockPinnedFetch, + dispatcher: { destroy: mockDestroy }, + })), + mockCreateSsrfGuardedMcpFetch: vi.fn(() => mockGuardedFetch), + } +}) vi.mock('@/lib/core/security/input-validation.server', () => ({ - createPinnedFetch: mockCreatePinnedFetch, + createPinnedFetchWithDispatcher: mockCreatePinnedFetchWithDispatcher, })) vi.mock('@/lib/mcp/pinned-fetch', () => ({ createSsrfGuardedMcpFetch: mockCreateSsrfGuardedMcpFetch, @@ -48,7 +58,7 @@ describe('detectMcpAuthType — connection pinning (SSRF / DNS-rebinding)', () = const authType = await detectMcpAuthType('https://rebind.example.com/mcp', '203.0.113.10') expect(authType).toBe('none') - expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10') + expect(mockCreatePinnedFetchWithDispatcher).toHaveBeenCalledWith('203.0.113.10') expect(mockCreateSsrfGuardedMcpFetch).not.toHaveBeenCalled() expect(mockPinnedFetch).toHaveBeenCalledTimes(1) // The unpinned global fetch must never be used — that was the SSRF sink. @@ -62,7 +72,7 @@ describe('detectMcpAuthType — connection pinning (SSRF / DNS-rebinding)', () = expect(authType).toBe('none') expect(mockCreateSsrfGuardedMcpFetch).toHaveBeenCalledTimes(1) - expect(mockCreatePinnedFetch).not.toHaveBeenCalled() + expect(mockCreatePinnedFetchWithDispatcher).not.toHaveBeenCalled() expect(mockGuardedFetch).toHaveBeenCalledTimes(1) expect(globalFetchSpy).not.toHaveBeenCalled() }) @@ -88,7 +98,7 @@ describe('detectMcpAuthType — connection pinning (SSRF / DNS-rebinding)', () = const authType = await detectMcpAuthType('http://example.com/mcp', '203.0.113.10') expect(authType).toBe('headers') - expect(mockCreatePinnedFetch).not.toHaveBeenCalled() + expect(mockCreatePinnedFetchWithDispatcher).not.toHaveBeenCalled() expect(mockCreateSsrfGuardedMcpFetch).not.toHaveBeenCalled() expect(globalFetchSpy).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/mcp/oauth/probe.ts b/apps/sim/lib/mcp/oauth/probe.ts index d6aac4d19c1..fc23cd7efbb 100644 --- a/apps/sim/lib/mcp/oauth/probe.ts +++ b/apps/sim/lib/mcp/oauth/probe.ts @@ -1,7 +1,7 @@ import { extractWWWAuthenticateParams } from '@modelcontextprotocol/sdk/client/auth.js' import type { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js' import { createLogger } from '@sim/logger' -import { createPinnedFetch } from '@/lib/core/security/input-validation.server' +import { createPinnedFetchWithDispatcher } from '@/lib/core/security/input-validation.server' import { isLoopbackHostname } from '@/lib/core/utils/urls' import { createSsrfGuardedMcpFetch } from '@/lib/mcp/pinned-fetch' import type { McpAuthType } from '@/lib/mcp/types' @@ -33,12 +33,17 @@ export async function detectMcpAuthType( return 'headers' } - const probeFetch: FetchLike = resolvedIP - ? createPinnedFetch(resolvedIP) - : createSsrfGuardedMcpFetch() + // When the caller pre-validated the IP we pin to it directly and own the Agent's + // lifetime; the SSRF-guarded fetch (used when we must validate ourselves) manages its + // own per-request Agent teardown internally. + const pinned = resolvedIP ? createPinnedFetchWithDispatcher(resolvedIP) : undefined + const probeFetch: FetchLike = pinned?.fetch ?? createSsrfGuardedMcpFetch() const controller = new AbortController() const timer = setTimeout(() => controller.abort(), PROBE_TIMEOUT_MS) + // Best-effort session cleanup runs after we return the classification; the pinned + // Agent must outlive it, so we tear it down once cleanup settles. + let sessionClose: Promise = Promise.resolve() try { const res = await probeFetch(url, { @@ -63,7 +68,7 @@ export async function detectMcpAuthType( const sessionId = res.headers.get('mcp-session-id') if (sessionId) { - void closeMcpSession(url, sessionId, probeFetch) + sessionClose = closeMcpSession(url, sessionId, probeFetch) } if (res.status === 401) { @@ -82,6 +87,11 @@ export async function detectMcpAuthType( return 'headers' } finally { clearTimeout(timer) + // Destroy (not close) so a hung request can't make teardown itself hang; wait for + // the best-effort session-close hop first so it isn't aborted mid-flight. + if (pinned) { + void sessionClose.finally(() => pinned.dispatcher.destroy().catch(() => {})) + } } } diff --git a/apps/sim/lib/mcp/oauth/revoke.test.ts b/apps/sim/lib/mcp/oauth/revoke.test.ts index ccd0caba98b..beb3a9dd0ab 100644 --- a/apps/sim/lib/mcp/oauth/revoke.test.ts +++ b/apps/sim/lib/mcp/oauth/revoke.test.ts @@ -31,7 +31,10 @@ const { })) vi.mock('@/lib/core/security/input-validation.server', () => ({ - createPinnedFetch: vi.fn(() => mockUndiciFetch), + createPinnedFetchWithDispatcher: vi.fn(() => ({ + fetch: mockUndiciFetch, + dispatcher: { destroy: vi.fn(() => Promise.resolve()) }, + })), })) vi.mock('@/lib/mcp/domain-check', () => ({ validateMcpServerSsrf: mockValidateMcpServerSsrf, diff --git a/apps/sim/lib/mcp/pinned-fetch.test.ts b/apps/sim/lib/mcp/pinned-fetch.test.ts index f4fd72e6741..415a2a1b358 100644 --- a/apps/sim/lib/mcp/pinned-fetch.test.ts +++ b/apps/sim/lib/mcp/pinned-fetch.test.ts @@ -4,14 +4,20 @@ import { sleep } from '@sim/utils/helpers' import { beforeEach, describe, expect, it, vi } from 'vitest' -const { mockCreatePinnedFetch, mockValidateMcpServerSsrf, sentinelFetch } = vi.hoisted(() => ({ - mockCreatePinnedFetch: vi.fn(), +const { + mockCreatePinnedFetchWithDispatcher, + mockValidateMcpServerSsrf, + sentinelFetch, + mockDestroy, +} = vi.hoisted(() => ({ + mockCreatePinnedFetchWithDispatcher: vi.fn(), mockValidateMcpServerSsrf: vi.fn(), sentinelFetch: vi.fn(), + mockDestroy: vi.fn(), })) vi.mock('@/lib/core/security/input-validation.server', () => ({ - createPinnedFetch: mockCreatePinnedFetch, + createPinnedFetchWithDispatcher: mockCreatePinnedFetchWithDispatcher, })) vi.mock('@/lib/mcp/domain-check', () => ({ validateMcpServerSsrf: mockValidateMcpServerSsrf, @@ -19,11 +25,18 @@ vi.mock('@/lib/mcp/domain-check', () => ({ import { createSsrfGuardedMcpFetch } from '@/lib/mcp/pinned-fetch' +/** The per-request pinned Agent is always built with a DoS-backstop response cap. */ +const withResponseCap = expect.objectContaining({ maxResponseSize: expect.any(Number) }) + describe('createSsrfGuardedMcpFetch', () => { beforeEach(() => { vi.clearAllMocks() - mockCreatePinnedFetch.mockReturnValue(sentinelFetch) - sentinelFetch.mockResolvedValue(new Response('ok')) + mockDestroy.mockResolvedValue(undefined) + mockCreatePinnedFetchWithDispatcher.mockReturnValue({ + fetch: sentinelFetch, + dispatcher: { destroy: mockDestroy }, + }) + sentinelFetch.mockImplementation(async () => new Response('ok')) }) it('validates each request URL and pins to the resolved IP', async () => { @@ -32,13 +45,140 @@ describe('createSsrfGuardedMcpFetch', () => { await fetchLike('https://attacker.example/revoke', { method: 'POST' }) expect(mockValidateMcpServerSsrf).toHaveBeenCalledWith('https://attacker.example/revoke') - expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10') + // The pinned Agent is always built with the DoS-backstop response-size cap. + expect(mockCreatePinnedFetchWithDispatcher).toHaveBeenCalledWith( + '203.0.113.10', + withResponseCap + ) expect(sentinelFetch).toHaveBeenCalledWith( 'https://attacker.example/revoke', expect.objectContaining({ method: 'POST', signal: expect.any(AbortSignal) }) ) }) + it('relabels an oversized response to a descriptive McpError', async () => { + mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') + // undici surfaces the cap breach as a fetch TypeError with a coded cause. + sentinelFetch.mockRejectedValue( + Object.assign(new TypeError('fetch failed'), { + cause: { code: 'UND_ERR_RES_EXCEEDED_MAX_SIZE' }, + }) + ) + const fetchLike = createSsrfGuardedMcpFetch() + + await expect(fetchLike('https://as.example/token', { method: 'POST' })).rejects.toThrow( + /exceeded \d+ bytes/ + ) + expect(mockDestroy).toHaveBeenCalledTimes(1) + }) + + it('tears down the per-request pinned Agent after a successful request', async () => { + mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') + const fetchLike = createSsrfGuardedMcpFetch() + await fetchLike('https://attacker.example/token', { method: 'POST' }) + + expect(mockDestroy).toHaveBeenCalledTimes(1) + }) + + it('tears down the pinned Agent even when the request fails', async () => { + mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') + sentinelFetch.mockRejectedValue(new Error('socket hang up')) + const fetchLike = createSsrfGuardedMcpFetch() + + await expect(fetchLike('https://attacker.example/token')).rejects.toThrow('socket hang up') + expect(mockDestroy).toHaveBeenCalledTimes(1) + }) + + it('returns a detached, in-memory copy of the response body', async () => { + mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') + sentinelFetch.mockImplementation( + async () => + new Response(JSON.stringify({ token_endpoint: 'https://as.example/token' }), { + headers: { 'content-type': 'application/json' }, + }) + ) + const fetchLike = createSsrfGuardedMcpFetch() + const res = await fetchLike('https://as.example/.well-known/oauth-authorization-server') + + // The body is readable even though the underlying socket/Agent is already destroyed. + expect(mockDestroy).toHaveBeenCalledTimes(1) + await expect(res.json()).resolves.toEqual({ token_endpoint: 'https://as.example/token' }) + }) + + it('reconstructs a null-body (204) response without throwing', async () => { + mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') + // A 204 has no body; the detached copy must not pass an (empty) body to Response. + sentinelFetch.mockImplementation(async () => new Response(null, { status: 204 })) + const fetchLike = createSsrfGuardedMcpFetch() + const res = await fetchLike('https://as.example/revoke', { method: 'POST' }) + + expect(res.status).toBe(204) + expect(mockDestroy).toHaveBeenCalledTimes(1) + }) + + it('returns a streaming response live (un-buffered) over the unpinned fallback', async () => { + // resolvedIP null → global fetch; a text/event-stream reply (the auth-type probe) + // must be handed back as-is so the caller reads headers without draining the stream. + // Identity (same object) proves it was NOT re-wrapped into a buffered copy. + mockValidateMcpServerSsrf.mockResolvedValue(null) + const streamingRes = new Response(new ReadableStream({ start() {} }), { + headers: { 'content-type': 'text/event-stream' }, + }) + const globalFetch = vi.spyOn(globalThis, 'fetch').mockImplementation(async () => streamingRes) + try { + const fetchLike = createSsrfGuardedMcpFetch() + const res = await fetchLike('https://allowed.internal/mcp', { method: 'POST' }) + + expect(res).toBe(streamingRes) + // No per-request Agent on the unpinned path, so nothing to tear down. + expect(mockDestroy).not.toHaveBeenCalled() + } finally { + globalFetch.mockRestore() + } + }) + + it('streams (does not buffer) a pinned text/event-stream reply and tears down after it drains', async () => { + // The guard resolves the IP itself, so the probe's initialize over the guarded path + // DOES get a pinned Agent. A streaming reply must still be handed back live (not + // buffered — that could stall/misclassify), with the Agent torn down once it drains. + mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') + const sseRes = new Response( + new ReadableStream({ + start(c) { + c.enqueue(new TextEncoder().encode('event: ready\ndata: {}\n\n')) + c.close() + }, + }), + { headers: { 'content-type': 'text/event-stream', 'mcp-session-id': 'sess-1' } } + ) + sentinelFetch.mockImplementation(async () => sseRes) + const fetchLike = createSsrfGuardedMcpFetch() + const res = await fetchLike('https://mcp.example/mcp', { method: 'POST' }) + + // Live (tee'd, not a buffered copy), headers preserved for the probe's classification. + expect(res).not.toBe(sseRes) + expect(res.headers.get('content-type')).toBe('text/event-stream') + expect(res.headers.get('mcp-session-id')).toBe('sess-1') + // Teardown happens in the background once the stream drains. + await vi.waitFor(() => expect(mockDestroy).toHaveBeenCalledTimes(1)) + }) + + it('buffers (does not return live) a non-streaming JSON response', async () => { + // Contrast with the streaming case: a JSON body is re-wrapped into a detached copy, + // so the returned object is NOT the original. + mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') + const jsonRes = new Response('{"ok":true}', { + headers: { 'content-type': 'application/json' }, + }) + sentinelFetch.mockImplementation(async () => jsonRes) + const fetchLike = createSsrfGuardedMcpFetch() + const res = await fetchLike('https://as.example/token', { method: 'POST' }) + + expect(res).not.toBe(jsonRes) + await expect(res.json()).resolves.toEqual({ ok: true }) + expect(mockDestroy).toHaveBeenCalledTimes(1) + }) + it('attaches an abort signal to every guarded request even without a caller signal', async () => { mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') const fetchLike = createSsrfGuardedMcpFetch() @@ -67,6 +207,22 @@ describe('createSsrfGuardedMcpFetch', () => { await expect(fetchLike('https://slow.example/token', { method: 'POST' })).rejects.toThrow( /timed out after 5ms/ ) + expect(mockDestroy).toHaveBeenCalledTimes(1) + }) + + it('bounds a stalled response body read by the deadline, not just time-to-headers', async () => { + mockValidateMcpServerSsrf.mockResolvedValue('203.0.113.10') + // Headers arrive immediately, but the body never completes — the exact shape of the + // "Connecting… forever" hang. The deadline must still fire. + sentinelFetch.mockImplementation( + async () => new Response(new ReadableStream({ start() {} })) + ) + const fetchLike = createSsrfGuardedMcpFetch(5) + + await expect(fetchLike('https://slow-body.example/token', { method: 'POST' })).rejects.toThrow( + /timed out after 5ms/ + ) + expect(mockDestroy).toHaveBeenCalledTimes(1) }) it('bounds a stalled SSRF/DNS validation by the deadline', async () => { @@ -75,9 +231,10 @@ describe('createSsrfGuardedMcpFetch', () => { const fetchLike = createSsrfGuardedMcpFetch(5) await expect(fetchLike('https://slow-dns.example/token')).rejects.toThrow(/timed out after 5ms/) - // Never got past validation, so no request was issued. - expect(mockCreatePinnedFetch).not.toHaveBeenCalled() + // Never got past validation, so no request was issued and no Agent was created. + expect(mockCreatePinnedFetchWithDispatcher).not.toHaveBeenCalled() expect(sentinelFetch).not.toHaveBeenCalled() + expect(mockDestroy).not.toHaveBeenCalled() }) it('does not orphan the validation promise when the signal is already aborted', async () => { @@ -91,7 +248,7 @@ describe('createSsrfGuardedMcpFetch', () => { await expect( fetchLike('https://slow.example/token', { signal: controller.signal }) ).rejects.toThrow('pre-aborted') - expect(mockCreatePinnedFetch).not.toHaveBeenCalled() + expect(mockCreatePinnedFetchWithDispatcher).not.toHaveBeenCalled() // Let the swallowed validation rejection settle so a leak would surface here. await sleep(0) }) @@ -105,7 +262,7 @@ describe('createSsrfGuardedMcpFetch', () => { controller.abort(new Error('caller cancelled')) await expect(pending).rejects.toThrow('caller cancelled') - expect(mockCreatePinnedFetch).not.toHaveBeenCalled() + expect(mockCreatePinnedFetchWithDispatcher).not.toHaveBeenCalled() }) it('propagates a caller-initiated abort unchanged (composed with the deadline)', async () => { @@ -137,7 +294,7 @@ describe('createSsrfGuardedMcpFetch', () => { await expect( fetchLike('http://169.254.169.254/latest/meta-data/', { method: 'POST' }) ).rejects.toThrow('blocked') - expect(mockCreatePinnedFetch).not.toHaveBeenCalled() + expect(mockCreatePinnedFetchWithDispatcher).not.toHaveBeenCalled() expect(sentinelFetch).not.toHaveBeenCalled() }) @@ -147,14 +304,27 @@ describe('createSsrfGuardedMcpFetch', () => { await fetchLike(new URL('https://attacker.example/discover')) expect(mockValidateMcpServerSsrf).toHaveBeenCalledWith('https://attacker.example/discover') - expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10') + expect(mockCreatePinnedFetchWithDispatcher).toHaveBeenCalledWith( + '203.0.113.10', + withResponseCap + ) }) it('falls back to global fetch when validation returns no IP', async () => { mockValidateMcpServerSsrf.mockResolvedValue(null) - const fetchLike = createSsrfGuardedMcpFetch() - await fetchLike('https://allowed.internal/mcp') + const globalFetch = vi + .spyOn(globalThis, 'fetch') + .mockImplementation(async () => new Response('ok')) + try { + const fetchLike = createSsrfGuardedMcpFetch() + await fetchLike('https://allowed.internal/mcp') - expect(mockCreatePinnedFetch).not.toHaveBeenCalled() + expect(mockCreatePinnedFetchWithDispatcher).not.toHaveBeenCalled() + expect(globalFetch).toHaveBeenCalledTimes(1) + // No pinned Agent was created, so there is nothing to tear down. + expect(mockDestroy).not.toHaveBeenCalled() + } finally { + globalFetch.mockRestore() + } }) }) diff --git a/apps/sim/lib/mcp/pinned-fetch.ts b/apps/sim/lib/mcp/pinned-fetch.ts index 0d12fcf0192..563296f10c3 100644 --- a/apps/sim/lib/mcp/pinned-fetch.ts +++ b/apps/sim/lib/mcp/pinned-fetch.ts @@ -1,8 +1,6 @@ import type { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js' -import { - createPinnedFetch, - createPinnedFetchWithDispatcher, -} from '@/lib/core/security/input-validation.server' +import type { Agent } from 'undici' +import { createPinnedFetchWithDispatcher } from '@/lib/core/security/input-validation.server' import { validateMcpServerSsrf } from '@/lib/mcp/domain-check' import { McpError } from '@/lib/mcp/types' @@ -47,15 +45,35 @@ export function createPinnedMcpFetch(resolvedIP: string): PinnedMcpFetch { const OAUTH_FETCH_TIMEOUT_MS = 30_000 /** - * Awaits `promise` but rejects with the signal's reason if `signal` aborts first. - * Bounds `dns.lookup`-based SSRF validation (which accepts no signal) by the - * composed deadline + caller signal. Removes the abort listener once `promise` - * settles so a late abort can't surface as an unhandled rejection. + * Cap on a guarded OAuth response body. Discovery/registration/token/refresh/revocation + * replies are always well under 1 KB; this ceiling is purely a DoS backstop so a + * malicious authorization server (reached via attacker-controllable metadata URLs) + * can't stream a multi-GB body within the deadline and exhaust memory. undici rejects + * with `UND_ERR_RES_EXCEEDED_MAX_SIZE` once the decoded body exceeds it. */ -function raceWithSignal(promise: Promise, signal: AbortSignal): Promise { +const MAX_OAUTH_RESPONSE_BYTES = 1_048_576 + +/** True for undici's response-too-large rejection, however it's wrapped by `fetch`. */ +function isResponseTooLarge(error: unknown): boolean { + const e = error as { code?: string; cause?: { code?: string } } | null + return ( + e?.code === 'UND_ERR_RES_EXCEEDED_MAX_SIZE' || + e?.cause?.code === 'UND_ERR_RES_EXCEEDED_MAX_SIZE' + ) +} + +/** + * Bounds `promise` by the composed deadline/caller `signal`, rejecting with the signal's + * reason if it aborts first. Holds every guarded phase inside the deadline — the + * `dns.lookup`-based SSRF validation (which takes no signal of its own), the HTTP + * request, and the response body read. Either path attaches a rejection handler to + * `promise` (the pre-aborted `.catch`, or `.then(resolve, reject)`), so a settlement + * arriving after the deadline has fired — once we've stopped awaiting it and are tearing + * the request down — can't surface as an unhandled rejection. The abort listener is + * removed once `promise` settles. + */ +function withDeadline(promise: Promise, signal: AbortSignal): Promise { if (signal.aborted) { - // The promise is already in flight; adopt its settlement so a later rejection - // (SSRF/DNS failure) can't surface as an unhandled rejection once we've aborted. promise.catch(() => {}) return Promise.reject(signal.reason) } @@ -76,27 +94,97 @@ function raceWithSignal(promise: Promise, signal: AbortSignal): Promise } /** - * Builds a `FetchLike` that validates every outbound request URL against the - * MCP SSRF policy before issuing it, then pins the connection to the resolved - * IP. Unlike the live transport — where the server URL is validated once up - * front — OAuth discovery and RFC 7009 revocation follow URLs taken verbatim - * from attacker-controllable authorization-server metadata - * (`authorization_servers`, `token_endpoint`, `revocation_endpoint`, …). Each - * such hop must be re-validated, so this guard runs `validateMcpServerSsrf` - * per request and rejects private/reserved/loopback targets (honoring - * `ALLOWED_MCP_DOMAINS` and self-hosted localhost rules). + * Reads a guarded OAuth response fully under the deadline, then returns a detached, + * in-memory copy. The MCP SDK reads discovery/registration/token/refresh bodies + * lazily — AFTER `auth()`'s injected fetch resolves — and applies no timeout of its + * own, so returning the live response would let the body read escape every deadline + * (the "Connecting… forever" hang). Buffering here brings the body read inside the + * wall-clock `signal`: undici's `bodyTimeout` only measures idle gaps between chunks + * and cannot bound a slow-drip or stalled body. These bodies are always small JSON. + */ +async function bufferUnderDeadline(response: Response, signal: AbortSignal): Promise { + const body = await withDeadline(response.arrayBuffer(), signal) + const headers = new Headers(response.headers) + // The buffered copy is decoded and detached from the socket; drop framing headers + // that would misdescribe it. + headers.delete('content-encoding') + headers.delete('content-length') + const nullBody = response.status === 204 || response.status === 205 || response.status === 304 + return new Response(nullBody ? null : body, { + status: response.status, + statusText: response.statusText, + headers, + }) +} + +/** + * Hands a streaming guarded response back live instead of buffering it. The only + * streaming reply a guarded call can see is the auth-type probe's `initialize` + * (`text/event-stream`), and the probe classifies from headers alone — buffering it + * would drain the stream or stall on a server that holds it open (misclassifying auth). + * The per-request pinned Agent (if any) is torn down once the stream ends, the caller + * cancels it, or the deadline aborts, so the socket is never stranded; the `tee` keeps + * the returned body fully readable meanwhile. Teardown ownership moves here, out of the + * caller's `finally`. + */ +function releaseStreamOnSettle( + response: Response, + dispatcher: Agent | undefined, + signal: AbortSignal +): Response { + if (!dispatcher || !response.body) { + void dispatcher?.destroy().catch(() => {}) + return response + } + const [drain, passthrough] = response.body.tee() + void (async () => { + const reader = drain.getReader() + const onAbort = () => void reader.cancel().catch(() => {}) + if (signal.aborted) onAbort() + else signal.addEventListener('abort', onAbort, { once: true }) + try { + while (!(await reader.read()).done) { + // Drain to end-of-stream so the Agent can be torn down once the reply completes. + } + } catch { + // Aborted or errored — the teardown below still runs. + } finally { + signal.removeEventListener('abort', onAbort) + void dispatcher.destroy().catch(() => {}) + } + })() + return new Response(passthrough, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }) +} + +/** + * Builds a `FetchLike` for one-shot MCP OAuth calls (discovery, dynamic client + * registration, token exchange/refresh, RFC 7009 revocation). It validates every + * outbound URL against the MCP SSRF policy before issuing it, then pins the + * connection to the resolved IP. Unlike the live transport — where the server URL is + * validated once up front — these hops follow URLs taken verbatim from + * attacker-controllable authorization-server metadata (`authorization_servers`, + * `token_endpoint`, `revocation_endpoint`, …), so each is re-validated and + * private/reserved/loopback targets are rejected (honoring `ALLOWED_MCP_DOMAINS` and + * self-hosted localhost rules). * - * Each request is bounded by a `timeoutMs` deadline via `AbortSignal.timeout`, - * composed with any caller-provided signal so cancellation still works. Only our - * own deadline is relabeled to an `McpError`; a caller abort or any other failure - * propagates unchanged. + * Three correctness guarantees, each of which the MCP SDK does NOT provide itself (it + * sets no timeout on any OAuth leg and reads bodies lazily): + * - **Bounded end-to-end.** The `timeoutMs` deadline (`AbortSignal.timeout`) composed + * with any caller signal covers SSRF validation, the request, AND the response body + * read — the body is buffered here so the SDK's later read can't outlive the + * deadline. Only our own deadline is relabeled to an `McpError`; a caller abort or + * any other failure propagates unchanged. + * - **No leaked sockets.** The per-request pinned Agent is `destroy()`ed on every path, + * releasing the keep-alive socket a one-shot flow would otherwise strand. + * - **Detached response.** The returned `Response` is an in-memory copy, safe to read + * after the underlying socket is gone. * - * Both the deadline and any caller-provided signal cover the whole guarded call — - * SSRF validation (whose `dns.lookup` takes no signal, so it's raced against the - * composed signal) and the HTTP request — so a caller awaiting this never waits - * past `timeoutMs` and can cancel at any point, including mid-validation. A - * stalled DNS resolution still runs to completion in the background, but its - * result is discarded. + * A stalled DNS resolution still runs to completion in the background, but its result + * is discarded. * * @param timeoutMs Per-request deadline in ms (defaults to 30s; override for tests). * @throws McpSsrfError if a request URL resolves to a blocked IP address @@ -106,22 +194,54 @@ export function createSsrfGuardedMcpFetch(timeoutMs: number = OAUTH_FETCH_TIMEOU return (async (url, init) => { const target = typeof url === 'string' ? url : url.href const timeoutSignal = AbortSignal.timeout(timeoutMs) - // Compose deadline + caller signal up front so both phases — SSRF validation - // and the HTTP request — are bounded by the deadline and caller cancellation. + // Compose deadline + caller signal up front so every phase — SSRF validation, the + // HTTP request, and the body read — is bounded by the deadline and caller cancellation. const signal = init?.signal ? AbortSignal.any([init.signal, timeoutSignal]) : timeoutSignal + // The per-request pinned Agent MUST be torn down: it holds a keep-alive socket the + // one-shot OAuth flow never reuses, so leaving it open leaks a socket per leg. + let dispatcher: Agent | undefined try { - const resolvedIP = await raceWithSignal(validateMcpServerSsrf(target), signal) - const pinnedFetch: FetchLike = resolvedIP ? createPinnedFetch(resolvedIP) : globalThis.fetch - return await pinnedFetch(url, { ...init, signal }) + const resolvedIP = await withDeadline(validateMcpServerSsrf(target), signal) + let response: Response + if (resolvedIP) { + const pinned = createPinnedFetchWithDispatcher(resolvedIP, { + maxResponseSize: MAX_OAUTH_RESPONSE_BYTES, + }) + dispatcher = pinned.dispatcher + response = await withDeadline(pinned.fetch(url, { ...init, signal }), signal) + } else { + // No pin (self-hosted allowlist) — global fetch over the shared dispatcher. + response = await withDeadline(globalThis.fetch(url, { ...init, signal }), signal) + } + // A `text/event-stream` reply is the auth-type probe's `initialize` (the only + // streaming case a guarded call sees); hand it back live with its own teardown so + // the caller reads headers without the buffer draining/stalling the stream. Every + // OAuth leg is single-shot JSON and falls through to be buffered and torn down. + const contentType = response.headers.get('content-type') ?? '' + if (contentType.includes('text/event-stream')) { + const streamed = releaseStreamOnSettle(response, dispatcher, signal) + dispatcher = undefined // teardown ownership handed to releaseStreamOnSettle + return streamed + } + return await bufferUnderDeadline(response, signal) } catch (error) { + const host = URL.canParse(target) ? new URL(target).host : target // Relabel only when our own deadline is what fired — identified by the // rejection reason's identity, not init.signal's state (which may abort // independently just after the deadline). if (timeoutSignal.aborted && error === timeoutSignal.reason) { - const host = URL.canParse(target) ? new URL(target).host : target throw new McpError(`MCP authorization request to ${host} timed out after ${timeoutMs}ms`) } + if (isResponseTooLarge(error)) { + throw new McpError( + `MCP authorization response from ${host} exceeded ${MAX_OAUTH_RESPONSE_BYTES} bytes` + ) + } throw error + } finally { + // Destroy (not close) so a hung leg can't make teardown itself hang; this releases + // the pooled keep-alive socket the per-request Agent would otherwise strand. + await dispatcher?.destroy().catch(() => {}) } }) satisfies FetchLike }