Cap ASGI middleware request body buffering

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."""

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