Spaces:
Running
Running
| 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}) | |