83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
"""Payload size middleware and rate-limit setup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from fastapi import Request, Response
|
|
from slowapi import Limiter
|
|
from slowapi.util import get_remote_address
|
|
|
|
from .auth import extract_bearer_token
|
|
|
|
PAYLOAD_LIMIT_BYTES = 2 * 1024 * 1024 # 2 MiB
|
|
|
|
_TOO_LARGE_BODY = json.dumps(
|
|
{
|
|
"error": "payload_too_large",
|
|
"detail": f"Request body exceeds {PAYLOAD_LIMIT_BYTES} bytes",
|
|
}
|
|
).encode()
|
|
|
|
|
|
def _get_token_or_ip(request: Request) -> str:
|
|
"""Key function for rate limiting: use Bearer token if present, else IP."""
|
|
return extract_bearer_token(request) or get_remote_address(request)
|
|
|
|
|
|
limiter = Limiter(key_func=_get_token_or_ip, default_limits=[])
|
|
|
|
|
|
def _too_large() -> Response:
|
|
return Response(
|
|
content=_TOO_LARGE_BODY,
|
|
status_code=413,
|
|
media_type="application/json",
|
|
)
|
|
|
|
|
|
async def payload_size_middleware(request: Request, call_next) -> Response:
|
|
"""Reject requests larger than PAYLOAD_LIMIT_BYTES.
|
|
|
|
Two-layer enforcement:
|
|
1. If `Content-Length` is present, reject up front (cheap and avoids
|
|
buffering anything).
|
|
2. Otherwise (chunked encoding / missing header) wrap `request.receive`
|
|
so the body stream is tallied as it arrives; abort the moment the
|
|
running total exceeds the cap. This is defense-in-depth against
|
|
clients that omit Content-Length.
|
|
"""
|
|
content_length = request.headers.get("content-length")
|
|
if content_length is not None:
|
|
try:
|
|
length = int(content_length)
|
|
except ValueError:
|
|
length = 0
|
|
if length > PAYLOAD_LIMIT_BYTES:
|
|
return _too_large()
|
|
# Trust the declared length; the ASGI server enforces it.
|
|
return await call_next(request)
|
|
|
|
# No Content-Length: tally bytes as they arrive.
|
|
received_total = 0
|
|
too_large = False
|
|
original_receive = request.receive
|
|
|
|
async def guarded_receive() -> dict:
|
|
nonlocal received_total, too_large
|
|
message = await original_receive()
|
|
if message.get("type") == "http.request":
|
|
body = message.get("body") or b""
|
|
received_total += len(body)
|
|
if received_total > PAYLOAD_LIMIT_BYTES:
|
|
too_large = True
|
|
# Tell downstream "no more body" — they may still try to
|
|
# parse what's been seen, but we'll override the response.
|
|
return {"type": "http.disconnect"}
|
|
return message
|
|
|
|
request._receive = guarded_receive # type: ignore[attr-defined]
|
|
response = await call_next(request)
|
|
if too_large:
|
|
return _too_large()
|
|
return response
|