diff --git a/nemo_retriever/src/nemo_retriever/service/app.py b/nemo_retriever/src/nemo_retriever/service/app.py index fec9ffd789..93a2c80098 100644 --- a/nemo_retriever/src/nemo_retriever/service/app.py +++ b/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: diff --git a/nemo_retriever/tests/test_gateway_body_cache.py b/nemo_retriever/tests/test_gateway_body_cache.py new file mode 100644 index 0000000000..9af3c4e0e3 --- /dev/null +++ b/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"