Limit gateway POST body buffering#2251
Conversation
Greptile SummaryThis PR adds request-body size enforcement to
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/service/app.py | Adds max_body_bytes parameter and enforcement to _GatewayBodyCacheMiddleware via header fast-path and streaming byte-count check; wires config.resources.max_upload_bytes at middleware registration time. |
| nemo_retriever/tests/test_gateway_body_cache.py | New test file covering all six key scenarios: unlimited mode, header-based rejection, exact-limit boundaries for both header and streaming paths, chunked oversize detection, and multipart replay. |
Reviews (2): Last reviewed commit: "Limit gateway POST body buffering" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def _content_length(scope: dict) -> int | None: | ||
| for name, value in scope.get("headers", []): | ||
| if name.lower() == b"content-length": | ||
| try: | ||
| return int(value) | ||
| except ValueError: | ||
| return None | ||
| return None |
There was a problem hiding this comment.
_content_length accepts non-positive values, silently bypassing the early check
int(value) can return 0 (for Content-Length: 0) or a negative integer (e.g. Content-Length: -1, which is technically invalid HTTP but sent by some misbehaving clients). Neither will satisfy declared_size > limit for any positive limit, so the early fast-path rejection is skipped and the request falls through to the streaming loop instead. The streaming loop still enforces the limit correctly, so there is no security bypass, but adding an explicit non-positive guard makes the method's contract unambiguous.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/service/app.py
Line: 258-265
Comment:
**`_content_length` accepts non-positive values, silently bypassing the early check**
`int(value)` can return `0` (for `Content-Length: 0`) or a negative integer (e.g. `Content-Length: -1`, which is technically invalid HTTP but sent by some misbehaving clients). Neither will satisfy `declared_size > limit` for any positive limit, so the early fast-path rejection is skipped and the request falls through to the streaming loop instead. The streaming loop still enforces the limit correctly, so there is no security bypass, but adding an explicit non-positive guard makes the method's contract unambiguous.
How can I resolve this? If you propose a fix, please make it concise.| await send({"type": "http.response.body", "body": b"accepted"}) | ||
|
|
||
|
|
||
| def test_gateway_body_cache_rejects_oversized_content_length_without_reading_body() -> None: |
There was a problem hiding this comment.
max_body_bytes=None path (unlimited mode). This is the constructor default and is also used directly in unit tests; a regression that accidentally re-enables the limit when None is passed would go undetected. Consider adding a case that verifies bodies of arbitrary size pass through and are cached when no limit is configured.
| def test_gateway_body_cache_rejects_oversized_content_length_without_reading_body() -> None: | |
| def test_gateway_body_cache_no_limit_passes_all_bodies_through() -> None: | |
| downstream = _MultipartParsingApp() | |
| app = _GatewayBodyCacheMiddleware(downstream, max_body_bytes=None) | |
| payload = b"x" * 1_000_000 # 1 MB — would exceed any typical small limit | |
| 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}], | |
| ) | |
| ) | |
| # Should NOT reject; downstream should receive the full body | |
| assert downstream.called is True | |
| assert downstream.cached_body == payload | |
| def test_gateway_body_cache_rejects_oversized_content_length_without_reading_body() -> None: |
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/tests/test_gateway_body_cache.py
Line: 93
Comment:
There is no test covering the `max_body_bytes=None` path (unlimited mode). This is the constructor default and is also used directly in unit tests; a regression that accidentally re-enables the limit when `None` is passed would go undetected. Consider adding a case that verifies bodies of arbitrary size pass through and are cached when no limit is configured.
```suggestion
def test_gateway_body_cache_no_limit_passes_all_bodies_through() -> None:
downstream = _MultipartParsingApp()
app = _GatewayBodyCacheMiddleware(downstream, max_body_bytes=None)
payload = b"x" * 1_000_000 # 1 MB — would exceed any typical small limit
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}],
)
)
# Should NOT reject; downstream should receive the full body
assert downstream.called is True
assert downstream.cached_body == payload
def test_gateway_body_cache_rejects_oversized_content_length_without_reading_body() -> None:
```
How can I resolve this? If you propose a fix, please make it concise.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| 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_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_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" |
There was a problem hiding this comment.
Both checks use strict >, so a body whose size is exactly max_body_bytes should be accepted. There's currently no test that sends a body of exactly limit bytes and verifies a 200-series response. Given that the header-check (declared_size > limit) and the streaming check (body_size > limit) are the two separate code paths, a test with len(body) == limit would confirm both boundaries behave consistently.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/tests/test_gateway_body_cache.py
Line: 93-172
Comment:
**No test for exact-limit boundary (off-by-one check)**
Both checks use strict `>`, so a body whose size is exactly `max_body_bytes` should be accepted. There's currently no test that sends a body of exactly `limit` bytes and verifies a 200-series response. Given that the header-check (`declared_size > limit`) and the streaming check (`body_size > limit`) are the two separate code paths, a test with `len(body) == limit` would confirm both boundaries behave consistently.
How can I resolve this? If you propose a fix, please make it concise.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
fallintoplace
force-pushed
the
fix-gateway-body-limit
branch
from
fdebf3f to
22aac4c
Compare
1 participant
Summary
This caps gateway POST body buffering with the existing
resources.max_upload_byteslimit.Content-Lengthrequests before reading the body.Root Cause
Gateway mode buffered every POST body into memory before route-level upload validation ran. Large or chunked requests could allocate memory before the configured upload limit had a chance to reject them.
Validation
uv run pytest tests/test_gateway_body_cache.py tests/test_service_ingest_router.pygit diff --check