Make request body limit middleware return reliable 413 responses

app/core/request_limits.py

@@ -3,10 +3,16 @@ from __future__ import annotations
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 
+class _RequestBodyTooLarge(Exception):
+    """Internal signal used to stop an application from consuming an oversized body."""
+
+
 class RequestBodyLimitMiddleware:
-    """Enforce a hard byte ceiling on HTTP request bodies, including chunked bodies."""
+    """Enforce a request-body ceiling using Content-Length and streamed byte counts."""
 
     def __init__(self, app: ASGIApp, max_bytes: int) -> None:
+        if max_bytes < 1:
+            raise ValueError("max_bytes must be positive")
         self.app = app
         self.max_bytes = max_bytes
 
@@ -17,19 +23,23 @@ class RequestBodyLimitMiddleware:
 
         headers = {key.lower(): value for key, value in scope.get("headers", [])}
         declared = headers.get(b"content-length")
-        if declared:
+        if declared is not None:
             try:
-                if int(declared) > self.max_bytes:
-                    await self._too_large(scope, receive, send)
-                    return
+                content_length = int(declared)
             except ValueError:
-                # Invalid Content-Length is left to the ASGI server/parser; the
-                # streamed byte counter below remains authoritative.
-                pass
+                await self._send_413(send)
+                return
+            if content_length < 0:
+                await self._send_413(send)
+                return
+            if content_length > self.max_bytes:
+                await self._send_413(send)
+                return
 
         received = 0
         exceeded = False
         response_started = False
+        response_body_sent = False
 
         async def limited_receive() -> Message:
             nonlocal received, exceeded
@@ -38,43 +48,52 @@ class RequestBodyLimitMiddleware:
                 received += len(message.get("body", b""))
                 if received > self.max_bytes:
                     exceeded = True
-                    # Do not pass an over-limit chunk into a multipart parser or
-                    # request.body(); it could otherwise spool/read the full body.
-                    return {"type": "http.disconnect"}
+                    raise _RequestBodyTooLarge
             return message
 
         async def guarded_send(message: Message) -> None:
-            nonlocal response_started
-            if message["type"] == "http.response.start":
-                if exceeded and not response_started:
+            nonlocal response_started, response_body_sent
+            if not exceeded:
+                if message["type"] == "http.response.start":
                     response_started = True
-                    await send({
-                        "type": "http.response.start",
-                        "status": 413,
-                        "headers": [
-                            (b"content-type", b"application/json"),
-                            (b"cache-control", b"no-store"),
-                        ],
-                    })
-                    return
-                response_started = True
-            if exceeded:
-                if message["type"] == "http.response.body" and not response_started:
-                    return
-                if message["type"] == "http.response.body":
-                    await send({"type": "http.response.body", "body": b'{"detail":"Request body too large"}', "more_body": False})
+                await send(message)
                 return
-            await send(message)
+
+            # Replace any application response after overflow with one clean 413.
+            if message["type"] == "http.response.start" and not response_started:
+                response_started = True
+                body = b'{"detail":"Request body too large"}'
+                await send({
+                    "type": "http.response.start",
+                    "status": 413,
+                    "headers": [
+                        (b"content-type", b"application/json"),
+                        (b"content-length", str(len(body)).encode("ascii")),
+                        (b"cache-control", b"no-store"),
+                    ],
+                })
+            elif message["type"] == "http.response.body" and not response_body_sent:
+                response_body_sent = True
+                await send({
+                    "type": "http.response.body",
+                    "body": b'{"detail":"Request body too large"}',
+                    "more_body": False,
+                })
 
         try:
             await self.app(scope, limited_receive, guarded_send)
-        except Exception:
-            if exceeded and not response_started:
-                await self._too_large(scope, receive, send)
-                return
-            raise
+        except _RequestBodyTooLarge:
+            if not response_started:
+                await self._send_413(send)
+            elif not response_body_sent:
+                response_body_sent = True
+                await send({
+                    "type": "http.response.body",
+                    "body": b'{"detail":"Request body too large"}',
+                    "more_body": False,
+                })
 
-    async def _too_large(self, scope: Scope, receive: Receive, send: Send) -> None:
+    async def _send_413(self, send: Send) -> None:
         body = b'{"detail":"Request body too large"}'
         await send({
             "type": "http.response.start",