fix(mcp): stream MCP response body with running cap (no 16 MiB pre-rejection allocation)#4051
Conversation
728088d to
3569aca
Compare
…llocations read_response_bytes_capped used response.bytes().await — which buffers the FULL body in memory before checking the cap. A server that omits Content-Length (chunked transfer) forces a fresh ~16 MiB allocation per request before rejection, creating real memory pressure under concurrent abuse. The audit of #3926 flagged this as a MEDIUM exploitation surface. Switch to streaming via reqwest::Response::chunk() (already in the crate's reqwest version, no extra feature flag). Loop: match response.chunk().await { Ok(Some(chunk)) if running + chunk > cap => err and abort Ok(Some(chunk)) => extend buf Ok(None) => break Err(e) => err } Pre-allocate Vec capacity from Content-Length when present, clamped to MAX_RESPONSE_BYTES so a malicious large hint can't bypass the fix. When Content-Length is missing the Vec grows organically; the running counter still aborts the moment the next chunk would breach the cap. Worst-case allocation is now MAX_RESPONSE_BYTES + 1 chunk size, not the buffered 16 MiB ceiling regardless of Content-Length.
…ed on 32-bit Fast-path `content_length as usize > MAX_RESPONSE_BYTES` truncates a u64 to usize on 32-bit / wasm32 targets, so a malicious Content-Length: 0x1_0000_0001 casts to 1 and slips past the "no bytes read" guarantee. The streaming loop's running counter saves us from OOM, but the fast-path semantic is lost on non-64-bit builds (notably the wasm32 sandbox tests). Compare in u64 instead. The `hint.min(MAX_RESPONSE_BYTES) as usize` reserve has the same shape — clamp in u64 before the cast for the same reason.
873da5b to
f097854
Compare
houko
left a comment
houko
left a comment
There was a problem hiding this comment.
LGTM — streaming + running cap fixes the pre-rejection allocation cleanly. Cross-file check: minimax.rs:286/578 and gemini.rs:135 use the same name MAX_RESPONSE_BYTES but they're post-parse field-size checks on already-extracted base64 strings, not body buffering — different shape, out of scope here.
Two minor nits, neither a blocker:
-
response.content_length()is called twice (fast-path reject + reserve hint). Cheap call, but the value could be bound once for clarity. -
PR description says "worst-case allocation is now MAX_RESPONSE_BYTES + 1 chunk size regardless of header presence" — strictly speaking, when an honest
Content-Lengthis present we stillreserve(min(cl, MAX)), so the upper bound is the same MAX as before. The real win is the no-header / chunked-transfer path, which previously allocated up to MAX unconditionally.
1 participant
Follow-up to #3926 (cap SSE body).
Problem
read_response_bytes_cappedusedresponse.bytes().await— which buffers the full body in memory before the length check:A server omitting
Content-Length(HTTP/1.1 chunked transfer) forces a fresh ~16 MiB allocation per request before rejection. Under concurrent abuse this creates real memory pressure — the audit of #3926 flagged it as a MEDIUM exploitation surface.Fix
Stream via
reqwest::Response::chunk()(already available in the crate's reqwest version, no new feature flag):Veccapacity fromContent-Lengthwhen present, clamped toMAX_RESPONSE_BYTESso a malicious large hint can't bypass the streaming fix.Content-Lengthis missing theVecgrows organically; the running counter aborts the moment the next chunk would breach the cap.MAX_RESPONSE_BYTES + 1 chunk sizeregardless of header presence — not the unconditional 16 MiB ceiling.