wehub-resource-sync
123 lines
4.5 KiB
Python
123 lines
4.5 KiB
Python
import os
|
|
|
|
from loguru import logger
|
|
from starlette.datastructures import Headers
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
# Matches the file-size guidance already documented in
|
|
# docs/3-USER-GUIDE/adding-sources.md ("Very large files (>100MB) - Timeout").
|
|
DEFAULT_MAX_UPLOAD_SIZE_MB = 100
|
|
|
|
|
|
def get_max_upload_size_bytes() -> int:
|
|
"""Read the configured max request body size, in bytes.
|
|
|
|
Configurable via OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB for deployments that
|
|
need larger audio/video uploads; falls back to the default on unset,
|
|
malformed, or non-positive values (a zero/negative limit would reject
|
|
every request that has a body).
|
|
"""
|
|
raw = os.environ.get("OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB", "").strip()
|
|
try:
|
|
mb = float(raw) if raw else DEFAULT_MAX_UPLOAD_SIZE_MB
|
|
except ValueError:
|
|
mb = DEFAULT_MAX_UPLOAD_SIZE_MB
|
|
if mb <= 0:
|
|
logger.warning(
|
|
f"OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB={raw!r} is not a positive size; "
|
|
f"using the default of {DEFAULT_MAX_UPLOAD_SIZE_MB}MB"
|
|
)
|
|
mb = DEFAULT_MAX_UPLOAD_SIZE_MB
|
|
return int(mb * 1024 * 1024)
|
|
|
|
|
|
class _RequestBodyTooLarge(Exception):
|
|
pass
|
|
|
|
|
|
class MaxBodySizeMiddleware:
|
|
"""
|
|
Raw ASGI middleware rejecting requests whose body exceeds a configured
|
|
maximum, so a large upload can't exhaust memory/disk on a deployment
|
|
with no fronting proxy enforcing its own limit (e.g. the shipped
|
|
docker-compose.yml, which exposes the API directly).
|
|
|
|
Implemented at the raw ASGI level (not BaseHTTPMiddleware) so the check
|
|
can run ahead of FastAPI's own body/form parsing instead of after it.
|
|
Rejects on the `Content-Length` header up front when present (the common
|
|
case, and cheap), and also counts bytes as the body streams in - a
|
|
client can lie about Content-Length or omit it entirely with chunked
|
|
transfer-encoding.
|
|
"""
|
|
|
|
def __init__(self, app: ASGIApp, max_body_size: int) -> None:
|
|
self.app = app
|
|
self.max_body_size = max_body_size
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
headers = Headers(scope=scope)
|
|
content_length = headers.get("content-length")
|
|
if content_length is not None:
|
|
try:
|
|
if int(content_length) > self.max_body_size:
|
|
logger.warning(
|
|
f"Rejected {scope.get('method', '?')} {scope.get('path', '?')}: "
|
|
f"declared body of {content_length} bytes exceeds the "
|
|
f"{self.max_body_size}-byte limit"
|
|
)
|
|
await _send_413(send)
|
|
return
|
|
except ValueError:
|
|
pass # malformed header - fall through to streaming enforcement
|
|
|
|
total_size = 0
|
|
response_started = False
|
|
|
|
async def send_wrapper(message: Message) -> None:
|
|
nonlocal response_started
|
|
if message["type"] == "http.response.start":
|
|
response_started = True
|
|
await send(message)
|
|
|
|
async def receive_wrapper() -> Message:
|
|
nonlocal total_size
|
|
message = await receive()
|
|
if message["type"] == "http.request":
|
|
total_size += len(message.get("body") or b"")
|
|
if total_size > self.max_body_size:
|
|
raise _RequestBodyTooLarge()
|
|
return message
|
|
|
|
try:
|
|
await self.app(scope, receive_wrapper, send_wrapper)
|
|
except _RequestBodyTooLarge:
|
|
logger.warning(
|
|
f"Rejected {scope.get('method', '?')} {scope.get('path', '?')}: "
|
|
f"streamed body exceeded the {self.max_body_size}-byte limit"
|
|
)
|
|
if not response_started:
|
|
await _send_413(send)
|
|
# Else the app already started responding - nothing safe to send;
|
|
# let the connection drop rather than violate the ASGI protocol
|
|
# with a second response.start.
|
|
|
|
|
|
async def _send_413(send: Send) -> None:
|
|
await send(
|
|
{
|
|
"type": "http.response.start",
|
|
"status": 413,
|
|
"headers": [(b"content-type", b"application/json")],
|
|
}
|
|
)
|
|
await send(
|
|
{
|
|
"type": "http.response.body",
|
|
"body": b'{"detail":"Request body exceeds the maximum allowed upload size"}',
|
|
}
|
|
)
|