diff --git a/crates/tare-py/python/tare/integrations.py b/crates/tare-py/python/tare/integrations.py index 73c66f3..2062d04 100644 --- a/crates/tare-py/python/tare/integrations.py +++ b/crates/tare-py/python/tare/integrations.py @@ -174,6 +174,23 @@ async def async_pre_call_hook( # ASGI middleware # --------------------------------------------------------------------------- +DEFAULT_MAX_BODY_BYTES = 32 * 1024 * 1024 + + +class _RequestBodyTooLarge(BaseException): + """Signal that an ASGI request exceeded the configured body limit. + + Subclasses ``BaseException`` (not ``Exception``) so that an inner + application's catch-all exception middleware cannot intercept it. If it + were an ordinary ``Exception``, wrapping an app that has its own error + middleware (e.g. ``CompressionMiddleware(existing_fastapi_app)``) would + let that middleware emit a 500 and mark the response started before the + signal reaches our handler, so oversized chunked requests would get a 500 + instead of the documented 413. Keeping it out of the ``Exception`` + hierarchy lets the signal propagate cleanly to the 413 handler below. + """ + + class CompressionMiddleware: """ASGI middleware that compresses JSON request bodies containing 'messages'. @@ -195,36 +212,93 @@ class CompressionMiddleware: The inner ASGI application. task: Optional task hint forwarded to :func:`compress_messages`. + max_body_bytes: + Maximum request body size to buffer. Larger requests receive a 413 + response. Defaults to 32 MiB, matching the tare proxy. """ - def __init__(self, app: Any, task: str = "") -> None: + def __init__( + self, + app: Any, + task: str = "", + max_body_bytes: int = DEFAULT_MAX_BODY_BYTES, + ) -> None: + if max_body_bytes < 1: + raise ValueError("max_body_bytes must be positive") self.app = app self.task = task + self.max_body_bytes = max_body_bytes async def __call__(self, scope: dict, receive: Any, send: Any) -> None: if scope.get("type") != "http": await self.app(scope, receive, send) return - consumed: list[bytes] = [] + content_length = next( + ( + value + for name, value in scope.get("headers", []) + if name.lower() == b"content-length" + ), + None, + ) + if content_length is not None: + try: + body_is_too_large = int(content_length) > self.max_body_bytes + except ValueError: + body_is_too_large = False + if body_is_too_large: + await self._send_payload_too_large(send) + return + + body: bytes | None = None + response_started = False async def patched_receive() -> dict: - if consumed: - return { - "type": "http.request", - "body": b"".join(consumed), - "more_body": False, - } + nonlocal body + if body is not None: + return await receive() + + body_buffer = bytearray() event = await receive() - body = event.get("body", b"") - while event.get("more_body", False): + while True: + chunk = event.get("body", b"") + if len(body_buffer) + len(chunk) > self.max_body_bytes: + raise _RequestBodyTooLarge + body_buffer.extend(chunk) + if not event.get("more_body", False): + break event = await receive() - body += event.get("body", b"") - body = self._maybe_compress_body(body) - consumed.append(body) + body = self._maybe_compress_body(bytes(body_buffer)) return {"type": "http.request", "body": body, "more_body": False} - await self.app(scope, patched_receive, send) + async def patched_send(event: dict) -> None: + nonlocal response_started + if event.get("type") == "http.response.start": + response_started = True + await send(event) + + try: + await self.app(scope, patched_receive, patched_send) + except _RequestBodyTooLarge: + if response_started: + raise + await self._send_payload_too_large(send) + + @staticmethod + async def _send_payload_too_large(send: Any) -> None: + body = b"Request body too large" + await send( + { + "type": "http.response.start", + "status": 413, + "headers": [ + (b"content-type", b"text/plain; charset=utf-8"), + (b"content-length", str(len(body)).encode()), + ], + } + ) + await send({"type": "http.response.body", "body": body}) def _maybe_compress_body(self, body: bytes) -> bytes: """Return compressed body if it contains a 'messages' key; else passthrough.""" diff --git a/crates/tare-py/tests/test_integrations.py b/crates/tare-py/tests/test_integrations.py index 73e584a..2f3a110 100644 --- a/crates/tare-py/tests/test_integrations.py +++ b/crates/tare-py/tests/test_integrations.py @@ -12,6 +12,7 @@ from __future__ import annotations +import asyncio import json import sys import types @@ -146,6 +147,92 @@ async def dummy_app(scope: dict, receive: object, send: object) -> None: assert mw.task == "test task" +def test_compression_middleware_rejects_large_chunked_body() -> None: + from tare.integrations import CompressionMiddleware + + chunks = [b"1234", b"5678"] + received_events = iter( + [ + {"type": "http.request", "body": chunks[0], "more_body": True}, + {"type": "http.request", "body": chunks[1], "more_body": False}, + ] + ) + sent_events: list[dict] = [] + + async def receive() -> dict: + return next(received_events) + + async def send(event: dict) -> None: + sent_events.append(event) + + async def app(scope: dict, app_receive: object, app_send: object) -> None: + await app_receive() # type: ignore[operator] + + middleware = CompressionMiddleware(app, max_body_bytes=7) + asyncio.run(middleware({"type": "http"}, receive, send)) + + assert sent_events[0]["status"] == 413 + assert sent_events[1]["body"] == b"Request body too large" + + +def test_compression_middleware_does_not_replace_started_response() -> None: + import pytest + + from tare.integrations import CompressionMiddleware, _RequestBodyTooLarge + + received_events = iter( + [ + {"type": "http.request", "body": b"1234", "more_body": True}, + {"type": "http.request", "body": b"5678", "more_body": False}, + ] + ) + sent_events: list[dict] = [] + + async def receive() -> dict: + return next(received_events) + + async def send(event: dict) -> None: + sent_events.append(event) + + async def app(scope: dict, app_receive: object, app_send: object) -> None: + await app_send({"type": "http.response.start", "status": 200}) # type: ignore[operator] + await app_receive() # type: ignore[operator] + + middleware = CompressionMiddleware(app, max_body_bytes=7) + with pytest.raises(_RequestBodyTooLarge): + asyncio.run(middleware({"type": "http"}, receive, send)) + + assert sent_events == [{"type": "http.response.start", "status": 200}] + + +def test_compression_middleware_rejects_large_content_length_early() -> None: + from tare.integrations import CompressionMiddleware + + receive_called = False + app_called = False + sent_events: list[dict] = [] + + async def receive() -> dict: + nonlocal receive_called + receive_called = True + return {"type": "http.request", "body": b""} + + async def send(event: dict) -> None: + sent_events.append(event) + + async def app(scope: dict, app_receive: object, app_send: object) -> None: + nonlocal app_called + app_called = True + + middleware = CompressionMiddleware(app, max_body_bytes=7) + scope = {"type": "http", "headers": [(b"content-length", b"8")]} + asyncio.run(middleware(scope, receive, send)) + + assert sent_events[0]["status"] == 413 + assert not receive_called + assert not app_called + + def test_compression_middleware_compresses_json_body() -> None: """_maybe_compress_body reduces total message content for redundant input.""" from tare.integrations import CompressionMiddleware