51 lines
1.9 KiB
Python
51 lines
1.9 KiB
Python
"""Bound security-sensitive request bodies before JSON/multipart parsing."""
|
|
|
|
from starlette.exceptions import HTTPException
|
|
from starlette.responses import JSONResponse
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
|
|
# Encrypted backup limit is 32 MiB. Allow a bounded margin for the multipart
|
|
# envelope; count streamed chunks as well as checking the untrusted header.
|
|
RESTORE_BODY_LIMIT = 34 * 1024 * 1024
|
|
BOOTSTRAP_BODY_LIMIT = 16 * 1024
|
|
|
|
|
|
class InstallationBodyLimitMiddleware:
|
|
def __init__(self, app: ASGIApp) -> None:
|
|
self.app = app
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
if scope["type"] != "http" or scope.get("method") != "POST":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
path = scope.get("path", "").rstrip("/")
|
|
limit = {
|
|
"/admin/backups/restore": RESTORE_BODY_LIMIT,
|
|
"/admin/backups/export": BOOTSTRAP_BODY_LIMIT,
|
|
"/setup/bootstrap": BOOTSTRAP_BODY_LIMIT,
|
|
}.get(path)
|
|
if limit is None:
|
|
await self.app(scope, receive, send)
|
|
return
|
|
headers = dict(scope.get("headers", []))
|
|
try:
|
|
length = int(headers.get(b"content-length", b"0"))
|
|
except ValueError:
|
|
length = -1
|
|
if length < 0 or length > limit:
|
|
await JSONResponse({"detail": "Request body is too large or has an invalid length."}, status_code=413)(scope, receive, send)
|
|
return
|
|
received = 0
|
|
|
|
async def bounded_receive() -> Message:
|
|
nonlocal received
|
|
message = await receive()
|
|
if message["type"] == "http.request":
|
|
received += len(message.get("body", b""))
|
|
if received > limit:
|
|
raise HTTPException(status_code=413, detail="Request body is too large.")
|
|
return message
|
|
|
|
await self.app(scope, bounded_receive, send)
|