From 2f02a5f7f36ecd27f5f643afb0ccbf74a00a124b Mon Sep 17 00:00:00 2001 From: Chukwuebuka-2003 Date: Mon, 3 Aug 2026 00:44:14 +0000 Subject: [PATCH 1/7] fix(proxy): cap decompressed request body size to prevent decompression-bomb DoS _read_request_body_bytes expands Content-Encoding bodies (zstd/gzip/ deflate/br) with uncapped one-shot calls (gzip.decompress, zlib.decompress, brotli.decompress, stream_reader().read()), so a small compressed request could balloon into an unbounded in-memory buffer and OOM the proxy. Decompress incrementally in bounded slices against a new MAX_DECOMPRESSED_BODY_BYTES cap (mirrors MAX_REQUEST_BODY_SIZE) and raise RequestBodyTooLarge (a ValueError, so existing 400 paths are unchanged) when the running total exceeds it. gzip now goes through zlib's bounded decompressobj API instead of gzip.decompress. --- headroom/proxy/helpers.py | 127 +++++++++++-- tests/test_proxy_request_decompression_cap.py | 172 ++++++++++++++++++ 2 files changed, 288 insertions(+), 11 deletions(-) create mode 100644 tests/test_proxy_request_decompression_cap.py diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index faa1d0ed5..bf33d0c96 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -504,6 +504,14 @@ MAX_REQUEST_BODY_SIZE = 100 * 1024 * 1024 # Maximum SSE buffer size (10MB - prevents memory exhaustion from malformed streams) MAX_SSE_BUFFER_SIZE = 10 * 1024 * 1024 +# Hard ceiling on the *decompressed* size of request bodies. Compressed +# bodies (zstd/gzip/deflate/br) are expanded against this cap so a tiny +# compressed request cannot balloon into an unbounded in-memory buffer +# (decompression-bomb DoS). Mirrors MAX_REQUEST_BODY_SIZE: any payload that +# would exceed the uncompressed request budget is rejected outright. Read at +# call time so tests can shrink it without a module reload. +MAX_DECOMPRESSED_BODY_BYTES = MAX_REQUEST_BODY_SIZE + # Per-event SSE size cap (PR-A8 / P1-8). Configurable via # HEADROOM_SSE_BUFFER_MAX_BYTES. Guards against pathological huge events # (a single event > 1 MB by default is treated as an upstream protocol bug @@ -1889,12 +1897,99 @@ def apply_session_sticky_ccr_tool( return tools_out, True +class RequestBodyTooLarge(ValueError): + """A compressed request body expanded beyond the decompression cap. + + Subclasses ``ValueError`` so existing callers that translate + decompression failures into a client error (400) keep working. The + pristine message matters: it names the limit instead of the wrapped + decompressor error. + """ + + +def _enforce_decompression_cap(size: int, label: str) -> None: + """Raise :class:`RequestBodyTooLarge` if *size* exceeds the cap.""" + max_output_bytes = MAX_DECOMPRESSED_BODY_BYTES + if size > max_output_bytes: + raise RequestBodyTooLarge( + f"{label} request body exceeds the {max_output_bytes}-byte decompression limit" + ) + + +def _decompress_capped(decompressor: Any, data: bytes, *, label: str) -> bytes: + """Incrementally decompress *data* (zlib ``decompressobj``) under a hard cap. + + ``zlib.decompressobj.decompress(data, max_length)`` returns at most + ``max_length`` bytes per call and parks unconsumed input on + ``unconsumed_tail``, so feeding in bounded slices and checking the + running total keeps peak memory under ``MAX_DECOMPRESSED_BODY_BYTES``. + A decompression bomb becomes a clean :class:`RequestBodyTooLarge` + instead of an OOM kill. + """ + out = bytearray() + remaining = data + while remaining and not decompressor.eof: + chunk = decompressor.decompress(remaining, 64 * 1024) + out.extend(chunk) + _enforce_decompression_cap(len(out), label) + remaining = decompressor.unconsumed_tail + # Drain output still buffered inside the decompressor. flush() also + # surfaces truncated-stream errors, matching the old one-shot calls. + while True: + chunk = decompressor.flush(64 * 1024) + if not chunk: + break + out.extend(chunk) + _enforce_decompression_cap(len(out), label) + return bytes(out) + + +def _decompress_zstd_capped(zstandard: Any, data: bytes) -> bytes: + """Decompress a zstd frame via its streaming reader, capped on output. + + ``ZstdDecompressor.stream_reader`` can emit an unbounded output from a + small input; read it in bounded slices and enforce the same cap as the + zlib/brotli paths. + """ + reader = zstandard.ZstdDecompressor().stream_reader(data) + try: + out = bytearray() + while True: + chunk = reader.read(64 * 1024) + if not chunk: + break + out.extend(chunk) + _enforce_decompression_cap(len(out), "zstd") + return bytes(out) + finally: + reader.close() + + +def _decompress_brotli_capped(brotli: Any, data: bytes) -> bytes: + """Decompress a brotli stream incrementally, capped on output. + + ``brotli.decompress`` has no output cap and ``Decompressor.process`` + materializes everything it can from the given input, so feed the input + in bounded slices and enforce the cap on the running total. + """ + decompressor = brotli.Decompressor() + out = bytearray() + for i in range(0, len(data), 64 * 1024): + out.extend(decompressor.process(data[i : i + 64 * 1024])) + _enforce_decompression_cap(len(out), "br") + if not decompressor.is_finished(): + raise ValueError("Failed to decompress br request body: truncated or corrupt stream") + return bytes(out) + + async def _read_request_body_bytes(request: Request) -> bytes: """Read and (if needed) decompress the request body, returning raw UTF-8 bytes. Mirrors ``_read_request_json`` but returns the bytes pre-parse so forwarders can implement byte-faithful passthrough (PR-A3, fixes P0-2). - Raises ``ValueError`` on any decompression failure. + Raises ``ValueError`` on any decompression failure. Decompression is + bounded by ``MAX_DECOMPRESSED_BODY_BYTES`` so a small compressed body + cannot balloon into an unbounded in-memory buffer (decompression bomb). """ encoding = (request.headers.get("content-encoding") or "").lower().strip() raw = await request.body() @@ -1902,41 +1997,51 @@ async def _read_request_body_bytes(request: Request) -> bytes: if encoding in ("zstd", "zstandard"): try: import zstandard - - dctx = zstandard.ZstdDecompressor() - reader = dctx.stream_reader(raw) - raw = reader.read() - reader.close() except ImportError: raise ValueError( "Request body is zstd-compressed but the 'zstandard' package is not installed. " "Install it with: pip install zstandard" ) from None + try: + raw = _decompress_zstd_capped(zstandard, raw) + except RequestBodyTooLarge: + raise except Exception as exc: raise ValueError(f"Failed to decompress zstd request body: {exc}") from exc elif encoding == "gzip": - import gzip as _gzip + import zlib try: - raw = _gzip.decompress(raw) + # zlib's gzip mode (wbits=31) gives us a bounded incremental API; + # gzip.decompress() has no output cap and would materialize a + # decompression bomb in full before we could reject it. + decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16) + raw = _decompress_capped(decompressor, raw, label="gzip") + except RequestBodyTooLarge: + raise except Exception as exc: raise ValueError(f"Failed to decompress gzip request body: {exc}") from exc elif encoding == "deflate": import zlib try: - raw = zlib.decompress(raw) + decompressor = zlib.decompressobj() + raw = _decompress_capped(decompressor, raw, label="deflate") + except RequestBodyTooLarge: + raise except Exception as exc: raise ValueError(f"Failed to decompress deflate request body: {exc}") from exc elif encoding == "br": try: import brotli - - raw = brotli.decompress(raw) except ImportError: raise ValueError( "Request body is brotli-compressed but the 'brotli' package is not installed." ) from None + try: + raw = _decompress_brotli_capped(brotli, raw) + except RequestBodyTooLarge: + raise except Exception as exc: raise ValueError(f"Failed to decompress brotli request body: {exc}") from exc elif encoding and encoding != "identity": diff --git a/tests/test_proxy_request_decompression_cap.py b/tests/test_proxy_request_decompression_cap.py new file mode 100644 index 000000000..8029a11fa --- /dev/null +++ b/tests/test_proxy_request_decompression_cap.py @@ -0,0 +1,172 @@ +"""Decompression of compressed request bodies must be bounded. + +``_read_request_body_bytes`` (headroom/proxy/helpers.py) expands +zstd/gzip/deflate/br request bodies before forwarding. The expansion used +to be unbounded (``gzip.decompress`` / ``zlib.decompress`` / +``brotli.decompress`` / ``stream_reader().read()``), so a tiny compressed +body could balloon into an unbounded in-memory buffer — a decompression-bomb +DoS. Every format is now fed incrementally against +``MAX_DECOMPRESSED_BODY_BYTES`` and raises ``RequestBodyTooLarge`` (a +``ValueError``) when the cap is exceeded. +""" + +import asyncio +import gzip +import json +import zlib + +import pytest + +from headroom.proxy.helpers import ( + MAX_DECOMPRESSED_BODY_BYTES, + RequestBodyTooLarge, + _read_request_body_bytes, +) + +_PAYLOAD = b"the quick brown fox jumps over the lazy dog " * 5000 # ~225 KB + + +class _FakeHeaders: + def __init__(self, d=None): + self._d = {k.lower(): v for k, v in (d or {}).items()} + + def get(self, k, default=None): + return self._d.get(k.lower(), default) + + +class _FakeRequest: + def __init__(self, raw, headers=None): + self._raw = raw + self.headers = _FakeHeaders(headers) + + async def body(self): + return self._raw + + +def _read(raw, encoding): + return asyncio.run(_read_request_body_bytes(_FakeRequest(raw, {"content-encoding": encoding}))) + + +# --------------------------------------------------------------------------- +# Round-trips: compressed bodies still decompress byte-for-byte +# --------------------------------------------------------------------------- + + +def test_gzip_round_trip(): + raw = gzip.compress(_PAYLOAD) + assert _read(raw, "gzip") == _PAYLOAD + + +def test_deflate_round_trip(): + raw = zlib.compress(_PAYLOAD) + assert _read(raw, "deflate") == _PAYLOAD + + +def test_zstd_round_trip(): + zstandard = pytest.importorskip("zstandard") + raw = zstandard.ZstdCompressor().compress(_PAYLOAD) + assert _read(raw, "zstd") == _PAYLOAD + + +def test_brotli_round_trip(): + brotli = pytest.importorskip("brotli") + raw = brotli.compress(_PAYLOAD) + assert _read(raw, "br") == _PAYLOAD + + +# --------------------------------------------------------------------------- +# Decompression bombs: small compressed input expanding past the cap is +# rejected with RequestBodyTooLarge (a ValueError) instead of OOMing +# --------------------------------------------------------------------------- + + +def _bomb_under_cap(monkeypatch, compress): + # Shrink the cap so the test stays fast and memory-cheap: the bomb below + # expands 1 MB -> ~1 MB of output from ~1 KB of compressed input, which + # dwarfs the 4 KB test cap. + monkeypatch.setattr("headroom.proxy.helpers.MAX_DECOMPRESSED_BODY_BYTES", 4096) + bomb = compress(b"A" * (1024 * 1024)) + assert len(bomb) < 4096 # sanity: the input itself is tiny + return bomb + + +def test_gzip_bomb_rejected(monkeypatch): + with pytest.raises(RequestBodyTooLarge): + _read(_bomb_under_cap(monkeypatch, gzip.compress), "gzip") + + +def test_deflate_bomb_rejected(monkeypatch): + with pytest.raises(RequestBodyTooLarge): + _read(_bomb_under_cap(monkeypatch, zlib.compress), "deflate") + + +def test_zstd_bomb_rejected(monkeypatch): + zstandard = pytest.importorskip("zstandard") + with pytest.raises(RequestBodyTooLarge): + _read( + _bomb_under_cap(monkeypatch, zstandard.ZstdCompressor().compress), + "zstd", + ) + + +def test_brotli_bomb_rejected(monkeypatch): + brotli = pytest.importorskip("brotli") + with pytest.raises(RequestBodyTooLarge): + _read(_bomb_under_cap(monkeypatch, brotli.compress), "br") + + +def test_payload_at_cap_is_accepted(monkeypatch): + monkeypatch.setattr("headroom.proxy.helpers.MAX_DECOMPRESSED_BODY_BYTES", 1024) + payload = b"x" * 1024 + assert _read(gzip.compress(payload), "gzip") == payload + + +# --------------------------------------------------------------------------- +# Error behavior preserved +# --------------------------------------------------------------------------- + + +def test_corrupt_gzip_still_raises_value_error(): + with pytest.raises(ValueError): + _read(b"this is not a gzip stream at all", "gzip") + + +def test_truncated_brotli_still_raises_value_error(): + brotli = pytest.importorskip("brotli") + truncated = brotli.compress(_PAYLOAD)[:16] + with pytest.raises(ValueError): + _read(truncated, "br") + + +def test_unsupported_encoding_raises(): + with pytest.raises(ValueError): + _read(b"data", "lzma") + + +def test_identity_and_missing_encoding_passthrough(): + payload = b'{"model": "x"}' + assert _read(payload, "identity") == payload + assert _read(payload, "") == payload + + +# --------------------------------------------------------------------------- +# Full path: read_request_json_with_bytes still decodes compressed JSON +# --------------------------------------------------------------------------- + + +def test_read_request_json_with_bytes_decompresses_gzip(): + from headroom.proxy.helpers import read_request_json_with_bytes + + body = {"model": "x", "messages": [{"role": "user", "content": "hi"}]} + raw = gzip.compress(json.dumps(body).encode("utf-8")) + result, out_raw = asyncio.run( + read_request_json_with_bytes(_FakeRequest(raw, {"content-encoding": "gzip"})) + ) + assert result == body + assert json.loads(out_raw) == body + + +def test_cap_is_sane_default(): + # The default cap mirrors the uncompressed request budget: a compressed + # body must never expand beyond what the handler would accept anyway. + assert MAX_DECOMPRESSED_BODY_BYTES > 0 From 2a7f5854d221c0e3bb56912023baf23d69effd29 Mon Sep 17 00:00:00 2001 From: Chukwuebuka-2003 Date: Mon, 3 Aug 2026 14:12:44 +0000 Subject: [PATCH 2/7] fix(proxy): reject trailing data after gzip stream; close decompression-cap coverage gaps zlib's gzip mode stops at the end of the first member and silently ignores whatever follows, so multi-member gzip bodies were truncated and trailing garbage accepted where gzip.decompress() previously handled or rejected them. Raise on decompressor.unused_data instead. Restructure _decompress_capped around an explicit eof check: the flush loop's output branch is unreachable (unconsumed_tail always drains it), and the check makes truncated-stream detection explicit and testable. Add tests for the paths Codecov flagged: zstd/brotli ImportError handlers, corrupt zstd/deflate streams, truncated gzip, multi-member gzip, and trailing-garbage rejection. Signed-off-by: Chukwuebuka-2003 --- headroom/proxy/helpers.py | 25 ++++++--- tests/test_proxy_request_decompression_cap.py | 54 +++++++++++++++++++ 2 files changed, 71 insertions(+), 8 deletions(-) diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index bf33d0c96..63b5d629a 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1933,14 +1933,13 @@ def _decompress_capped(decompressor: Any, data: bytes, *, label: str) -> bytes: out.extend(chunk) _enforce_decompression_cap(len(out), label) remaining = decompressor.unconsumed_tail - # Drain output still buffered inside the decompressor. flush() also - # surfaces truncated-stream errors, matching the old one-shot calls. - while True: - chunk = decompressor.flush(64 * 1024) - if not chunk: - break - out.extend(chunk) - _enforce_decompression_cap(len(out), label) + # Input exhausted without the stream end marker: truncated. The old + # one-shot calls (gzip.decompress / zlib.decompress) raised on this; + # keep that behavior so the generic handler turns it into a 400. + # (The main loop always drains pending output via unconsumed_tail, so + # there is no buffered output left to flush here.) + if not decompressor.eof: + raise ValueError(f"truncated {label} stream") return bytes(out) @@ -2017,6 +2016,16 @@ async def _read_request_body_bytes(request: Request) -> bytes: # decompression bomb in full before we could reject it. decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16) raw = _decompress_capped(decompressor, raw, label="gzip") + if decompressor.unused_data: + # zlib stops at the end of the first gzip member and ignores + # whatever follows. gzip.decompress() used to handle + # multi-member files and reject trailing garbage; reject + # loudly here instead of silently forwarding a truncated + # body that differs from what the client sent. + raise ValueError( + f"unexpected trailing data after gzip stream " + f"({len(decompressor.unused_data)} bytes)" + ) except RequestBodyTooLarge: raise except Exception as exc: diff --git a/tests/test_proxy_request_decompression_cap.py b/tests/test_proxy_request_decompression_cap.py index 8029a11fa..3d6a2c7a4 100644 --- a/tests/test_proxy_request_decompression_cap.py +++ b/tests/test_proxy_request_decompression_cap.py @@ -13,6 +13,7 @@ DoS. Every format is now fed incrementally against import asyncio import gzip import json +import sys import zlib import pytest @@ -131,6 +132,59 @@ def test_corrupt_gzip_still_raises_value_error(): _read(b"this is not a gzip stream at all", "gzip") +def test_corrupt_deflate_still_raises_value_error(): + with pytest.raises(ValueError): + _read(b"this is not a deflate stream at all", "deflate") + + +def test_corrupt_zstd_still_raises_value_error(): + zstandard = pytest.importorskip("zstandard") + del zstandard # bytes only; the ImportError path is covered separately + with pytest.raises(ValueError): + _read(b"this is not a zstd frame at all", "zstd") + + +def test_zstd_not_installed_raises(monkeypatch): + monkeypatch.setitem(sys.modules, "zstandard", None) + with pytest.raises(ValueError, match="not installed"): + _read(b"ignored", "zstd") + + +def test_brotli_not_installed_raises(monkeypatch): + monkeypatch.setitem(sys.modules, "brotli", None) + with pytest.raises(ValueError, match="not installed"): + _read(b"ignored", "br") + + +def test_gzip_multi_member_rejected(): + # zlib stops at the first gzip member; anything after it must not be + # silently dropped (gzip.decompress() used to decompress all members). + m1 = gzip.compress(b"first member " * 100) + m2 = gzip.compress(b"second member " * 100) + with pytest.raises(ValueError, match="trailing data"): + _read(m1 + m2, "gzip") + + +def test_gzip_trailing_garbage_rejected(): + with pytest.raises(ValueError, match="trailing data"): + _read(gzip.compress(b"ok") + b"NOTGZIP", "gzip") + + +def test_truncated_gzip_still_raises_value_error(): + # Header + partial body, no end marker: the old one-shot call raised + # BadGzipFile; the incremental path must raise too. + truncated = gzip.compress(_PAYLOAD)[:32] + with pytest.raises(ValueError): + _read(truncated, "gzip") + + +def test_gzip_payload_across_slice_boundary_round_trips(): + # Exactly 64 KiB + a few bytes: exercises the multi-iteration path where + # the final output is emitted after the cap-sized slice. + payload = b"x" * (64 * 1024 + 7) + assert _read(gzip.compress(payload), "gzip") == payload + + def test_truncated_brotli_still_raises_value_error(): brotli = pytest.importorskip("brotli") truncated = brotli.compress(_PAYLOAD)[:16] From 5b1830d85e081c9dcf44564b0f3d8171621f63ff Mon Sep 17 00:00:00 2001 From: Chukwuebuka-2003 Date: Mon, 3 Aug 2026 14:16:24 +0000 Subject: [PATCH 3/7] fix(proxy): support multi-member gzip bodies under the decompression cap zlib's gzip mode stops at the first member and parks the remainder on unused_data, which the previous guard rejected outright. Restore gzip.decompress() parity by looping the members against one shared running total so the cap applies to the whole payload (not per member), rejecting only trailing bytes that are not another gzip member. _decompress_capped gains an optional out buffer so the cumulative total is enforced incrementally across members. Adds multi-member round-trip and cumulative-cap tests. Signed-off-by: Chukwuebuka-2003 --- headroom/proxy/helpers.py | 50 +++++++++++++------ tests/test_proxy_request_decompression_cap.py | 21 +++++--- 2 files changed, 49 insertions(+), 22 deletions(-) diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 63b5d629a..70628ce06 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1916,7 +1916,9 @@ def _enforce_decompression_cap(size: int, label: str) -> None: ) -def _decompress_capped(decompressor: Any, data: bytes, *, label: str) -> bytes: +def _decompress_capped( + decompressor: Any, data: bytes, *, label: str, out: bytearray | None = None +) -> bytes: """Incrementally decompress *data* (zlib ``decompressobj``) under a hard cap. ``zlib.decompressobj.decompress(data, max_length)`` returns at most @@ -1924,9 +1926,10 @@ def _decompress_capped(decompressor: Any, data: bytes, *, label: str) -> bytes: ``unconsumed_tail``, so feeding in bounded slices and checking the running total keeps peak memory under ``MAX_DECOMPRESSED_BODY_BYTES``. A decompression bomb becomes a clean :class:`RequestBodyTooLarge` - instead of an OOM kill. + instead of an OOM kill. Pass *out* to accumulate across calls (used + for multi-member gzip) so the cap applies to the whole payload. """ - out = bytearray() + out = out if out is not None else bytearray() remaining = data while remaining and not decompressor.eof: chunk = decompressor.decompress(remaining, 64 * 1024) @@ -1943,6 +1946,34 @@ def _decompress_capped(decompressor: Any, data: bytes, *, label: str) -> bytes: return bytes(out) +def _decompress_gzip_capped(data: bytes) -> bytes: + """Decompress *data*, including multi-member gzip, under the hard cap. + + zlib's gzip mode stops at the end of the first member and parks the + remainder on ``unused_data``, whereas ``gzip.decompress()`` used to + decompress every member. Loop the members against one shared running + total so the cap applies to the whole payload, and reject trailing + bytes that are not another gzip member (old behavior raised + ``BadGzipFile`` on garbage). + """ + import zlib + + out = bytearray() + remaining = data + while True: + decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16) + _decompress_capped(decompressor, remaining, label="gzip", out=out) + remaining = decompressor.unused_data + if not remaining: + break + if not remaining.startswith(b"\x1f\x8b"): + raise ValueError( + f"unexpected trailing data after gzip stream " + f"({len(remaining)} bytes)" + ) + return bytes(out) + + def _decompress_zstd_capped(zstandard: Any, data: bytes) -> bytes: """Decompress a zstd frame via its streaming reader, capped on output. @@ -2014,18 +2045,7 @@ async def _read_request_body_bytes(request: Request) -> bytes: # zlib's gzip mode (wbits=31) gives us a bounded incremental API; # gzip.decompress() has no output cap and would materialize a # decompression bomb in full before we could reject it. - decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16) - raw = _decompress_capped(decompressor, raw, label="gzip") - if decompressor.unused_data: - # zlib stops at the end of the first gzip member and ignores - # whatever follows. gzip.decompress() used to handle - # multi-member files and reject trailing garbage; reject - # loudly here instead of silently forwarding a truncated - # body that differs from what the client sent. - raise ValueError( - f"unexpected trailing data after gzip stream " - f"({len(decompressor.unused_data)} bytes)" - ) + raw = _decompress_gzip_capped(raw) except RequestBodyTooLarge: raise except Exception as exc: diff --git a/tests/test_proxy_request_decompression_cap.py b/tests/test_proxy_request_decompression_cap.py index 3d6a2c7a4..2be70f6af 100644 --- a/tests/test_proxy_request_decompression_cap.py +++ b/tests/test_proxy_request_decompression_cap.py @@ -156,13 +156,20 @@ def test_brotli_not_installed_raises(monkeypatch): _read(b"ignored", "br") -def test_gzip_multi_member_rejected(): - # zlib stops at the first gzip member; anything after it must not be - # silently dropped (gzip.decompress() used to decompress all members). - m1 = gzip.compress(b"first member " * 100) - m2 = gzip.compress(b"second member " * 100) - with pytest.raises(ValueError, match="trailing data"): - _read(m1 + m2, "gzip") +def test_gzip_multi_member_round_trips(): + # gzip.decompress() handled multi-member files; the incremental path + # must decompress every member against the same cap. + m1 = b"first member " * 100 + m2 = b"second member " * 100 + assert _read(gzip.compress(m1) + gzip.compress(m2), "gzip") == m1 + m2 + + +def test_gzip_multi_member_cumulative_cap(monkeypatch): + # The cap applies to the whole payload, not per member. + monkeypatch.setattr("headroom.proxy.helpers.MAX_DECOMPRESSED_BODY_BYTES", 4096) + member = gzip.compress(b"A" * 2048) # 2048 B decompressed each + with pytest.raises(RequestBodyTooLarge): + _read(member + member + member, "gzip") def test_gzip_trailing_garbage_rejected(): From a7061c7b4da69a1f263bfdbfc31fbd284d9ec8ff Mon Sep 17 00:00:00 2001 From: Chukwuebuka-2003 Date: Mon, 3 Aug 2026 22:18:13 +0000 Subject: [PATCH 4/7] style(proxy): satisfy ruff format on trailing-data error Signed-off-by: Chukwuebuka-2003 --- headroom/proxy/helpers.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 70628ce06..8df53d696 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1967,10 +1967,7 @@ def _decompress_gzip_capped(data: bytes) -> bytes: if not remaining: break if not remaining.startswith(b"\x1f\x8b"): - raise ValueError( - f"unexpected trailing data after gzip stream " - f"({len(remaining)} bytes)" - ) + raise ValueError(f"unexpected trailing data after gzip stream ({len(remaining)} bytes)") return bytes(out) From b4b3a02b056fc5a40367c299b779eb0b0b2ec5c1 Mon Sep 17 00:00:00 2001 From: Chukwuebuka-2003 Date: Tue, 4 Aug 2026 09:20:32 +0000 Subject: [PATCH 5/7] fix(proxy): reject br request bodies and return 413 on decompression-cap violations Addresses review feedback on the decompression-bomb boundary: - Brotli request bodies are rejected outright. The Python brotli bindings (brotli 1.2.0, brotlicffi) expose no output-bounded streaming API: Decompressor.process(input_slice) returns ALL output produced from a slice before the cap check runs, so a highly compressible sub-64 KiB stream could materialize far beyond MAX_DECOMPRESSED_BODY_BYTES inside a single call. Rejecting the encoding is the only way the process-wide boundary holds for br; clients can use gzip/deflate/zstd, which have bounded APIs. Removes the unsafe _decompress_brotli_capped helper entirely. - RequestBodyTooLarge is no longer translated to 400 invalid_json. Every handler path that reads a request body now returns the configured body-too-large status (get_body_too_large_status, default 413) with a request_too_large error code, before the generic ValueError catch. The Bedrock adapter no longer fail-opens on it: the compressed bomb is rejected instead of being forwarded verbatim upstream. The batch passthrough best-effort read re-raises it so the handler's 413 path sees it. Tests: br rejection + peak-input regression in test_proxy_request_decompression_cap.py; new test_proxy_decompression_cap_status.py covers the affected handler paths (OpenAI chat/responses, Anthropic messages, Bedrock invoke with no-fail-open assertion, batch create, Gemini) plus the configurable status and the batch best-effort audit. Signed-off-by: Chukwuebuka-2003 --- headroom/proxy/handlers/anthropic.py | 32 ++- headroom/proxy/handlers/batch.py | 40 ++- headroom/proxy/handlers/bedrock.py | 21 ++ headroom/proxy/handlers/gemini.py | 65 ++++- headroom/proxy/handlers/openai.py | 43 ++- headroom/proxy/helpers.py | 54 ++-- tests/test_proxy_decompression_cap_status.py | 263 ++++++++++++++++++ tests/test_proxy_request_decompression_cap.py | 51 +++- 8 files changed, 520 insertions(+), 49 deletions(-) create mode 100644 tests/test_proxy_decompression_cap_status.py diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index e255ab886..aa84b89e0 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -533,8 +533,10 @@ class AnthropicHandlerMixin: from headroom.proxy.helpers import ( MAX_MESSAGE_ARRAY_LENGTH, MAX_REQUEST_BODY_SIZE, + RequestBodyTooLarge, _get_image_compressor, compute_turn_id, + get_body_too_large_status, read_request_json_with_bytes, ) from headroom.proxy.modes import is_cache_mode, is_token_mode @@ -683,6 +685,18 @@ class AnthropicHandlerMixin: try: async with stage_timer.measure("read_request_json"): body, original_body_bytes = await read_request_json_with_bytes(request) + except RequestBodyTooLarge as e: + await _finalize_pre_upstream() + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "type": "error", + "error": { + "type": "request_too_large", + "message": f"{e!s}", + }, + }, + ) except (json.JSONDecodeError, ValueError) as e: await _finalize_pre_upstream() return JSONResponse( @@ -3819,7 +3833,12 @@ class AnthropicHandlerMixin: from fastapi.responses import JSONResponse, Response from headroom.ccr import CCRToolInjector - from headroom.proxy.helpers import MAX_REQUEST_BODY_SIZE, _read_request_json + from headroom.proxy.helpers import ( + MAX_REQUEST_BODY_SIZE, + RequestBodyTooLarge, + _read_request_json, + get_body_too_large_status, + ) from headroom.proxy.modes import is_cache_mode from headroom.utils import extract_user_query @@ -3843,6 +3862,17 @@ class AnthropicHandlerMixin: # Parse request try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "type": "error", + "error": { + "type": "request_too_large", + "message": f"{e!s}", + }, + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, diff --git a/headroom/proxy/handlers/batch.py b/headroom/proxy/handlers/batch.py index 3681294a4..a4e1c6229 100644 --- a/headroom/proxy/handlers/batch.py +++ b/headroom/proxy/handlers/batch.py @@ -53,7 +53,12 @@ class BatchHandlerMixin: from fastapi.responses import JSONResponse, Response from headroom.ccr import CCRToolInjector - from headroom.proxy.helpers import MAX_REQUEST_BODY_SIZE, _read_request_json + from headroom.proxy.helpers import ( + MAX_REQUEST_BODY_SIZE, + RequestBodyTooLarge, + _read_request_json, + get_body_too_large_status, + ) from headroom.utils import extract_user_query start_time = time.time() @@ -76,6 +81,17 @@ class BatchHandlerMixin: # Parse request try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "code": get_body_too_large_status(), + "message": f"{e!s}", + "status": "INVALID_ARGUMENT", + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, @@ -804,7 +820,11 @@ class BatchHandlerMixin: """ from fastapi.responses import JSONResponse, Response - from headroom.proxy.helpers import _read_request_json + from headroom.proxy.helpers import ( + RequestBodyTooLarge, + _read_request_json, + get_body_too_large_status, + ) start_time = time.time() request_id = await self._next_request_id() @@ -812,6 +832,17 @@ class BatchHandlerMixin: # Parse request try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "message": f"{e!s}", + "type": "invalid_request_error", + "code": "request_too_large", + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, @@ -1207,6 +1238,7 @@ class BatchHandlerMixin: from headroom.proxy.body_forwarding import prepare_outbound_body_bytes from headroom.proxy.helpers import ( + RequestBodyTooLarge, _read_request_body_bytes, _strip_internal_headers, log_outbound_headers, @@ -1230,8 +1262,12 @@ class BatchHandlerMixin: # Best effort: capture the original (decompressed) bytes so the # passthrough is truly byte-faithful. If the body was already # consumed upstream we fall through to canonical re-serialization. + # A size-policy rejection is never best-effort: the handler's 413 + # path must see it, so re-raise instead of swallowing it here. try: original_body_bytes: bytes | None = await _read_request_body_bytes(request) + except RequestBodyTooLarge: + raise except Exception: original_body_bytes = None diff --git a/headroom/proxy/handlers/bedrock.py b/headroom/proxy/handlers/bedrock.py index 688e1befb..c0d9a566c 100644 --- a/headroom/proxy/handlers/bedrock.py +++ b/headroom/proxy/handlers/bedrock.py @@ -76,9 +76,11 @@ class BedrockHandlerMixin: from headroom.proxy.helpers import ( COMPRESSION_TIMEOUT_SECONDS, MAX_MESSAGE_ARRAY_LENGTH, + RequestBodyTooLarge, _headroom_bypass_enabled, _strip_internal_headers, extract_tags, + get_body_too_large_status, read_request_json_with_bytes, ) from headroom.proxy.modes import is_cache_mode @@ -126,8 +128,27 @@ class BedrockHandlerMixin: # Read the body up front so we can fail open to a verbatim forward on any # parse error (a malformed body is the gateway's problem, not ours). + # A size-policy rejection (decompression cap) is NOT a parse error: + # it must never fail open or forward the expanding body, or the + # decompression-bomb boundary is meaningless. try: body, raw = await read_request_json_with_bytes(request) + except RequestBodyTooLarge as e: + logger.warning( + "[%s] %s rejecting request body past decompression cap: %s", + request_id, + LOG_TAG, + e, + ) + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "type": "request_too_large", + "message": f"{e!s}", + } + }, + ) except Exception as err: from starlette.requests import ClientDisconnect diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index 199a79fd1..9e0c51343 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -269,7 +269,12 @@ class GeminiHandlerMixin: from fastapi import HTTPException from fastapi.responses import JSONResponse, Response - from headroom.proxy.helpers import MAX_REQUEST_BODY_SIZE, _read_request_json + from headroom.proxy.helpers import ( + MAX_REQUEST_BODY_SIZE, + RequestBodyTooLarge, + _read_request_json, + get_body_too_large_status, + ) from headroom.utils import extract_user_query start_time = time.time() @@ -291,6 +296,16 @@ class GeminiHandlerMixin: # Parse request try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "message": f"{e!s}", + "code": get_body_too_large_status(), + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, @@ -828,7 +843,11 @@ class GeminiHandlerMixin: """Handle Pi/OpenClaw Google Cloud Code Assist and Antigravity streaming requests.""" from fastapi.responses import JSONResponse - from headroom.proxy.helpers import _read_request_json + from headroom.proxy.helpers import ( + RequestBodyTooLarge, + _read_request_json, + get_body_too_large_status, + ) from headroom.utils import extract_user_query start_time = time.time() @@ -836,6 +855,16 @@ class GeminiHandlerMixin: try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "message": f"{e!s}", + "code": get_body_too_large_status(), + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, @@ -994,7 +1023,11 @@ class GeminiHandlerMixin: """Handle Gemini streaming endpoint /v1beta/models/{model}:streamGenerateContent.""" from fastapi.responses import JSONResponse - from headroom.proxy.helpers import _read_request_json + from headroom.proxy.helpers import ( + RequestBodyTooLarge, + _read_request_json, + get_body_too_large_status, + ) start_time = time.time() request_id = await self._next_request_id() @@ -1002,6 +1035,16 @@ class GeminiHandlerMixin: # Parse request try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "message": f"{e!s}", + "code": get_body_too_large_status(), + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, @@ -1084,7 +1127,11 @@ class GeminiHandlerMixin: """ from fastapi.responses import JSONResponse, Response - from headroom.proxy.helpers import _read_request_json + from headroom.proxy.helpers import ( + RequestBodyTooLarge, + _read_request_json, + get_body_too_large_status, + ) from headroom.utils import extract_user_query start_time = time.time() @@ -1093,6 +1140,16 @@ class GeminiHandlerMixin: # Parse request try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "message": f"{e!s}", + "code": get_body_too_large_status(), + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 74bedd16a..0ef53ba47 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -2618,7 +2618,9 @@ class OpenAIHandlerMixin: COMPRESSION_TIMEOUT_SECONDS, MAX_MESSAGE_ARRAY_LENGTH, MAX_REQUEST_BODY_SIZE, + RequestBodyTooLarge, _read_request_json, + get_body_too_large_status, ) from headroom.proxy.modes import is_cache_mode, is_token_mode from headroom.utils import extract_user_query @@ -2651,6 +2653,17 @@ class OpenAIHandlerMixin: # Parse request try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "message": f"{e!s}", + "type": "invalid_request_error", + "code": "request_too_large", + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, @@ -4451,6 +4464,8 @@ class OpenAIHandlerMixin: from headroom.proxy.body_forwarding import BodyMutationTracker from headroom.proxy.helpers import ( MAX_REQUEST_BODY_SIZE, + RequestBodyTooLarge, + get_body_too_large_status, read_request_json_with_bytes, ) from headroom.utils import extract_user_query @@ -4487,6 +4502,17 @@ class OpenAIHandlerMixin: # (#1542); byte-faithful passthrough avoids that. try: body, original_body_bytes = await read_request_json_with_bytes(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "message": f"{e!s}", + "type": "invalid_request_error", + "code": "request_too_large", + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, @@ -8353,12 +8379,27 @@ class OpenAIHandlerMixin: """ from fastapi.responses import JSONResponse - from headroom.proxy.helpers import _read_request_json + from headroom.proxy.helpers import ( + RequestBodyTooLarge, + _read_request_json, + get_body_too_large_status, + ) # Check bypass header if request.headers.get("x-headroom-bypass", "").lower() == "true": try: body = await _read_request_json(request) + except RequestBodyTooLarge as e: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "message": f"{e!s}", + "type": "invalid_request_error", + "code": "request_too_large", + } + }, + ) except (json.JSONDecodeError, ValueError) as e: return JSONResponse( status_code=400, diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 8df53d696..5660177d9 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1900,8 +1900,13 @@ def apply_session_sticky_ccr_tool( class RequestBodyTooLarge(ValueError): """A compressed request body expanded beyond the decompression cap. - Subclasses ``ValueError`` so existing callers that translate - decompression failures into a client error (400) keep working. The + Subclasses ``ValueError`` so the exception always degrades to a client + error even if a caller forgets to handle it explicitly. Handlers that + read request bodies MUST translate this into + :func:`get_body_too_large_status` (default 413) — never a 400, and + never a fail-open forward: a size-policy rejection is a client error + distinct from malformed JSON, and the whole point of the cap is that + the expanded body must never be materialized or forwarded. The pristine message matters: it names the limit instead of the wrapped decompressor error. """ @@ -1992,23 +1997,6 @@ def _decompress_zstd_capped(zstandard: Any, data: bytes) -> bytes: reader.close() -def _decompress_brotli_capped(brotli: Any, data: bytes) -> bytes: - """Decompress a brotli stream incrementally, capped on output. - - ``brotli.decompress`` has no output cap and ``Decompressor.process`` - materializes everything it can from the given input, so feed the input - in bounded slices and enforce the cap on the running total. - """ - decompressor = brotli.Decompressor() - out = bytearray() - for i in range(0, len(data), 64 * 1024): - out.extend(decompressor.process(data[i : i + 64 * 1024])) - _enforce_decompression_cap(len(out), "br") - if not decompressor.is_finished(): - raise ValueError("Failed to decompress br request body: truncated or corrupt stream") - return bytes(out) - - async def _read_request_body_bytes(request: Request) -> bytes: """Read and (if needed) decompress the request body, returning raw UTF-8 bytes. @@ -2058,18 +2046,22 @@ async def _read_request_body_bytes(request: Request) -> bytes: except Exception as exc: raise ValueError(f"Failed to decompress deflate request body: {exc}") from exc elif encoding == "br": - try: - import brotli - except ImportError: - raise ValueError( - "Request body is brotli-compressed but the 'brotli' package is not installed." - ) from None - try: - raw = _decompress_brotli_capped(brotli, raw) - except RequestBodyTooLarge: - raise - except Exception as exc: - raise ValueError(f"Failed to decompress brotli request body: {exc}") from exc + # Brotli request bodies are rejected outright. The Python brotli + # bindings (brotli 1.2.0 and brotlicffi) expose no output-bounded + # streaming API: ``Decompressor.process(input_slice)`` returns ALL + # output produced from that slice before returning, so a highly + # compressible stream smaller than any feed slice can expand well + # past ``MAX_DECOMPRESSED_BODY_BYTES`` inside a single call — the + # exact allocation the cap exists to prevent. Rejecting the + # encoding (rather than decompressing) is the only way the + # process-wide decompression-bomb boundary holds for br; clients + # that send compressed request bodies can use gzip, deflate, or + # zstd, which all have bounded streaming APIs. + raise ValueError( + "Request body is brotli-compressed, but brotli request bodies are not " + "accepted: the Python brotli API cannot bound decompressed output " + "(decompression-bomb risk). Use Content-Encoding: gzip, deflate, or zstd." + ) elif encoding and encoding != "identity": raise ValueError(f"Unsupported Content-Encoding: {encoding}") diff --git a/tests/test_proxy_decompression_cap_status.py b/tests/test_proxy_decompression_cap_status.py new file mode 100644 index 000000000..e5f4afa83 --- /dev/null +++ b/tests/test_proxy_decompression_cap_status.py @@ -0,0 +1,263 @@ +"""Decompression-cap rejections must surface as body-too-large, never 400. + +``RequestBodyTooLarge`` is a ``ValueError`` subclass, so a generic +``except (json.JSONDecodeError, ValueError)`` handler clause would translate +a decompression-bomb rejection into ``400 invalid_json``. Every handler +path that reads a request body must instead return the configured +body-too-large status (default 413, see ``get_body_too_large_status``), and +the Bedrock adapter must never fail open: its old ``except Exception`` +path forwarded the compressed bomb verbatim to the upstream — exactly the +allocation this cap exists to prevent. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from headroom.proxy.helpers import ( + RequestBodyTooLarge, + get_body_too_large_status, +) +from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin +from headroom.proxy.handlers.batch import BatchHandlerMixin +from headroom.proxy.handlers.bedrock import BedrockHandlerMixin +from headroom.proxy.handlers.gemini import GeminiHandlerMixin +from headroom.proxy.handlers.openai import OpenAIHandlerMixin + +_TOO_LARGE = RequestBodyTooLarge( + "gzip request body exceeds the 100-byte decompression limit" +) + + +class _FakeState: + auth_mode = None + + +class _FakeRequest: + """Minimal Starlette Request stand-in: headers, state, url, body().""" + + def __init__(self, *, path: str = "/v1/messages", headers: dict | None = None) -> None: + self.headers = headers or {} + self.state = _FakeState() + self.url = SimpleNamespace(path=path, query="") + self.query_params: dict[str, str] = {} + + async def body(self) -> bytes: + return b"{}" + + +async def _raise_too_large(request) -> None: # noqa: ANN001 + raise _TOO_LARGE + + +def _error_payload(response) -> dict: + return json.loads(response.body) + + +class _FakeHttpClient: + def __init__(self) -> None: + self.posts: list[dict] = [] + self.post_response = SimpleNamespace( + status_code=200, + content=b"{}", + headers={"content-type": "application/json"}, + ) + + async def post(self, url: str, **kwargs): # noqa: ANN003, ANN201 + self.posts.append({"url": url, **kwargs}) + return self.post_response + + +class _OpenAIHandler(OpenAIHandlerMixin): + async def _next_request_id(self) -> str: + return "req-1" + + +class _AnthropicHandler(AnthropicHandlerMixin): + async def _next_request_id(self) -> str: + return "req-1" + + +class _BedrockHandler(BedrockHandlerMixin): + def __init__(self) -> None: + self.forwarded: list[dict] = [] + + def _bedrock_upstream_base(self) -> str: + return "https://bedrock.example" + + async def _next_request_id(self) -> str: + return "req-1" + + async def _forward_bedrock(self, **kwargs): # noqa: ANN003, ANN201 + self.forwarded.append(kwargs) + from fastapi.responses import Response + + return Response(status_code=599, content=b"should-not-be-forwarded") + + +class _BatchHandler(BatchHandlerMixin): + OPENAI_API_URL = "https://openai.example" + GEMINI_API_URL = "https://gemini.example" + + def __init__(self) -> None: + self.http_client = _FakeHttpClient() + + async def _next_request_id(self) -> str: + return "req-1" + + +class _GeminiHandler(GeminiHandlerMixin): + async def _next_request_id(self) -> str: + return "req-1" + + +# --------------------------------------------------------------------------- +# OpenAI /v1/chat/completions and /v1/responses +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_openai_chat_bomb_returns_body_too_large_status(monkeypatch) -> None: + monkeypatch.setattr("headroom.proxy.helpers._read_request_json", _raise_too_large) + response = await _OpenAIHandler().handle_openai_chat( + _FakeRequest(path="/v1/chat/completions") + ) + assert response.status_code == get_body_too_large_status() + error = _error_payload(response)["error"] + assert error["code"] == "request_too_large" + assert "decompression limit" in error["message"] + + +@pytest.mark.asyncio +async def test_openai_responses_bomb_returns_body_too_large_status(monkeypatch) -> None: + monkeypatch.setattr( + "headroom.proxy.helpers.read_request_json_with_bytes", _raise_too_large + ) + response = await _OpenAIHandler().handle_openai_responses( + _FakeRequest(path="/v1/responses") + ) + assert response.status_code == get_body_too_large_status() + assert _error_payload(response)["error"]["code"] == "request_too_large" + + +@pytest.mark.asyncio +async def test_openai_chat_bomb_status_is_configurable(monkeypatch) -> None: + monkeypatch.setenv("HEADROOM_PROXY_BODY_TOO_LARGE_STATUS", "429") + monkeypatch.setattr("headroom.proxy.helpers._read_request_json", _raise_too_large) + response = await _OpenAIHandler().handle_openai_chat(_FakeRequest()) + assert response.status_code == 429 + + +# --------------------------------------------------------------------------- +# Anthropic /v1/messages +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_anthropic_messages_bomb_returns_body_too_large_status(monkeypatch) -> None: + async def _noop(*args, **kwargs) -> None: # noqa: ANN002, ANN003 + return None + + monkeypatch.setattr( + "headroom.proxy.helpers.read_request_json_with_bytes", _raise_too_large + ) + monkeypatch.setattr( + "headroom.proxy.handlers.anthropic.emit_stage_timings_log", _noop + ) + response = await _AnthropicHandler().handle_anthropic_messages( + _FakeRequest(path="/v1/messages") + ) + assert response.status_code == get_body_too_large_status() + assert _error_payload(response)["error"]["type"] == "request_too_large" + + +# --------------------------------------------------------------------------- +# Bedrock: must never fail open and forward the expanding body +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_bedrock_invoke_bomb_never_fails_open(monkeypatch) -> None: + monkeypatch.setattr( + "headroom.proxy.helpers.read_request_json_with_bytes", _raise_too_large + ) + handler = _BedrockHandler() + response = await handler.handle_bedrock_invoke( + _FakeRequest(path="/model/x/invoke"), "amazon.nova-pro-v1:0", stream=False + ) + assert response.status_code == get_body_too_large_status() + assert handler.forwarded == [] # the compressed bomb is never forwarded + + +# --------------------------------------------------------------------------- +# Batch paths (OpenAI create + Google create) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_openai_batch_create_bomb_returns_body_too_large_status(monkeypatch) -> None: + monkeypatch.setattr("headroom.proxy.helpers._read_request_json", _raise_too_large) + response = await _BatchHandler().handle_batch_create( + _FakeRequest(path="/v1/batches") + ) + assert response.status_code == get_body_too_large_status() + assert _error_payload(response)["error"]["code"] == "request_too_large" + + +@pytest.mark.asyncio +async def test_google_batch_create_bomb_returns_body_too_large_status(monkeypatch) -> None: + monkeypatch.setattr("headroom.proxy.helpers._read_request_json", _raise_too_large) + response = await _BatchHandler().handle_google_batch_create( + _FakeRequest(path="/v1beta/models/gemini-pro:batchGenerateContent"), + "gemini-pro", + ) + assert response.status_code == get_body_too_large_status() + assert _error_payload(response)["error"]["code"] == get_body_too_large_status() + + +# --------------------------------------------------------------------------- +# Gemini +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_gemini_generate_content_bomb_returns_body_too_large_status( + monkeypatch, +) -> None: + monkeypatch.setattr("headroom.proxy.helpers._read_request_json", _raise_too_large) + response = await _GeminiHandler().handle_gemini_generate_content( + _FakeRequest(path="/v1beta/models/gemini-pro:generateContent"), "gemini-pro" + ) + assert response.status_code == get_body_too_large_status() + assert _error_payload(response)["error"]["code"] == get_body_too_large_status() + + +# --------------------------------------------------------------------------- +# Batch passthrough best-effort read: size-policy rejections must propagate +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_batch_passthrough_re_raises_body_too_large(monkeypatch) -> None: + monkeypatch.setattr( + "headroom.proxy.helpers._read_request_body_bytes", _raise_too_large + ) + handler = _BatchHandler() + with pytest.raises(RequestBodyTooLarge): + await handler._batch_passthrough(_FakeRequest(), {"model": "x"}) + assert handler.http_client.posts == [] # nothing forwarded + + +@pytest.mark.asyncio +async def test_batch_passthrough_other_read_errors_still_best_effort(monkeypatch) -> None: + async def _boom(request) -> None: # noqa: ANN001 + raise RuntimeError("body already consumed upstream") + + monkeypatch.setattr("headroom.proxy.helpers._read_request_body_bytes", _boom) + handler = _BatchHandler() + response = await handler._batch_passthrough(_FakeRequest(), {"model": "x"}) + assert response.status_code == 200 + assert handler.http_client.posts # forwarded via canonical re-serialization diff --git a/tests/test_proxy_request_decompression_cap.py b/tests/test_proxy_request_decompression_cap.py index 2be70f6af..178dcb897 100644 --- a/tests/test_proxy_request_decompression_cap.py +++ b/tests/test_proxy_request_decompression_cap.py @@ -1,13 +1,19 @@ """Decompression of compressed request bodies must be bounded. ``_read_request_body_bytes`` (headroom/proxy/helpers.py) expands -zstd/gzip/deflate/br request bodies before forwarding. The expansion used +zstd/gzip/deflate request bodies before forwarding. The expansion used to be unbounded (``gzip.decompress`` / ``zlib.decompress`` / ``brotli.decompress`` / ``stream_reader().read()``), so a tiny compressed body could balloon into an unbounded in-memory buffer — a decompression-bomb -DoS. Every format is now fed incrementally against +DoS. Every supported format is now fed incrementally against ``MAX_DECOMPRESSED_BODY_BYTES`` and raises ``RequestBodyTooLarge`` (a ``ValueError``) when the cap is exceeded. + +Brotli (``br``) request bodies are rejected outright: the Python brotli +bindings expose no output-bounded streaming API, so a highly compressible +stream smaller than any feed slice can expand past the cap inside a single +``Decompressor.process()`` call — the exact allocation the cap exists to +prevent. Rejecting the encoding keeps the process-wide boundary intact. """ import asyncio @@ -69,10 +75,14 @@ def test_zstd_round_trip(): assert _read(raw, "zstd") == _PAYLOAD -def test_brotli_round_trip(): +def test_brotli_round_trip_rejected(): + # br is rejected outright: the Python brotli bindings have no + # output-bounded streaming API, so even a valid small stream could + # expand past the cap inside a single Decompressor.process() call. brotli = pytest.importorskip("brotli") raw = brotli.compress(_PAYLOAD) - assert _read(raw, "br") == _PAYLOAD + with pytest.raises(ValueError, match="brotli"): + _read(raw, "br") # --------------------------------------------------------------------------- @@ -110,10 +120,24 @@ def test_zstd_bomb_rejected(monkeypatch): ) -def test_brotli_bomb_rejected(monkeypatch): +def test_brotli_peak_input_never_reaches_decompressor(monkeypatch): + """Regression: peak-producing br input cannot cross the cap boundary. + + Before this PR the brotli path fed 64 KiB input slices to + ``Decompressor.process()``, which returns ALL output produced from a + slice before the cap check runs — a highly compressible sub-64 KiB + stream could materialize far beyond the cap in one call. Brotli is now + rejected before any decompression happens, so the boundary holds by + construction: the bomb is refused as an unsupported encoding + (ValueError), never decompressed, and never surfaces as + RequestBodyTooLarge (which would imply a decompressor ran). + """ brotli = pytest.importorskip("brotli") - with pytest.raises(RequestBodyTooLarge): - _read(_bomb_under_cap(monkeypatch, brotli.compress), "br") + monkeypatch.setattr("headroom.proxy.helpers.MAX_DECOMPRESSED_BODY_BYTES", 4096) + bomb = brotli.compress(b"A" * (1024 * 1024)) + assert len(bomb) < 4096 # sanity: the input itself is tiny + with pytest.raises(ValueError, match="brotli"): + _read(bomb, "br") def test_payload_at_cap_is_accepted(monkeypatch): @@ -150,9 +174,13 @@ def test_zstd_not_installed_raises(monkeypatch): _read(b"ignored", "zstd") -def test_brotli_not_installed_raises(monkeypatch): +def test_brotli_rejected_regardless_of_install(monkeypatch): + # The br rejection is unconditional: the encoding is refused before any + # brotli import, so whether the package is installed or not changes + # nothing (and a missing package can no longer silently allow a bomb + # through as "not installed"). monkeypatch.setitem(sys.modules, "brotli", None) - with pytest.raises(ValueError, match="not installed"): + with pytest.raises(ValueError, match="brotli"): _read(b"ignored", "br") @@ -193,9 +221,12 @@ def test_gzip_payload_across_slice_boundary_round_trips(): def test_truncated_brotli_still_raises_value_error(): + # br is rejected before any decompression, so even a truncated stream is + # refused with the same encoding rejection (still a client-error + # ValueError, matching the pre-PR contract). brotli = pytest.importorskip("brotli") truncated = brotli.compress(_PAYLOAD)[:16] - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="brotli"): _read(truncated, "br") From f5afe84626917a2e8154b5443f9d9c64617742d1 Mon Sep 17 00:00:00 2001 From: Chukwuebuka-2003 Date: Wed, 5 Aug 2026 02:27:26 +0000 Subject: [PATCH 6/7] fix(proxy): preserve body-too-large status in /v1/compress handler handle_compress (both bypass and non-bypass branches) now catches RequestBodyTooLarge explicitly and returns the configured body-too-large status (413 default) instead of a generic 400. Add integration tests for /v1/compress with gzip bombs covering both normal and bypass paths. Signed-off-by: Chukwuebuka-2003 --- headroom/proxy/handlers/openai.py | 10 +++ tests/test_proxy_request_decompression_cap.py | 83 +++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 0ef53ba47..ed65c8859 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -8420,6 +8420,16 @@ class OpenAIHandlerMixin: try: body = await _read_request_json(request) + except RequestBodyTooLarge: + return JSONResponse( + status_code=get_body_too_large_status(), + content={ + "error": { + "type": "invalid_request", + "message": "Request body too large.", + } + }, + ) except Exception: return JSONResponse( status_code=400, diff --git a/tests/test_proxy_request_decompression_cap.py b/tests/test_proxy_request_decompression_cap.py index 178dcb897..01d54388d 100644 --- a/tests/test_proxy_request_decompression_cap.py +++ b/tests/test_proxy_request_decompression_cap.py @@ -262,3 +262,86 @@ def test_cap_is_sane_default(): # The default cap mirrors the uncompressed request budget: a compressed # body must never expand beyond what the handler would accept anyway. assert MAX_DECOMPRESSED_BODY_BYTES > 0 + + +# --------------------------------------------------------------------------- +# Integration: /v1/compress preserves body-too-large status (413 / configured) +# --------------------------------------------------------------------------- + +try: + from fastapi.testclient import TestClient # noqa: F811 + + _HAS_FASTAPI = True +except ImportError: + _HAS_FASTAPI = False + + +@pytest.mark.skipif(not _HAS_FASTAPI, reason="fastapi not installed") +def test_compress_endpoint_body_too_large_gzip(monkeypatch): + """A gzip bomb sent to /v1/compress returns the configured status, not 400.""" + from headroom.proxy.server import ProxyConfig, create_app + from headroom.proxy.helpers import get_body_too_large_status + + config = ProxyConfig( + optimize=True, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + + # Build a gzip payload whose decompressed size exceeds the cap. + # The bomb is just a highly compressible string (all 'A') so the + # compressed wire size stays tiny. + bomb_body = json.dumps({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "A" * (MAX_DECOMPRESSED_BODY_BYTES + 1)}], + }).encode("utf-8") + compressed = gzip.compress(bomb_body) + expected_status = get_body_too_large_status() + + with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client: + resp = client.post( + "/v1/compress", + content=compressed, + headers={"Content-Encoding": "gzip", "Content-Type": "application/json"}, + ) + assert resp.status_code == expected_status, ( + f"expected {expected_status}, got {resp.status_code}: {resp.text}" + ) + + +@pytest.mark.skipif(not _HAS_FASTAPI, reason="fastapi not installed") +def test_compress_bypass_body_too_large_gzip(monkeypatch): + """A gzip bomb sent to /v1/compress with x-headroom-bypass also returns the configured status.""" + from headroom.proxy.server import ProxyConfig, create_app + from headroom.proxy.helpers import get_body_too_large_status + + config = ProxyConfig( + optimize=True, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + + bomb_body = json.dumps({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "A" * (MAX_DECOMPRESSED_BODY_BYTES + 1)}], + }).encode("utf-8") + compressed = gzip.compress(bomb_body) + expected_status = get_body_too_large_status() + + with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client: + resp = client.post( + "/v1/compress", + content=compressed, + headers={ + "Content-Encoding": "gzip", + "Content-Type": "application/json", + "x-headroom-bypass": "true", + }, + ) + assert resp.status_code == expected_status, ( + f"expected {expected_status}, got {resp.status_code}: {resp.text}" + ) From df78777b5616204dceeffdb2cea9977952a3e2f5 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Tue, 4 Aug 2026 21:32:19 -0500 Subject: [PATCH 7/7] style(tests): order decompression cap imports --- tests/test_proxy_request_decompression_cap.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_proxy_request_decompression_cap.py b/tests/test_proxy_request_decompression_cap.py index 01d54388d..6e2663f36 100644 --- a/tests/test_proxy_request_decompression_cap.py +++ b/tests/test_proxy_request_decompression_cap.py @@ -279,8 +279,8 @@ except ImportError: @pytest.mark.skipif(not _HAS_FASTAPI, reason="fastapi not installed") def test_compress_endpoint_body_too_large_gzip(monkeypatch): """A gzip bomb sent to /v1/compress returns the configured status, not 400.""" - from headroom.proxy.server import ProxyConfig, create_app from headroom.proxy.helpers import get_body_too_large_status + from headroom.proxy.server import ProxyConfig, create_app config = ProxyConfig( optimize=True, @@ -314,8 +314,8 @@ def test_compress_endpoint_body_too_large_gzip(monkeypatch): @pytest.mark.skipif(not _HAS_FASTAPI, reason="fastapi not installed") def test_compress_bypass_body_too_large_gzip(monkeypatch): """A gzip bomb sent to /v1/compress with x-headroom-bypass also returns the configured status.""" - from headroom.proxy.server import ProxyConfig, create_app from headroom.proxy.helpers import get_body_too_large_status + from headroom.proxy.server import ProxyConfig, create_app config = ProxyConfig( optimize=True,