studio: enforce body cap on chunked uploads and drop unsafe-inline from script-src
MaxBodyMiddleware previously only inspected the declared Content-Length header; clients omitting it or sending Transfer-Encoding: chunked bypassed the cap and could still drive an OOM via the downstream JSON / file readers on /v1/chat/completions, /api/inference, /api/data-recipe, /api/datasets, /api/train, /api/export. Rewrite as a raw ASGI middleware that drains and counts http.request frames, replies 413 once the running total exceeds UNSLOTH_STUDIO_MAX_BODY_MB before invoking the FastAPI handler, and replays the buffered body to downstream so route code that calls request.json() / await request.body() works unchanged. CSP previously included 'unsafe-inline' on script-src, which defeats the main XSS protection. The frontend bundle does not need inline scripts; the only inline <script> the backend ever emits is _inject_bootstrap, which is opt-in via UNSLOTH_STUDIO_INJECT_BOOTSTRAP. Drop 'unsafe-inline' from script-src by default; when _inject_bootstrap fires, generate a per-response nonce, embed it on the inlined <script>, and have SecurityHeadersMiddleware splice 'nonce-XXX' into the CSP for that one response (the internal x-internal-script-nonce header is popped before the response leaves the server). 'unsafe-inline' stays on style-src for Vite-injected styles.
This commit is contained in:
parent
7e98120499
commit
06ec088a56
1 changed files with 173 additions and 58 deletions
|
|
@ -249,27 +249,43 @@ from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402
|
||||||
from starlette.requests import Request as _StarletteRequest # noqa: E402
|
from starlette.requests import Request as _StarletteRequest # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
_CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
|
||||||
|
|
||||||
|
|
||||||
|
def _build_csp(script_nonce: "str | None" = None) -> str:
|
||||||
|
script_src = "script-src 'self'"
|
||||||
|
if script_nonce:
|
||||||
|
script_src += f" 'nonce-{script_nonce}'"
|
||||||
|
return (
|
||||||
|
"default-src 'self'; "
|
||||||
|
"img-src 'self' data: blob: https://t0.gstatic.com "
|
||||||
|
"https://t1.gstatic.com https://t2.gstatic.com "
|
||||||
|
"https://t3.gstatic.com; "
|
||||||
|
"connect-src 'self'; "
|
||||||
|
"style-src 'self' 'unsafe-inline'; "
|
||||||
|
f"{script_src}; "
|
||||||
|
"font-src 'self' data:; "
|
||||||
|
"frame-ancestors 'none'; "
|
||||||
|
"form-action 'self'; "
|
||||||
|
"base-uri 'self'"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
|
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
|
||||||
"""Set baseline security headers on every response."""
|
"""Set baseline security headers on every response.
|
||||||
|
|
||||||
|
The CSP defaults to ``script-src 'self'``. When a downstream handler
|
||||||
|
needs to emit one inline ``<script>`` (currently only the opt-in
|
||||||
|
``_inject_bootstrap`` path), it attaches the nonce via the internal
|
||||||
|
header :data:`_CSP_SCRIPT_NONCE_HEADER`; we splice ``'nonce-XXX'`` into
|
||||||
|
the CSP for that response and strip the header before sending so the
|
||||||
|
nonce never leaves the server.
|
||||||
|
"""
|
||||||
|
|
||||||
async def dispatch(self, request: _StarletteRequest, call_next):
|
async def dispatch(self, request: _StarletteRequest, call_next):
|
||||||
response = await call_next(request)
|
response = await call_next(request)
|
||||||
response.headers.setdefault(
|
nonce = response.headers.pop(_CSP_SCRIPT_NONCE_HEADER, None)
|
||||||
"Content-Security-Policy",
|
response.headers.setdefault("Content-Security-Policy", _build_csp(nonce))
|
||||||
(
|
|
||||||
"default-src 'self'; "
|
|
||||||
"img-src 'self' data: blob: https://t0.gstatic.com "
|
|
||||||
"https://t1.gstatic.com https://t2.gstatic.com "
|
|
||||||
"https://t3.gstatic.com; "
|
|
||||||
"connect-src 'self'; "
|
|
||||||
"style-src 'self' 'unsafe-inline'; "
|
|
||||||
"script-src 'self' 'unsafe-inline'; "
|
|
||||||
"font-src 'self' data:; "
|
|
||||||
"frame-ancestors 'none'; "
|
|
||||||
"form-action 'self'; "
|
|
||||||
"base-uri 'self'"
|
|
||||||
),
|
|
||||||
)
|
|
||||||
response.headers.setdefault("X-Frame-Options", "DENY")
|
response.headers.setdefault("X-Frame-Options", "DENY")
|
||||||
response.headers.setdefault("X-Content-Type-Options", "nosniff")
|
response.headers.setdefault("X-Content-Type-Options", "nosniff")
|
||||||
response.headers.setdefault("Referrer-Policy", "no-referrer")
|
response.headers.setdefault("Referrer-Policy", "no-referrer")
|
||||||
|
|
@ -287,7 +303,6 @@ app.add_middleware(SecurityHeadersMiddleware)
|
||||||
# Cap upload body size on inference/dataset/training POSTs to prevent
|
# Cap upload body size on inference/dataset/training POSTs to prevent
|
||||||
# OOM via large base64-encoded attachments. Default 100 MB, env-tunable.
|
# OOM via large base64-encoded attachments. Default 100 MB, env-tunable.
|
||||||
import json as _json_for_413 # noqa: E402
|
import json as _json_for_413 # noqa: E402
|
||||||
from starlette.responses import JSONResponse as _JSONResponse # noqa: E402
|
|
||||||
|
|
||||||
|
|
||||||
_MAX_BODY_BYTES = int(os.environ.get("UNSLOTH_STUDIO_MAX_BODY_MB", "100")) * 1024 * 1024
|
_MAX_BODY_BYTES = int(os.environ.get("UNSLOTH_STUDIO_MAX_BODY_MB", "100")) * 1024 * 1024
|
||||||
|
|
@ -302,33 +317,126 @@ _BODY_PROTECTED_PREFIXES = (
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class MaxBodyMiddleware(BaseHTTPMiddleware):
|
async def _send_413(send, total_bytes: int) -> None:
|
||||||
async def dispatch(self, request: _StarletteRequest, call_next):
|
"""Send a 413 JSON response over a raw ASGI ``send`` channel."""
|
||||||
method = request.method.upper()
|
payload = _json_for_413.dumps(
|
||||||
path = request.url.path
|
{
|
||||||
if method in ("POST", "PUT", "PATCH") and any(
|
"detail": (
|
||||||
path.startswith(p) for p in _BODY_PROTECTED_PREFIXES
|
f"Request body too large "
|
||||||
|
f"({total_bytes:,} bytes; max {_MAX_BODY_BYTES:,})."
|
||||||
|
)
|
||||||
|
},
|
||||||
|
).encode("utf-8")
|
||||||
|
await send(
|
||||||
|
{
|
||||||
|
"type": "http.response.start",
|
||||||
|
"status": 413,
|
||||||
|
"headers": [
|
||||||
|
(b"content-type", b"application/json"),
|
||||||
|
(b"content-length", str(len(payload)).encode("ascii")),
|
||||||
|
],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
await send({"type": "http.response.body", "body": payload, "more_body": False})
|
||||||
|
|
||||||
|
|
||||||
|
class MaxBodyMiddleware:
|
||||||
|
"""Reject oversized request bodies on protected POST/PUT/PATCH prefixes.
|
||||||
|
|
||||||
|
The previous BaseHTTPMiddleware version only checked the declared
|
||||||
|
``Content-Length`` header; clients that omit it or send chunked
|
||||||
|
transfer-encoding bypassed the cap and could still drive an OOM via
|
||||||
|
the downstream JSON / file readers. This rewrite is a raw ASGI
|
||||||
|
middleware:
|
||||||
|
|
||||||
|
1. If ``Content-Length`` is declared and exceeds the cap, reply 413
|
||||||
|
immediately (same as before).
|
||||||
|
2. Otherwise the middleware buffers the request body itself, counting
|
||||||
|
bytes as ``http.request`` frames arrive. Once the running total
|
||||||
|
exceeds ``_MAX_BODY_BYTES`` the request is rejected with 413
|
||||||
|
before the FastAPI handler ever runs. The bounded buffer
|
||||||
|
(``_MAX_BODY_BYTES + 1`` worst case) is then replayed to
|
||||||
|
downstream as a single ``http.request`` so route code that calls
|
||||||
|
``request.json()`` / ``await request.body()`` works unchanged.
|
||||||
|
|
||||||
|
Non-protected prefixes, non-mutating methods, and non-HTTP scopes
|
||||||
|
bypass the cap as before.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, app, max_bytes: int, protected_prefixes: tuple):
|
||||||
|
self.app = app
|
||||||
|
self.max_bytes = max_bytes
|
||||||
|
self.protected_prefixes = protected_prefixes
|
||||||
|
|
||||||
|
async def __call__(self, scope, receive, send):
|
||||||
|
if scope["type"] != "http":
|
||||||
|
await self.app(scope, receive, send)
|
||||||
|
return
|
||||||
|
method = scope.get("method", "").upper()
|
||||||
|
path = scope.get("path", "")
|
||||||
|
if method not in ("POST", "PUT", "PATCH") or not any(
|
||||||
|
path.startswith(p) for p in self.protected_prefixes
|
||||||
):
|
):
|
||||||
cl = request.headers.get("content-length")
|
await self.app(scope, receive, send)
|
||||||
if cl:
|
return
|
||||||
|
|
||||||
|
declared = None
|
||||||
|
for name, value in scope.get("headers", []):
|
||||||
|
if name == b"content-length":
|
||||||
try:
|
try:
|
||||||
declared = int(cl)
|
declared = int(value.decode("latin-1"))
|
||||||
except ValueError:
|
except (ValueError, UnicodeDecodeError):
|
||||||
declared = None
|
declared = None
|
||||||
if declared is not None and declared > _MAX_BODY_BYTES:
|
break
|
||||||
return _JSONResponse(
|
if declared is not None and declared > self.max_bytes:
|
||||||
status_code = 413,
|
await _send_413(send, declared)
|
||||||
content = {
|
return
|
||||||
"detail": (
|
|
||||||
f"Request body too large "
|
# Drain the body up to max_bytes + 1, counting as we go.
|
||||||
f"({declared:,} bytes; max {_MAX_BODY_BYTES:,})."
|
chunks: list = []
|
||||||
)
|
total = 0
|
||||||
},
|
while True:
|
||||||
)
|
msg = await receive()
|
||||||
return await call_next(request)
|
mtype = msg.get("type")
|
||||||
|
if mtype == "http.disconnect":
|
||||||
|
# Client gave up mid-upload; bail without invoking downstream.
|
||||||
|
return
|
||||||
|
if mtype != "http.request":
|
||||||
|
# Unexpected ASGI message; forward as-is to downstream is unsafe
|
||||||
|
# because we are mid-stream — drop and abort.
|
||||||
|
return
|
||||||
|
body = msg.get("body", b"") or b""
|
||||||
|
if body:
|
||||||
|
total += len(body)
|
||||||
|
if total > self.max_bytes:
|
||||||
|
await _send_413(send, total)
|
||||||
|
return
|
||||||
|
chunks.append(body)
|
||||||
|
if not msg.get("more_body", False):
|
||||||
|
break
|
||||||
|
|
||||||
|
replayed = {"sent": False}
|
||||||
|
|
||||||
|
async def replay_receive():
|
||||||
|
if not replayed["sent"]:
|
||||||
|
replayed["sent"] = True
|
||||||
|
return {
|
||||||
|
"type": "http.request",
|
||||||
|
"body": b"".join(chunks),
|
||||||
|
"more_body": False,
|
||||||
|
}
|
||||||
|
# Subsequent receive() calls after the body — fall through to
|
||||||
|
# the real channel so http.disconnect etc still propagate.
|
||||||
|
return await receive()
|
||||||
|
|
||||||
|
await self.app(scope, replay_receive, send)
|
||||||
|
|
||||||
|
|
||||||
app.add_middleware(MaxBodyMiddleware)
|
app.add_middleware(
|
||||||
|
MaxBodyMiddleware,
|
||||||
|
max_bytes = _MAX_BODY_BYTES,
|
||||||
|
protected_prefixes = _BODY_PROTECTED_PREFIXES,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
# /recipes -> /data-recipes redirect so the bare URL works either way.
|
# /recipes -> /data-recipes redirect so the bare URL works either way.
|
||||||
|
|
@ -550,14 +658,21 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes:
|
||||||
return html.encode("utf-8")
|
return html.encode("utf-8")
|
||||||
|
|
||||||
|
|
||||||
def _inject_bootstrap(html_bytes: bytes, app: FastAPI) -> bytes:
|
def _inject_bootstrap(html_bytes: bytes, app: FastAPI):
|
||||||
"""Inject bootstrap credentials into HTML when password change is required.
|
"""Inject bootstrap credentials into HTML when password change is required.
|
||||||
|
|
||||||
The script tag is only injected while the default admin account still
|
The script tag is only injected while the default admin account still
|
||||||
has ``must_change_password=True``. Once the user changes the password
|
has ``must_change_password=True``. Once the user changes the password
|
||||||
the HTML is served clean — no credentials leak.
|
the HTML is served clean — no credentials leak.
|
||||||
|
|
||||||
|
Returns ``(html_bytes, script_nonce_or_None)``. When a nonce is
|
||||||
|
returned, callers MUST attach it via the internal CSP nonce header so
|
||||||
|
:class:`SecurityHeadersMiddleware` can splice ``'nonce-XXX'`` into the
|
||||||
|
``script-src`` directive for this one response. Without that the
|
||||||
|
inline ``<script>`` is blocked by the default ``script-src 'self'``.
|
||||||
"""
|
"""
|
||||||
import json as _json
|
import json as _json
|
||||||
|
import secrets as _secrets
|
||||||
|
|
||||||
# Opt-in only: the inject default leaked the plaintext bootstrap
|
# Opt-in only: the inject default leaked the plaintext bootstrap
|
||||||
# password to any LAN caller / browser extension. Users now read
|
# password to any LAN caller / browser extension. Users now read
|
||||||
|
|
@ -568,14 +683,14 @@ def _inject_bootstrap(html_bytes: bytes, app: FastAPI) -> bytes:
|
||||||
"yes",
|
"yes",
|
||||||
"on",
|
"on",
|
||||||
):
|
):
|
||||||
return html_bytes
|
return html_bytes, None
|
||||||
|
|
||||||
if not storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME):
|
if not storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME):
|
||||||
return html_bytes
|
return html_bytes, None
|
||||||
|
|
||||||
bootstrap_pw = getattr(app.state, "bootstrap_password", None)
|
bootstrap_pw = getattr(app.state, "bootstrap_password", None)
|
||||||
if not bootstrap_pw:
|
if not bootstrap_pw:
|
||||||
return html_bytes
|
return html_bytes, None
|
||||||
|
|
||||||
payload = _json.dumps(
|
payload = _json.dumps(
|
||||||
{
|
{
|
||||||
|
|
@ -583,10 +698,11 @@ def _inject_bootstrap(html_bytes: bytes, app: FastAPI) -> bytes:
|
||||||
"password": bootstrap_pw,
|
"password": bootstrap_pw,
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
tag = f"<script>window.__UNSLOTH_BOOTSTRAP__={payload}</script>"
|
nonce = _secrets.token_urlsafe(16)
|
||||||
|
tag = f'<script nonce="{nonce}">window.__UNSLOTH_BOOTSTRAP__={payload}</script>'
|
||||||
html = html_bytes.decode("utf-8")
|
html = html_bytes.decode("utf-8")
|
||||||
html = html.replace("</head>", f"{tag}</head>", 1)
|
html = html.replace("</head>", f"{tag}</head>", 1)
|
||||||
return html.encode("utf-8")
|
return html.encode("utf-8"), nonce
|
||||||
|
|
||||||
|
|
||||||
def setup_frontend(app: FastAPI, build_path: Path):
|
def setup_frontend(app: FastAPI, build_path: Path):
|
||||||
|
|
@ -599,17 +715,23 @@ def setup_frontend(app: FastAPI, build_path: Path):
|
||||||
if assets_dir.exists():
|
if assets_dir.exists():
|
||||||
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
|
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
|
||||||
|
|
||||||
@app.get("/")
|
def _build_index_response() -> Response:
|
||||||
async def serve_root():
|
|
||||||
content = (build_path / "index.html").read_bytes()
|
content = (build_path / "index.html").read_bytes()
|
||||||
content = _strip_crossorigin(content)
|
content = _strip_crossorigin(content)
|
||||||
content = _inject_bootstrap(content, app)
|
content, nonce = _inject_bootstrap(content, app)
|
||||||
|
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"}
|
||||||
|
if nonce:
|
||||||
|
headers[_CSP_SCRIPT_NONCE_HEADER] = nonce
|
||||||
return Response(
|
return Response(
|
||||||
content = content,
|
content = content,
|
||||||
media_type = "text/html",
|
media_type = "text/html",
|
||||||
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"},
|
headers = headers,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@app.get("/")
|
||||||
|
async def serve_root():
|
||||||
|
return _build_index_response()
|
||||||
|
|
||||||
@app.get("/{full_path:path}")
|
@app.get("/{full_path:path}")
|
||||||
async def serve_frontend(full_path: str):
|
async def serve_frontend(full_path: str):
|
||||||
if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")):
|
if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")):
|
||||||
|
|
@ -625,13 +747,6 @@ def setup_frontend(app: FastAPI, build_path: Path):
|
||||||
return FileResponse(file_path)
|
return FileResponse(file_path)
|
||||||
|
|
||||||
# Serve index.html as bytes — avoids Content-Length mismatch
|
# Serve index.html as bytes — avoids Content-Length mismatch
|
||||||
content = (build_path / "index.html").read_bytes()
|
return _build_index_response()
|
||||||
content = _strip_crossorigin(content)
|
|
||||||
content = _inject_bootstrap(content, app)
|
|
||||||
return Response(
|
|
||||||
content = content,
|
|
||||||
media_type = "text/html",
|
|
||||||
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"},
|
|
||||||
)
|
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue