Spaces:
Running
Running
File size: 2,427 Bytes
faaafef | 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 | """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)
|