From c8e47dcb43cd293140de9f01dd84d9e50f35184b Mon Sep 17 00:00:00 2001 From: oss-agent-shin Date: Wed, 6 May 2026 12:29:11 -0700 Subject: [PATCH] Fix early proxy request size enforcement (#27311) * Add early proxy request size guard Co-authored-by: ishaan-berri * Address request size review feedback Co-authored-by: ishaan-berri --------- Co-authored-by: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com> Co-authored-by: ishaan-berri --- .../request_size_limit_middleware.py | 121 ++++++++++++++++ litellm/proxy/proxy_server.py | 8 ++ .../test_request_size_limit_middleware.py | 135 ++++++++++++++++++ 3 files changed, 264 insertions(+) create mode 100644 litellm/proxy/middleware/request_size_limit_middleware.py create mode 100644 tests/proxy_unit_tests/test_request_size_limit_middleware.py diff --git a/litellm/proxy/middleware/request_size_limit_middleware.py b/litellm/proxy/middleware/request_size_limit_middleware.py new file mode 100644 index 0000000000..78a38e3572 --- /dev/null +++ b/litellm/proxy/middleware/request_size_limit_middleware.py @@ -0,0 +1,121 @@ +import json +from typing import Callable, Optional, Union + +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +MaxRequestSizeGetter = Callable[[], Optional[Union[int, float]]] +RequestSizeLimitEnabledGetter = Callable[[], bool] + + +class RequestEntityTooLarge(Exception): + pass + + +class RequestSizeLimitMiddleware: + """ + Reject oversized requests before downstream auth/routes parse the body. + + Content-Length can be rejected without reading any body bytes. Requests + without Content-Length are counted as the ASGI stream is consumed, limiting + memory exposure to the configured threshold plus the current chunk. + """ + + def __init__( + self, + app: ASGIApp, + get_max_request_size_mb: MaxRequestSizeGetter, + is_request_size_limit_enabled: RequestSizeLimitEnabledGetter, + ) -> None: + self.app = app + self.get_max_request_size_mb = get_max_request_size_mb + self.is_request_size_limit_enabled = is_request_size_limit_enabled + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + max_request_size_mb = self.get_max_request_size_mb() + max_request_size_bytes = _mb_to_bytes(max_request_size_mb) + if max_request_size_bytes is None or not self.is_request_size_limit_enabled(): + await self.app(scope, receive, send) + return + + content_length = _get_content_length(scope=scope) + if content_length is not None and content_length > max_request_size_bytes: + await _send_request_too_large( + send=send, max_request_size_mb=max_request_size_mb + ) + return + + received_body_bytes = 0 + response_started = False + + async def limited_receive() -> Message: + nonlocal received_body_bytes + + message = await receive() + if message["type"] != "http.request": + return message + + received_body_bytes += len(message.get("body", b"")) + if received_body_bytes > max_request_size_bytes: + raise RequestEntityTooLarge + return message + + async def tracking_send(message: Message) -> None: + nonlocal response_started + + if message["type"] == "http.response.start": + response_started = True + await send(message) + + try: + await self.app(scope, limited_receive, tracking_send) + except RequestEntityTooLarge: + if response_started: + raise + await _send_request_too_large( + send=send, max_request_size_mb=max_request_size_mb + ) + + +def _mb_to_bytes(max_request_size_mb: Optional[Union[int, float]]) -> Optional[int]: + if max_request_size_mb is None: + return None + if max_request_size_mb <= 0: + return None + return int(max_request_size_mb * 1024 * 1024) + + +def _get_content_length(scope: Scope) -> Optional[int]: + headers = dict(scope.get("headers") or []) + raw_content_length = headers.get(b"content-length") + if raw_content_length is None: + return None + + try: + return int(raw_content_length) + except ValueError: + return None + + +async def _send_request_too_large( + send: Send, + max_request_size_mb: Optional[Union[int, float]], +) -> None: + body = json.dumps( + {"error": f"Request size is too large. Max size is {max_request_size_mb} MB"}, + separators=(",", ":"), + ).encode("utf-8") + await send( + { + "type": "http.response.start", + "status": 413, + "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(body)).encode("latin-1")), + ], + } + ) + await send({"type": "http.response.body", "body": body, "more_body": False}) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a5905765c6..5a379183e3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -403,6 +403,9 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware +from litellm.proxy.middleware.request_size_limit_middleware import ( + RequestSizeLimitMiddleware, +) from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, @@ -14881,6 +14884,11 @@ app.include_router(ui_discovery_endpoints_router) app.include_router(google_router) attach_lazy_features(app) +app.add_middleware( + RequestSizeLimitMiddleware, + get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"), + is_request_size_limit_enabled=lambda: premium_user is True, +) async def _stream_mcp_asgi_response( diff --git a/tests/proxy_unit_tests/test_request_size_limit_middleware.py b/tests/proxy_unit_tests/test_request_size_limit_middleware.py new file mode 100644 index 0000000000..3e8792e017 --- /dev/null +++ b/tests/proxy_unit_tests/test_request_size_limit_middleware.py @@ -0,0 +1,135 @@ +import pytest +from starlette.responses import JSONResponse +from starlette.testclient import TestClient +from starlette.types import Message + +from litellm.proxy.middleware.request_size_limit_middleware import ( + RequestSizeLimitMiddleware, +) + + +def test_request_size_limit_middleware_rejects_content_length_before_body_read(): + downstream_called = False + + async def app(scope, receive, send): + nonlocal downstream_called + downstream_called = True + response = JSONResponse({"ok": True}) + await response(scope, receive, send) + + client = TestClient( + RequestSizeLimitMiddleware( + app, + get_max_request_size_mb=lambda: 1, + is_request_size_limit_enabled=lambda: True, + ) + ) + + response = client.post( + "/chat/completions", + content=b"x" * (1024 * 1024 + 1), + headers={"content-type": "application/json"}, + ) + + assert response.status_code == 413 + assert response.json() == {"error": "Request size is too large. Max size is 1 MB"} + assert response.headers["content-length"] == str(len(response.content)) + assert downstream_called is False + + +def test_request_size_limit_middleware_zero_limit_disables_guard(): + downstream_called = False + + async def app(scope, receive, send): + nonlocal downstream_called + downstream_called = True + response = JSONResponse({"ok": True}) + await response(scope, receive, send) + + client = TestClient( + RequestSizeLimitMiddleware( + app, + get_max_request_size_mb=lambda: 0, + is_request_size_limit_enabled=lambda: True, + ) + ) + + response = client.post( + "/chat/completions", + content=b"x", + headers={"content-type": "application/json"}, + ) + + assert response.status_code == 200 + assert response.json() == {"ok": True} + assert downstream_called is True + + +@pytest.mark.asyncio +async def test_request_size_limit_middleware_rejects_streamed_body_without_content_length(): + received_body_bytes = 0 + + async def app(scope, receive, send): + nonlocal received_body_bytes + while True: + message = await receive() + if message["type"] == "http.disconnect": + break + received_body_bytes += len(message.get("body", b"")) + if not message.get("more_body", False): + break + + response = JSONResponse({"ok": True}) + await response(scope, receive, send) + + middleware = RequestSizeLimitMiddleware( + app, + get_max_request_size_mb=lambda: 1, + is_request_size_limit_enabled=lambda: True, + ) + sent_messages: list[Message] = [] + receive_messages: list[Message] = [ + { + "type": "http.request", + "body": b"x" * (1024 * 1024), + "more_body": True, + }, + { + "type": "http.request", + "body": b"y", + "more_body": False, + }, + ] + + async def receive(): + return receive_messages.pop(0) + + async def send(message): + sent_messages.append(message) + + await middleware( + { + "type": "http", + "method": "POST", + "path": "/chat/completions", + "headers": [(b"content-type", b"application/json")], + }, + receive, + send, + ) + + expected_body = b'{"error":"Request size is too large. Max size is 1 MB"}' + assert sent_messages[0] == { + "type": "http.response.start", + "status": 413, + "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(expected_body)).encode("latin-1")), + ], + } + assert sent_messages[1] == { + "type": "http.response.body", + "body": expected_body, + "more_body": False, + } + assert received_body_bytes == 1024 * 1024