From 411a5e63be1e6e48232bbae399b0e284d866b6d8 Mon Sep 17 00:00:00 2001 From: Alix-007 Date: Mon, 29 Jun 2026 08:06:32 +0800 Subject: [PATCH] fix(zai): bound Z.AI endpoint-probe error body reads to prevent OOM (#97540) Co-authored-by: Vincent Koc --- extensions/zai/detect.test.ts | 151 ++++++++++++++++++++++++++++++++++ extensions/zai/detect.ts | 10 ++- 2 files changed, 160 insertions(+), 1 deletion(-) diff --git a/extensions/zai/detect.test.ts b/extensions/zai/detect.test.ts index 74f34b175df..b89a85a9167 100644 --- a/extensions/zai/detect.test.ts +++ b/extensions/zai/detect.test.ts @@ -1,10 +1,77 @@ // Zai tests cover detect plugin behavior. import { MAX_TIMER_TIMEOUT_MS } from "openclaw/plugin-sdk/number-runtime"; +import { readResponseWithLimit } from "openclaw/plugin-sdk/response-limit-runtime"; import { afterEach, describe, expect, it, vi } from "vitest"; import { detectZaiEndpoint } from "./detect.js"; type FetchResponse = { status: number; body?: unknown }; +const ZAI_DETECT_ERROR_BODY_MAX_BYTES = 16 * 1024 * 1024; + +/** + * Builds a streaming error Response whose body is far larger than the 16 MiB cap. + * Tracks how many bytes were actually pulled and whether the consumer cancelled + * the stream, so tests can prove the read is bounded (fail-closed) rather than + * draining the whole untrusted body into memory. + */ +function makeOversizedStreamFetch(params: { + url: string; + status: number; + chunkBytes?: number; + hardCeilingBytes?: number; +}) { + const chunkBytes = params.chunkBytes ?? 1024 * 1024; + const hardCeilingBytes = params.hardCeilingBytes ?? 64 * 1024 * 1024; + const state = { enqueuedBytes: 0, cancelled: false }; + + const fetchFn = (async (url: string) => { + if (url !== params.url) { + throw new Error(`unexpected url: ${url}`); + } + const body = new ReadableStream({ + pull(controller) { + if (state.enqueuedBytes >= hardCeilingBytes) { + // Safety stop: with an unbounded reader this point would be reached + // (and the test would fail on the bounded-bytes assertion below). + controller.close(); + return; + } + state.enqueuedBytes += chunkBytes; + controller.enqueue(new Uint8Array(chunkBytes)); + }, + cancel() { + state.cancelled = true; + }, + }); + return new Response(body, { + status: params.status, + headers: { "content-type": "application/json" }, + }); + }) as typeof fetch; + + return { fetchFn, state }; +} + +/** + * Builds a fetch returning a single raw (possibly non-JSON) error body, keyed by + * `${url}::${model}`. Used to drive the new bounded decode path with small, + * well-formed, empty, and malformed sub-cap bodies that must behave exactly as + * the previous `res.json()` path did. + */ +function makeRawBodyFetch(map: Record) { + return (async (url: string, init?: RequestInit) => { + const rawBody = typeof init?.body === "string" ? JSON.parse(init.body) : null; + const entry = map[`${url}::${rawBody?.model ?? ""}`] ?? map[url]; + if (!entry) { + throw new Error(`unexpected url: ${url} model=${String(rawBody?.model ?? "")}`); + } + return new Response(entry.raw, { + status: entry.status, + headers: { "content-type": "application/json" }, + }); + }) as typeof fetch; +} + function makeFetch(map: Record) { return (async (url: string, init?: RequestInit) => { const rawBody = typeof init?.body === "string" ? JSON.parse(init.body) : null; @@ -184,4 +251,88 @@ describe("detectZaiEndpoint", () => { expect(timeoutSpy).toHaveBeenCalledWith(expect.any(Function), MAX_TIMER_TIMEOUT_MS); }); + + it("still parses well-formed sub-cap error bodies to drive endpoint classification", async () => { + // Happy path: model-not-found errors must still be decoded from the bounded + // body so the probe classifies them as unsupported and walks to the GLM-4.7 + // fallback. The error message that drives classification lives only inside + // the body, so a passing fallback proves the new bounded reader decoded it. + const codingGlobal = "https://api.z.ai/api/coding/paas/v4/chat/completions"; + const detected = await detectZaiEndpoint({ + apiKey: "sk-test", // pragma: allowlist secret + endpoint: "coding-global", + fetchFn: makeRawBodyFetch({ + [`${codingGlobal}::glm-5.2`]: { + status: 400, + raw: JSON.stringify({ error: { message: "model not found for this plan" } }), + }, + [`${codingGlobal}::glm-5.1`]: { + status: 400, + raw: JSON.stringify({ code: 1211, msg: "model does not exist" }), + }, + [`${codingGlobal}::glm-4.7`]: { status: 200, raw: "{}" }, + }), + }); + + expect(detected?.endpoint).toBe("coding-global"); + expect(detected?.modelId).toBe("glm-4.7"); + }); + + it("swallows malformed and empty sub-cap error bodies and falls back on status", async () => { + // Regression: a non-JSON or empty error body must not throw out of the + // probe. JSON.parse fails, the existing try/catch swallows it, and the + // probe degrades to status-only classification (404 => unsupported model), + // so the GLM-4.7 fallback still resolves exactly as before. + const codingGlobal = "https://api.z.ai/api/coding/paas/v4/chat/completions"; + const detected = await detectZaiEndpoint({ + apiKey: "sk-test", // pragma: allowlist secret + endpoint: "coding-global", + fetchFn: makeRawBodyFetch({ + [`${codingGlobal}::glm-5.2`]: { status: 404, raw: "gateway error" }, + [`${codingGlobal}::glm-5.1`]: { status: 404, raw: "" }, + [`${codingGlobal}::glm-4.7`]: { status: 200, raw: "{}" }, + }), + }); + + expect(detected?.endpoint).toBe("coding-global"); + expect(detected?.modelId).toBe("glm-4.7"); + }); + + it("fails closed on oversized probe error bodies without buffering unbounded", async () => { + const { fetchFn, state } = makeOversizedStreamFetch({ + url: "https://api.z.ai/api/paas/v4/chat/completions", + status: 400, + }); + + const detected = await detectZaiEndpoint({ + apiKey: "sk-test", // pragma: allowlist secret + endpoint: "global", + fetchFn, + }); + + // Probe swallows the bounded-read overflow and falls back to status-only, + // so the oversized error body cannot promote this endpoint. + expect(detected).toBeNull(); + // The stream was cancelled (fail-closed) instead of being drained to the + // 64 MiB safety ceiling, proving the read stops near the 16 MiB cap. + expect(state.cancelled).toBe(true); + expect(state.enqueuedBytes).toBeLessThanOrEqual( + ZAI_DETECT_ERROR_BODY_MAX_BYTES + 2 * 1024 * 1024, + ); + }); + + it("rejects oversized bodies via the shared bounded reader the probe uses", async () => { + const { fetchFn } = makeOversizedStreamFetch({ + url: "https://api.z.ai/api/paas/v4/chat/completions", + status: 400, + }); + const res = await fetchFn("https://api.z.ai/api/paas/v4/chat/completions"); + + await expect( + readResponseWithLimit(res, ZAI_DETECT_ERROR_BODY_MAX_BYTES, { + onOverflow: ({ maxBytes }) => + new Error(`Z.AI probe error body exceeded size limit (${maxBytes} bytes)`), + }), + ).rejects.toThrow(/exceeded size limit/); + }); }); diff --git a/extensions/zai/detect.ts b/extensions/zai/detect.ts index c3f063bb7b6..c88999b3715 100644 --- a/extensions/zai/detect.ts +++ b/extensions/zai/detect.ts @@ -1,5 +1,6 @@ // Zai plugin module implements detect behavior. import { resolveTimerTimeoutMs } from "openclaw/plugin-sdk/number-runtime"; +import { readResponseWithLimit } from "openclaw/plugin-sdk/response-limit-runtime"; import { ZAI_CN_BASE_URL, ZAI_CODING_CN_BASE_URL, @@ -36,6 +37,9 @@ type ProbeCandidate = ZaiDetectedEndpoint & { const UNSUPPORTED_MODEL_ERROR_CODES = new Set(["1211", "1311"]); +/** Cap for the Z.AI probe error body; bounds untrusted error responses to avoid unbounded buffering/OOM. */ +const ZAI_DETECT_ERROR_BODY_MAX_BYTES = 16 * 1024 * 1024; + function isUnsupportedModelResult(result: ProbeResult): boolean { if (result.ok) { return false; @@ -106,7 +110,11 @@ async function probeZaiChatCompletions(params: { let errorCode: string | undefined; let errorMessage: string | undefined; try { - const json = (await res.json()) as { + const bytes = await readResponseWithLimit(res, ZAI_DETECT_ERROR_BODY_MAX_BYTES, { + onOverflow: ({ maxBytes }) => + new Error(`Z.AI probe error body exceeded size limit (${maxBytes} bytes)`), + }); + const json = JSON.parse(new TextDecoder().decode(bytes)) as { error?: { code?: unknown; message?: unknown }; code?: unknown; msg?: unknown;