Spaces:
Running
Running
| """ASGI body limit before JSON parsing, including chunked uploads.""" | |
| import asyncio | |
| import os | |
| from threading import Lock | |
| from starlette.responses import JSONResponse | |
| class BodyLimitMiddleware: | |
| def __init__(self, app, max_bytes=None, timeout=None): | |
| self.app = app | |
| self.active = 0 | |
| self.lock = Lock() | |
| self.max_concurrency = max(1, int(os.getenv("NEEDLE_BODY_CONCURRENCY", "16"))) | |
| self.max_bytes = max_bytes or int(os.getenv("NEEDLE_MAX_REQUEST_BYTES", str(8 * 1024 * 1024))) | |
| self.timeout = timeout or float(os.getenv("NEEDLE_BODY_TIMEOUT", "30")) | |
| async def __call__(self, scope, receive, send): | |
| if scope["type"] != "http" or scope["method"] not in {"POST", "PUT", "PATCH"}: | |
| return await self.app(scope, receive, send) | |
| with self.lock: | |
| admitted = self.active < self.max_concurrency | |
| if admitted: | |
| self.active += 1 | |
| if not admitted: | |
| return await JSONResponse({"detail": "Request concurrency limit"}, 429)(scope, receive, send) | |
| try: | |
| await self.handle(scope, receive, send) | |
| finally: | |
| with self.lock: | |
| self.active -= 1 | |
| async def handle(self, scope, receive, send): | |
| chunks, size = [], 0 | |
| try: | |
| async with asyncio.timeout(self.timeout): | |
| while True: | |
| message = await receive() | |
| if message["type"] == "http.disconnect": | |
| return | |
| body = message.get("body", b"") | |
| size += len(body) | |
| if size > self.max_bytes: | |
| return await JSONResponse({"detail": "Request body exceeds byte limit"}, 413)(scope, receive, send) | |
| chunks.append(body) | |
| if not message.get("more_body", False): | |
| break | |
| except TimeoutError: | |
| return await JSONResponse({"detail": "Request body timed out"}, 408)(scope, receive, send) | |
| scope["needle.body_bytes"] = size | |
| body = b"".join(chunks) | |
| delivered = False | |
| async def replay(): | |
| nonlocal delivered | |
| if not delivered: | |
| delivered = True | |
| return {"type": "http.request", "body": body, "more_body": False} | |
| return await receive() | |
| await self.app(scope, replay, send) | |