diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 5f5e4a510..4ff50b521 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -846,8 +846,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 @@ -996,6 +998,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( @@ -4876,7 +4890,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 @@ -4900,6 +4919,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 91c23f640..1b4129956 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, @@ -1218,6 +1249,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, @@ -1241,8 +1273,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 79463c579..7cf7e50e2 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -277,7 +277,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() @@ -299,6 +304,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, @@ -1006,7 +1021,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() @@ -1014,6 +1033,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, @@ -1175,7 +1204,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() @@ -1183,6 +1216,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, @@ -1268,7 +1311,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() @@ -1277,6 +1324,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 42ac1a8c0..b7f026b68 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -3251,7 +3251,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 @@ -3284,6 +3286,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, @@ -5349,6 +5362,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 @@ -5385,6 +5400,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, @@ -9557,12 +9583,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, @@ -9583,6 +9624,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/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 8d2a2f43c..1a2efa9cc 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1065,6 +1065,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 @@ -2537,12 +2545,114 @@ 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 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. + """ + + +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, out: bytearray | None = None +) -> 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. Pass *out* to accumulate across calls (used + for multi-member gzip) so the cap applies to the whole payload. + """ + out = out if out is not None else 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 + # 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) + + +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 ({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. + + ``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() + + 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() @@ -2550,43 +2660,56 @@ 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. + raw = _decompress_gzip_capped(raw) + 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 - 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 new file mode 100644 index 000000000..6e2663f36 --- /dev/null +++ b/tests/test_proxy_request_decompression_cap.py @@ -0,0 +1,347 @@ +"""Decompression of compressed request bodies must be bounded. + +``_read_request_body_bytes`` (headroom/proxy/helpers.py) expands +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 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 +import gzip +import json +import sys +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_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) + with pytest.raises(ValueError, match="brotli"): + _read(raw, "br") + + +# --------------------------------------------------------------------------- +# 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_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") + 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): + 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_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_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="brotli"): + _read(b"ignored", "br") + + +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(): + 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(): + # 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, match="brotli"): + _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 + + +# --------------------------------------------------------------------------- +# 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.helpers import get_body_too_large_status + from headroom.proxy.server import ProxyConfig, create_app + + 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.helpers import get_body_too_large_status + from headroom.proxy.server import ProxyConfig, create_app + + 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}" + )