mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Merge df78777b56 into 8884d87378
This commit is contained in:
Chukwuebuka (Gaus Octavio )
GitHub
commit
67c6decd34
8 changed files with 955 additions and 27 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
||||
|
|
|
|||
263
tests/test_proxy_decompression_cap_status.py
Normal file
263
tests/test_proxy_decompression_cap_status.py
Normal file
|
|
@ -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
|
||||
347
tests/test_proxy_request_decompression_cap.py
Normal file
347
tests/test_proxy_request_decompression_cap.py
Normal file
|
|
@ -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}"
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue