33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.requests import Request
|
|
|
|
from ..errors import error_response
|
|
from ..settings import get_settings
|
|
|
|
|
|
class BodyLimitMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(self, request: Request, call_next):
|
|
limit = get_settings().body_limit_bytes
|
|
cl = request.headers.get("content-length")
|
|
if cl is not None:
|
|
try:
|
|
if int(cl) > limit:
|
|
return error_response(request, 413, "payload_too_large", "body too large")
|
|
except ValueError:
|
|
return error_response(request, 400, "validation", "invalid content-length")
|
|
return await call_next(request)
|
|
|
|
total = 0
|
|
chunks: list[bytes] = []
|
|
async for chunk in request.stream():
|
|
total += len(chunk)
|
|
if total > limit:
|
|
return error_response(request, 413, "payload_too_large", "body too large")
|
|
chunks.append(chunk)
|
|
|
|
async def receive():
|
|
return {"type": "http.request", "body": b"".join(chunks), "more_body": False}
|
|
|
|
request._receive = receive # type: ignore[attr-defined]
|
|
return await call_next(request)
|