Limit gateway POST body buffering

nemo_retriever/src/nemo_retriever/service/app.py

@@ -194,36 +194,52 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
 
 
 class _GatewayBodyCacheMiddleware:
-    """Pure ASGI middleware: buffer POST bodies so the proxy can forward them.
+    """Pure ASGI middleware: cap and buffer POST bodies for proxy forwarding.
 
     FastAPI's dependency injection parses ``UploadFile`` / ``Form`` parameters
     by consuming the ASGI body stream *before* the route handler runs.  When
     the gateway's proxy later calls ``request.body()`` the stream is already
     exhausted and Starlette raises ``RuntimeError: Stream consumed``.
 
-    This middleware reads the entire body once, stores it on the ASGI scope
-    as ``scope["_cached_body"]``, and replays it through a synthetic
-    ``receive`` callable so that form parsing works normally.  The proxy
-    then reads from ``request.scope["_cached_body"]`` directly.
+    This middleware rejects bodies above ``max_body_bytes`` while reading,
+    stores accepted bodies on the ASGI scope as ``scope["_cached_body"]``,
+    and replays them through a synthetic ``receive`` callable so that form
+    parsing works normally.  The proxy then reads from
+    ``request.scope["_cached_body"]`` directly.
 
     Only active for ``POST`` requests; ``GET`` / ``OPTIONS`` etc. pass through
     untouched.
     """
 
-    def __init__(self, app: Any) -> None:
+    def __init__(self, app: Any, max_body_bytes: int | None = None) -> None:
         self.app = app
+        self.max_body_bytes = max_body_bytes
 
     async def __call__(self, scope: dict, receive: Any, send: Any) -> None:
         if scope["type"] != "http" or scope.get("method", "GET") != "POST":
             await self.app(scope, receive, send)
             return
 
+        limit = self.max_body_bytes
+        declared_size = self._content_length(scope)
+        if limit is not None and declared_size is not None and declared_size > limit:
+            await self._payload_too_large(scope, receive, send, limit, declared_size)
+            return
+
         body_parts: list[bytes] = []
+        body_size = 0
         while True:
             message = await receive()
-            body_parts.append(message.get("body", b""))
+            chunk = message.get("body", b"")
+            if chunk:
+                body_size += len(chunk)
+                if limit is not None and body_size > limit:
+                    await self._payload_too_large(scope, receive, send, limit, body_size)
+                    return
+                body_parts.append(chunk)
             if not message.get("more_body", False):
                 break
+
         body = b"".join(body_parts)
         scope["_cached_body"] = body
 
@@ -238,6 +254,33 @@ async def replay_receive() -> dict:
 
         await self.app(scope, replay_receive, send)
 
+    @staticmethod
+    def _content_length(scope: dict) -> int | None:
+        for name, value in scope.get("headers", []):
+            if name.lower() == b"content-length":
+                try:
+                    parsed = int(value)
+                except ValueError:
+                    return None
+                return parsed if parsed >= 0 else None
+        return None
+
+    @staticmethod
+    async def _payload_too_large(
+        scope: dict,
+        receive: Any,
+        send: Any,
+        limit: int,
+        actual_size: int,
+    ) -> None:
+        response = JSONResponse(
+            status_code=413,
+            content={
+                "detail": f"Request body size {actual_size:,} bytes exceeds limit of {limit:,} bytes",
+            },
+        )
+        await response(scope, receive, send)
+
 
 def create_app(config: ServiceConfig) -> FastAPI:
     """Build and return a fully-configured :class:`FastAPI` application."""
@@ -256,7 +299,7 @@ def create_app(config: ServiceConfig) -> FastAPI:
     app.add_middleware(_RequestIdMiddleware)
 
     if config.mode == "gateway":
-        app.add_middleware(_GatewayBodyCacheMiddleware)
+        app.add_middleware(_GatewayBodyCacheMiddleware, max_body_bytes=config.resources.max_upload_bytes)
         logger.info("Gateway body-cache middleware ENABLED")
 
     if config.auth.api_token:

nemo_retriever/tests/test_gateway_body_cache.py

@@ -0,0 +1,252 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from starlette.requests import Request
+
+from nemo_retriever.service.app import _GatewayBodyCacheMiddleware
+
+
+def _scope(headers: list[tuple[bytes, bytes]]) -> dict[str, Any]:
+    return {
+        "type": "http",
+        "asgi": {"version": "3.0"},
+        "http_version": "1.1",
+        "method": "POST",
+        "scheme": "http",
+        "path": "/v1/ingest/job/job-id/document",
+        "raw_path": b"/v1/ingest/job/job-id/document",
+        "query_string": b"",
+        "headers": headers,
+        "client": ("127.0.0.1", 12345),
+        "server": ("testserver", 80),
+    }
+
+
+async def _call_asgi(
+    app: Any,
+    *,
+    headers: list[tuple[bytes, bytes]],
+    messages: list[dict[str, Any]],
+) -> tuple[list[dict[str, Any]], int]:
+    sent: list[dict[str, Any]] = []
+    receive_calls = 0
+    pending = list(messages)
+
+    async def receive() -> dict[str, Any]:
+        nonlocal receive_calls
+        receive_calls += 1
+        if pending:
+            return pending.pop(0)
+        return {"type": "http.disconnect"}
+
+    async def send(message: dict[str, Any]) -> None:
+        sent.append(message)
+
+    await app(_scope(headers), receive, send)
+    return sent, receive_calls
+
+
+def _status(sent: list[dict[str, Any]]) -> int:
+    return next(message["status"] for message in sent if message["type"] == "http.response.start")
+
+
+def _body(sent: list[dict[str, Any]]) -> bytes:
+    return b"".join(message.get("body", b"") for message in sent if message["type"] == "http.response.body")
+
+
+class _NeverCalledApp:
+    called = False
+
+    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
+        self.called = True
+        await send({"type": "http.response.start", "status": 204, "headers": []})
+        await send({"type": "http.response.body", "body": b""})
+
+
+class _BodyReadingApp:
+    called = False
+    cached_body: bytes | None = None
+    body: bytes | None = None
+
+    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
+        self.called = True
+        self.cached_body = scope.get("_cached_body")
+        self.body = await Request(scope, receive).body()
+
+        await send({"type": "http.response.start", "status": 202, "headers": []})
+        await send({"type": "http.response.body", "body": b"accepted"})
+
+
+class _MultipartParsingApp:
+    called = False
+    cached_body: bytes | None = None
+    filename: str | None = None
+    metadata: str | None = None
+    file_content: bytes | None = None
+
+    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
+        self.called = True
+        self.cached_body = scope.get("_cached_body")
+
+        form = await Request(scope, receive).form()
+        upload = form["file"]
+        self.filename = upload.filename
+        self.metadata = str(form["metadata"])
+        self.file_content = await upload.read()
+
+        await send({"type": "http.response.start", "status": 202, "headers": []})
+        await send({"type": "http.response.body", "body": b"accepted"})
+
+
+def test_gateway_body_cache_no_limit_passes_large_body_through() -> None:
+    downstream = _BodyReadingApp()
+    app = _GatewayBodyCacheMiddleware(downstream, max_body_bytes=None)
+    payload = b"x" * 1_000_000
+
+    sent, receive_calls = asyncio.run(
+        _call_asgi(
+            app,
+            headers=[(b"content-length", str(len(payload)).encode())],
+            messages=[{"type": "http.request", "body": payload, "more_body": False}],
+        )
+    )
+
+    assert _status(sent) == 202
+    assert _body(sent) == b"accepted"
+    assert receive_calls == 1
+    assert downstream.called is True
+    assert downstream.cached_body == payload
+    assert downstream.body == payload
+
+
+def test_gateway_body_cache_rejects_oversized_content_length_without_reading_body() -> None:
+    downstream = _NeverCalledApp()
+    app = _GatewayBodyCacheMiddleware(downstream, max_body_bytes=4)
+
+    sent, receive_calls = asyncio.run(
+        _call_asgi(
+            app,
+            headers=[(b"content-length", b"5")],
+            messages=[{"type": "http.request", "body": b"12345", "more_body": False}],
+        )
+    )
+
+    assert _status(sent) == 413
+    assert b"exceeds limit" in _body(sent)
+    assert receive_calls == 0
+    assert downstream.called is False
+
+
+def test_gateway_body_cache_accepts_content_length_at_limit() -> None:
+    downstream = _BodyReadingApp()
+    payload = b"12345"
+    app = _GatewayBodyCacheMiddleware(downstream, max_body_bytes=len(payload))
+
+    sent, receive_calls = asyncio.run(
+        _call_asgi(
+            app,
+            headers=[(b"content-length", str(len(payload)).encode())],
+            messages=[{"type": "http.request", "body": payload, "more_body": False}],
+        )
+    )
+
+    assert _status(sent) == 202
+    assert _body(sent) == b"accepted"
+    assert receive_calls == 1
+    assert downstream.called is True
+    assert downstream.cached_body == payload
+    assert downstream.body == payload
+
+
+def test_gateway_body_cache_rejects_chunked_upload_after_limit_is_exceeded() -> None:
+    downstream = _NeverCalledApp()
+    app = _GatewayBodyCacheMiddleware(downstream, max_body_bytes=5)
+
+    sent, receive_calls = asyncio.run(
+        _call_asgi(
+            app,
+            headers=[(b"transfer-encoding", b"chunked")],
+            messages=[
+                {"type": "http.request", "body": b"abc", "more_body": True},
+                {"type": "http.request", "body": b"def", "more_body": True},
+                {"type": "http.request", "body": b"ghi", "more_body": False},
+            ],
+        )
+    )
+
+    assert _status(sent) == 413
+    assert b"6 bytes exceeds limit of 5 bytes" in _body(sent)
+    assert receive_calls == 2
+    assert downstream.called is False
+
+
+def test_gateway_body_cache_accepts_chunked_upload_at_limit() -> None:
+    downstream = _BodyReadingApp()
+    payload = b"abcdef"
+    app = _GatewayBodyCacheMiddleware(downstream, max_body_bytes=len(payload))
+
+    sent, receive_calls = asyncio.run(
+        _call_asgi(
+            app,
+            headers=[],
+            messages=[
+                {"type": "http.request", "body": payload[:3], "more_body": True},
+                {"type": "http.request", "body": payload[3:], "more_body": False},
+            ],
+        )
+    )
+
+    assert _status(sent) == 202
+    assert _body(sent) == b"accepted"
+    assert receive_calls == 2
+    assert downstream.called is True
+    assert downstream.cached_body == payload
+    assert downstream.body == payload
+
+
+def test_gateway_body_cache_replays_valid_multipart_upload_under_limit() -> None:
+    downstream = _MultipartParsingApp()
+    app = _GatewayBodyCacheMiddleware(downstream, max_body_bytes=1024)
+
+    boundary = b"----nemo-retriever-test-boundary"
+    body = (
+        b"--"
+        + boundary
+        + b'\r\nContent-Disposition: form-data; name="metadata"\r\n\r\n{}\r\n'
+        + b"--"
+        + boundary
+        + b'\r\nContent-Disposition: form-data; name="file"; filename="doc.txt"\r\n'
+        + b"Content-Type: text/plain\r\n\r\nhello gateway\r\n"
+        + b"--"
+        + boundary
+        + b"--\r\n"
+    )
+
+    sent, receive_calls = asyncio.run(
+        _call_asgi(
+            app,
+            headers=[
+                (b"content-type", b"multipart/form-data; boundary=" + boundary),
+                (b"content-length", str(len(body)).encode()),
+            ],
+            messages=[
+                {"type": "http.request", "body": body[:40], "more_body": True},
+                {"type": "http.request", "body": body[40:], "more_body": False},
+            ],
+        )
+    )
+
+    assert _status(sent) == 202
+    assert _body(sent) == b"accepted"
+    assert receive_calls == 2
+    assert downstream.called is True
+    assert downstream.cached_body == body
+    assert downstream.filename == "doc.txt"
+    assert downstream.metadata == "{}"
+    assert downstream.file_content == b"hello gateway"