Spaces:
Running
Running
File size: 4,136 Bytes
8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f bda3e6b 8939f6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | 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 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
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http" or scope["method"] in {"GET", "HEAD", "OPTIONS"}:
await self.app(scope, receive, send)
return
headers = {key.lower(): value for key, value in scope.get("headers", [])}
declared = headers.get(b"content-length")
if declared is not None:
try:
content_length = int(declared)
except ValueError:
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
message = await receive()
if message["type"] == "http.request":
received += len(message.get("body", b""))
if received > self.max_bytes:
exceeded = True
raise _RequestBodyTooLarge
return message
async def guarded_send(message: Message) -> None:
nonlocal response_started, response_body_sent
if not exceeded:
if message["type"] == "http.response.start":
response_started = True
await send(message)
return
# 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 _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 _send_413(self, send: Send) -> None:
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"),
],
})
await send({"type": "http.response.body", "body": body, "more_body": False})
|