diff --git a/.changeset/mcp-promql-response-cap.md b/.changeset/mcp-promql-response-cap.md new file mode 100644 index 0000000000..bd6042062e --- /dev/null +++ b/.changeset/mcp-promql-response-cap.md @@ -0,0 +1,5 @@ +--- +'@hyperdx/api': patch +--- + +fix: Refuse Prometheus endpoint responses over 8 MB in MCP PromQL label lookups and queries, instead of buffering them whole diff --git a/packages/api/src/controllers/__tests__/promql.test.ts b/packages/api/src/controllers/__tests__/promql.test.ts index c992b85312..770a55c275 100644 --- a/packages/api/src/controllers/__tests__/promql.test.ts +++ b/packages/api/src/controllers/__tests__/promql.test.ts @@ -1,4 +1,8 @@ -import { buildPrometheusRequest, runPromqlQuery } from '@/controllers/promql'; +import { + buildPrometheusRequest, + PROMETHEUS_MAX_RESPONSE_BYTES, + runPromqlQuery, +} from '@/controllers/promql'; const mockFetch = jest.mocked(global.fetch); @@ -182,4 +186,49 @@ describe('runPromqlQuery', () => { ).rejects.toMatchObject({ category: 'user' }); expect(mockFetch).not.toHaveBeenCalled(); }); + + describe('response size cap', () => { + it('reports a response over the cap as a server error naming the cap', async () => { + mockFetch.mockResolvedValueOnce( + new Response('{', { + headers: { + 'content-length': String(PROMETHEUS_MAX_RESPONSE_BYTES + 1), + }, + }), + ); + + await expect(runPromqlQuery(prometheus, range)).rejects.toMatchObject({ + message: + 'Prometheus at http://prom:9090/api/v1/query_range returned more than 8 MB. Narrow the query or time range.', + category: 'server', + }); + }); + + it('reads a chunked body under the cap as before', async () => { + const json = JSON.stringify({ status: 'success', data: matrix }); + const [head, tail] = [json.slice(0, 10), json.slice(10)]; + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(head)); + controller.enqueue(new TextEncoder().encode(tail)); + controller.close(); + }, + }); + mockFetch.mockResolvedValueOnce(new Response(body)); + + await expect(runPromqlQuery(prometheus, range)).resolves.toEqual(matrix); + }); + + it('reports a body that is not JSON by its HTTP status', async () => { + mockFetch.mockResolvedValueOnce( + new Response('Bad gateway', { status: 502 }), + ); + + await expect(runPromqlQuery(prometheus, range)).rejects.toMatchObject({ + message: + 'Prometheus at http://prom:9090/api/v1/query_range returned HTTP 502', + category: 'server', + }); + }); + }); }); diff --git a/packages/api/src/controllers/promql.ts b/packages/api/src/controllers/promql.ts index 96894609db..e68d9a084b 100644 --- a/packages/api/src/controllers/promql.ts +++ b/packages/api/src/controllers/promql.ts @@ -4,6 +4,7 @@ import { } from '@hyperdx/common-utils/dist/clickhouse'; import { getCurrentQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import type { PromqlQuery } from '@hyperdx/common-utils/dist/core/promqlChart'; +import { parseJSON } from '@hyperdx/common-utils/dist/core/utils'; import { SourceKind, type SQLInterval } from '@hyperdx/common-utils/dist/types'; import { z } from 'zod'; @@ -20,12 +21,28 @@ import { queryRangeViaTableFunction, } from '@/controllers/timeseriesEngine'; import { withSpan } from '@/utils/instrumentation'; +import { + readBoundedBody, + ResponseTooLargeError, +} from '@/utils/readBoundedBody'; type PromqlBackend = 'prometheus' | 'clickhouse'; const PROMETHEUS_FETCH_TIMEOUT_MS = 30_000; const PROMETHEUS_CH_TIMEOUT_MS = 30_000; export const PROMETHEUS_MAX_RESOLUTION = 11_000; +/** + * Largest Prometheus response body fetchPrometheusData will read into memory. + * + * Queries to ClickHouse are already capped by `max_result_rows` + * (PROMETHEUS_MAX_RESULT_ROWS). An external Prometheus has no equivalent: + * `/query` and `/query_range` take no row limit, and older servers ignore + * `limit` on the label APIs. So the only way to bound the memory a response + * can take is to stop reading at a byte limit. 8 MB fits roughly + * PROMETHEUS_MAX_RESULT_ROWS label values, so both backends allow about the + * same amount of data. + */ +export const PROMETHEUS_MAX_RESPONSE_BYTES = 8 * 1024 * 1024; // ClickHouse serves the Prometheus HTTP API under this prefix on its main HTTP // port when the `prometheus_api_v1` handler is configured (26.6+). @@ -176,9 +193,20 @@ async function fetchPrometheusData( ); } - const body = prometheusBodySchema.safeParse( - await resp.json().catch(() => undefined), - ); + let text: string; + try { + text = await readBoundedBody(resp, PROMETHEUS_MAX_RESPONSE_BYTES); + } catch (err) { + const tooLargeMessage = `Prometheus at ${redactedTarget} returned more than ${PROMETHEUS_MAX_RESPONSE_BYTES / (1024 * 1024)} MB. Narrow the query or time range.`; + const genericMessage = `Failed to read the response from Prometheus at ${redactedTarget}`; + throw new PromQLQueryError( + err instanceof ResponseTooLargeError ? tooLargeMessage : genericMessage, + 'server', + ); + } + // A body that is not JSON (a proxy's HTML error page, say) is reported by + // its HTTP status below. + const body = prometheusBodySchema.safeParse(parseJSON(text)); if (!resp.ok || !body.success || body.data.status !== 'success') { const isUserError = resp.status >= 400 && diff --git a/packages/api/src/mcp/__tests__/promqlQueryTool.int.test.ts b/packages/api/src/mcp/__tests__/promqlQueryTool.int.test.ts index 54fe5e2f2e..d9918ba018 100644 --- a/packages/api/src/mcp/__tests__/promqlQueryTool.int.test.ts +++ b/packages/api/src/mcp/__tests__/promqlQueryTool.int.test.ts @@ -1,6 +1,7 @@ import { Metadata } from '@hyperdx/common-utils/dist/core/metadata'; import { SourceKind } from '@hyperdx/common-utils/dist/types'; +import { PROMETHEUS_MAX_RESPONSE_BYTES } from '@/controllers/promql'; import { DEFAULT_DATABASE, dropTimeSeriesTable, @@ -307,6 +308,27 @@ describe('MCP clickstack_promql', () => { expect(getFirstText(result)).not.toContain('secret'); }); + it('reports a response too large to buffer as a server error', async () => { + mockFetch.mockResolvedValueOnce( + prometheusResponse( + { resultType: 'matrix', result: [] }, + { + headers: { + 'content-length': String(PROMETHEUS_MAX_RESPONSE_BYTES + 1), + }, + }, + ), + ); + const source = await createPrometheusSource(); + + const result = await promql({ sourceId: source._id.toString() }); + + expect(result.isError).toBe(true); + expect(getFirstText(result)).toContain( + 'returned more than 8 MB. Narrow the query or time range.', + ); + }); + it('rejects a result a chart cannot draw as a user error', async () => { mockFetch.mockResolvedValueOnce( prometheusResponse({ resultType: 'string', result: [START, 'hello'] }), diff --git a/packages/api/src/mcp/__tests__/promqlSources.int.test.ts b/packages/api/src/mcp/__tests__/promqlSources.int.test.ts index 6ec22dd529..acfc5c9da6 100644 --- a/packages/api/src/mcp/__tests__/promqlSources.int.test.ts +++ b/packages/api/src/mcp/__tests__/promqlSources.int.test.ts @@ -2,6 +2,7 @@ import { SourceKind } from '@hyperdx/common-utils/dist/types'; import { Client } from '@modelcontextprotocol/sdk/client/index.js'; import * as config from '@/config'; +import { PROMETHEUS_MAX_RESPONSE_BYTES } from '@/controllers/promql'; import { DEFAULT_DATABASE, dropTimeSeriesTable, @@ -18,8 +19,8 @@ const mockFetch = jest.mocked(global.fetch); const TABLE = 'mcp_promql_sources'; -const prometheusResponse = (data: unknown) => - new Response(JSON.stringify({ status: 'success', data })); +const prometheusResponse = (data: unknown, init?: ResponseInit) => + new Response(JSON.stringify({ status: 'success', data }), init); describe('MCP PromQL sources', () => { const server = getServer(); @@ -228,6 +229,26 @@ describe('MCP PromQL sources', () => { expect(output.error).toContain('http://prom.example.com/api/v1/'); expect(output.error).not.toContain('secret'); }); + + it('still returns the source config when a label list is too large to buffer', async () => { + const source = await createPrometheusSource(); + mockFetch.mockImplementation(async () => + prometheusResponse(['up'], { + headers: { + 'content-length': String(PROMETHEUS_MAX_RESPONSE_BYTES + 1), + }, + }), + ); + + const result = await callTool(client, 'clickstack_describe_source', { + sourceId: source._id.toString(), + }); + + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.partial).toBe(true); + expect(output.error).toContain('returned more than 8 MB'); + }); }); describe('with PromQL disabled', () => { diff --git a/packages/api/src/utils/__tests__/readBoundedBody.test.ts b/packages/api/src/utils/__tests__/readBoundedBody.test.ts new file mode 100644 index 0000000000..32ee993801 --- /dev/null +++ b/packages/api/src/utils/__tests__/readBoundedBody.test.ts @@ -0,0 +1,77 @@ +import { + readBoundedBody, + ResponseTooLargeError, +} from '@/utils/readBoundedBody'; + +const MAX = 4 * 1024; + +const streamOf = ( + pull: (controller: ReadableStreamDefaultController) => void, + onCancel: () => void, +) => + new ReadableStream({ + pull, + cancel() { + onCancel(); + }, + }); + +describe('readBoundedBody', () => { + it('refuses a declared Content-Length over the cap without reading the body', async () => { + let cancelled = false; + const body = streamOf( + c => c.enqueue(new TextEncoder().encode('{')), + () => (cancelled = true), + ); + const resp = new Response(body, { + headers: { 'content-length': String(MAX + 1) }, + }); + + await expect(readBoundedBody(resp, MAX)).rejects.toBeInstanceOf( + ResponseTooLargeError, + ); + expect(cancelled).toBe(true); + expect(resp.bodyUsed).toBe(true); + }); + + it('stops reading an undeclared-length body at the cap', async () => { + const chunk = new Uint8Array(1024).fill(0x20); + let pulled = 0; + let cancelled = false; + const body = streamOf( + c => { + pulled++; + c.enqueue(chunk); + }, + () => (cancelled = true), + ); + + await expect( + readBoundedBody(new Response(body), MAX), + ).rejects.toBeInstanceOf(ResponseTooLargeError); + expect(cancelled).toBe(true); + // One chunk per KB, plus the one that crossed the cap and the stream's + // read-ahead. + expect(pulled).toBeLessThanOrEqual(MAX / chunk.byteLength + 2); + }); + + it('joins a chunked body under the cap into one string', async () => { + const text = 'héllo wörld'; + const bytes = new TextEncoder().encode(text); + const body = new ReadableStream({ + start(c) { + c.enqueue(bytes.slice(0, 3)); + c.enqueue(bytes.slice(3)); + c.close(); + }, + }); + + await expect(readBoundedBody(new Response(body), MAX)).resolves.toBe(text); + }); + + it('returns an empty string for a response with no body', async () => { + await expect( + readBoundedBody(new Response(null, { status: 204 }), MAX), + ).resolves.toBe(''); + }); +}); diff --git a/packages/api/src/utils/readBoundedBody.ts b/packages/api/src/utils/readBoundedBody.ts new file mode 100644 index 0000000000..f4718036fb --- /dev/null +++ b/packages/api/src/utils/readBoundedBody.ts @@ -0,0 +1,41 @@ +/** Thrown by readBoundedBody when a response is larger than its cap. */ +export class ResponseTooLargeError extends Error { + constructor(maxBytes: number) { + super(`Response body exceeds ${maxBytes} bytes`); + this.name = 'ResponseTooLargeError'; + } +} + +/** + * Reads a fetch Response body as text without ever holding more than + * `maxBytes` of it. A declared Content-Length over the cap is refused before + * any of the body is read; an undeclared (chunked) body is read until the cap + * and then cancelled. Use it for bodies from hosts we do not control, where + * `resp.text()` / `resp.json()` would buffer whatever the server sends. + */ +export async function readBoundedBody( + resp: Response, + maxBytes: number, +): Promise { + const declared = Number(resp.headers.get('content-length')); + if (declared > maxBytes) { + await resp.body?.cancel(); + throw new ResponseTooLargeError(maxBytes); + } + if (!resp.body) return ''; + + const reader = resp.body.getReader(); + const chunks: Uint8Array[] = []; + let received = 0; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + received += value.byteLength; + if (received > maxBytes) { + await reader.cancel(); + throw new ResponseTooLargeError(maxBytes); + } + chunks.push(value); + } + return Buffer.concat(chunks).toString('utf8'); +}