102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
from collections.abc import Callable
|
|
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
|
|
SESSION_BODY_LIMIT = 16 * 1024
|
|
API_MUTATION_BODY_LIMIT = 64 * 1024
|
|
ISSUE_ATTACHMENT_BODY_LIMIT = 3 * 1024 * 1024
|
|
MUTATION_METHODS = frozenset({"POST", "PUT", "PATCH"})
|
|
|
|
|
|
def request_body_limit(method: str, path: str) -> int | None:
|
|
"""Return the admission limit for request bodies that FastAPI will parse."""
|
|
normalized_method = method.upper()
|
|
if normalized_method == "POST" and path == "/api/v1/session":
|
|
return SESSION_BODY_LIMIT
|
|
if (
|
|
normalized_method == "POST"
|
|
and path.startswith("/api/v1/repos/")
|
|
and "/issues/" in path
|
|
and path.endswith("/attachments")
|
|
):
|
|
return ISSUE_ATTACHMENT_BODY_LIMIT
|
|
if normalized_method in MUTATION_METHODS and path.startswith("/api/v1/"):
|
|
return API_MUTATION_BODY_LIMIT
|
|
return None
|
|
|
|
|
|
class RequestBodyLimitMiddleware:
|
|
"""Reject oversized API bodies before framework parsing or endpoint work."""
|
|
|
|
def __init__(
|
|
self,
|
|
app: ASGIApp,
|
|
limit_for: Callable[[str, str], int | None] = request_body_limit,
|
|
) -> None:
|
|
self.app = app
|
|
self.limit_for = limit_for
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
path = scope.get("path", "")
|
|
root_path = scope.get("root_path", "").rstrip("/")
|
|
if root_path and (path == root_path or path.startswith(root_path + "/")):
|
|
path = path[len(root_path) :] or "/"
|
|
limit = self.limit_for(scope.get("method", "GET"), path)
|
|
if limit is None:
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
declared_length = self._content_length(scope)
|
|
if declared_length is not None and declared_length > limit:
|
|
await self._reject(send)
|
|
return
|
|
|
|
messages: list[Message] = []
|
|
received = 0
|
|
while True:
|
|
message = await receive()
|
|
messages.append(message)
|
|
if message["type"] != "http.request":
|
|
break
|
|
received += len(message.get("body", b""))
|
|
if received > limit:
|
|
await self._reject(send)
|
|
return
|
|
if not message.get("more_body", False):
|
|
break
|
|
|
|
async def replay() -> Message:
|
|
if messages:
|
|
return messages.pop(0)
|
|
return {"type": "http.request", "body": b"", "more_body": False}
|
|
|
|
await self.app(scope, replay, send)
|
|
|
|
@staticmethod
|
|
def _content_length(scope: Scope) -> int | None:
|
|
for name, value in scope.get("headers", []):
|
|
if name.lower() != b"content-length":
|
|
continue
|
|
try:
|
|
parsed = int(value)
|
|
except ValueError:
|
|
return None
|
|
return max(0, parsed)
|
|
return None
|
|
|
|
@staticmethod
|
|
async def _reject(send: Send) -> None:
|
|
body = b'{"detail":"Request body too large"}'
|
|
headers = [
|
|
(b"content-type", b"application/json"),
|
|
(b"content-length", str(len(body)).encode("ascii")),
|
|
(b"cache-control", b"no-store"),
|
|
]
|
|
await send({"type": "http.response.start", "status": 413, "headers": headers})
|
|
await send({"type": "http.response.body", "body": body})
|