Spaces:
Running
Running
| """Early, bounded ingress for the batch-check multipart upload.""" | |
| from __future__ import annotations | |
| from starlette.types import ASGIApp, Message, Receive, Scope, Send | |
| _TOO_LARGE = b'{"ok":false,"message":"That file is too large to check."}' | |
| # The route still enforces an exact 8 MiB file. This separate ingress allowance leaves | |
| # bounded room for multipart field names, headers, boundary markers, and the final delimiter. | |
| DEFAULT_MAX_BODY_BYTES = 8_000_000 + 64 * 1024 | |
| class _UploadTooLarge(Exception): | |
| pass | |
| class BatchCheckUploadLimitMiddleware: | |
| """Stop an oversized batch-check body before Starlette spools multipart parts.""" | |
| def __init__(self, app: ASGIApp, *, max_bytes: int = DEFAULT_MAX_BODY_BYTES) -> None: | |
| self.app, self.max_bytes = app, max_bytes | |
| async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: | |
| if scope["type"] != "http" or scope.get("method") != "POST" or scope.get("path") != "/api/batch-check": | |
| await self.app(scope, receive, send) | |
| return | |
| declared = dict(scope.get("headers") or ()).get(b"content-length", b"") | |
| if declared.isdigit() and int(declared) > self.max_bytes: | |
| await self._send_refusal(send) | |
| return | |
| total = 0 | |
| overflow = False | |
| async def limited_receive() -> Message: | |
| nonlocal total, overflow | |
| event = await receive() | |
| if event["type"] == "http.request": | |
| total += len(event.get("body", b"")) | |
| if total > self.max_bytes: | |
| overflow = True | |
| raise _UploadTooLarge | |
| return event | |
| async def limited_send(event: Message) -> None: | |
| # Starlette turns the receive exception into a multipart 400. Once overflow is | |
| # known, discard that response and preserve the route's existing JSON contract. | |
| if not overflow: | |
| await send(event) | |
| try: | |
| await self.app(scope, limited_receive, limited_send) | |
| except _UploadTooLarge: | |
| overflow = True | |
| if overflow: | |
| await self._send_refusal(send) | |
| async def _send_refusal(send: Send) -> None: | |
| await send({"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"application/json"), (b"content-length", str(len(_TOO_LARGE)).encode())]}) | |
| await send({"type": "http.response.body", "body": _TOO_LARGE}) | |