a74ba71c07* Studio: convert SecurityHeadersMiddleware to pure ASGI SecurityHeadersMiddleware was the last BaseHTTPMiddleware in the global stack, so every response (including SSE streams) was wrapped in an anyio stream that penalizes streaming. Rewrite it as a pure-ASGI middleware that mutates the response-start headers, mirroring the logging-middleware rewrite in #6337. The header logic is unchanged: it uses MutableHeaders over the start message, so the same get/del/setdefault calls apply (CSP nonce splice and strip, X-Frame-Options skip on Colab and the artifact-preview frame, the baseline nosniff/Referrer-Policy/Permissions-Policy/server headers). The existing middleware tests cover it; added cases assert headers still apply to a streaming response and that the artifact-preview path omits X-Frame-Options. * Studio: harden ASGI header coercion in SecurityHeadersMiddleware Review follow-up. MutableHeaders mutates its raw list in place, so if a server sends http.response.start with tuple-valued or missing headers the mutation would raise. Coerce to a list (defaulting to empty) before wrapping, then inject the same security headers as before. Also drop a stray em dash in a comment. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
450 lines
16 KiB
Python
450 lines
16 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
"""Tests for MaxBodyMiddleware, SecurityHeadersMiddleware, and the /api/health auth gate."""
|
|
|
|
import asyncio
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
from fastapi.responses import Response
|
|
from fastapi.testclient import TestClient
|
|
|
|
|
|
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(_BACKEND_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_BACKEND_ROOT))
|
|
|
|
|
|
@pytest.fixture(scope = "module")
|
|
def main_module():
|
|
import main as _main # noqa: F401
|
|
return _main
|
|
|
|
|
|
# MaxBodyMiddleware
|
|
|
|
|
|
def _make_protected_app(
|
|
max_bytes: int,
|
|
main_module,
|
|
upload_passthrough_prefixes: tuple = (),
|
|
upload_passthrough_max_bytes_getter = None,
|
|
):
|
|
app = FastAPI()
|
|
app.add_middleware(
|
|
main_module.MaxBodyMiddleware,
|
|
max_bytes_getter = lambda: max_bytes,
|
|
protected_prefixes = ("/v1/chat/completions", "/api/settings", "/api/train"),
|
|
upload_passthrough_prefixes = upload_passthrough_prefixes,
|
|
upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter,
|
|
)
|
|
|
|
@app.post("/v1/chat/completions")
|
|
async def chat(payload: dict):
|
|
return {"ok": True, "n": len(payload.get("text", ""))}
|
|
|
|
@app.post("/api/other")
|
|
async def other(payload: dict):
|
|
return {"ok": True, "unprotected": True}
|
|
|
|
@app.put("/api/settings/upload-limit")
|
|
async def update_upload_limit(payload: dict):
|
|
return {"ok": True, "limit": payload.get("max_upload_size_mb")}
|
|
|
|
@app.post("/api/train/upload")
|
|
async def upload(request: Request):
|
|
total = 0
|
|
chunks = 0
|
|
async for chunk in request.stream():
|
|
if chunk:
|
|
chunks += 1
|
|
total += len(chunk)
|
|
return {"ok": True, "chunks": chunks, "total": total}
|
|
|
|
@app.get("/api/train/status")
|
|
async def status_get():
|
|
return {"ok": True, "get": True}
|
|
|
|
return app
|
|
|
|
|
|
class TestMaxBodyMiddleware:
|
|
def test_small_protected_body_passes(self, main_module):
|
|
app = _make_protected_app(1024, main_module)
|
|
c = TestClient(app)
|
|
r = c.post("/v1/chat/completions", json = {"text": "x" * 100})
|
|
assert r.status_code == 200
|
|
assert r.json()["n"] == 100
|
|
|
|
def test_large_declared_content_length_rejected(self, main_module):
|
|
app = _make_protected_app(1024, main_module)
|
|
c = TestClient(app)
|
|
r = c.post("/v1/chat/completions", json = {"text": "x" * 5000})
|
|
assert r.status_code == 413
|
|
assert "too large" in r.json()["detail"].lower()
|
|
|
|
def test_unprotected_prefix_passes_large_body(self, main_module):
|
|
app = _make_protected_app(1024, main_module)
|
|
c = TestClient(app)
|
|
r = c.post("/api/other", json = {"text": "x" * 5000})
|
|
assert r.status_code == 200
|
|
assert r.json()["unprotected"] is True
|
|
|
|
def test_settings_put_body_over_cap_rejected(self, main_module):
|
|
app = _make_protected_app(1024, main_module)
|
|
c = TestClient(app)
|
|
r = c.put(
|
|
"/api/settings/upload-limit",
|
|
json = {"max_upload_size_mb": 500, "padding": "x" * 5000},
|
|
)
|
|
assert r.status_code == 413
|
|
assert "too large" in r.json()["detail"].lower()
|
|
|
|
def test_chunked_upload_over_cap_rejected(self, main_module):
|
|
# Regression: declared-Content-Length-only check could be bypassed by
|
|
# chunked transfer-encoding.
|
|
app = _make_protected_app(1024, main_module)
|
|
c = TestClient(app)
|
|
|
|
def gen():
|
|
yield b'{"text":"'
|
|
yield b"x" * 800
|
|
yield b'"}'
|
|
yield b"\n" + b"y" * 500
|
|
|
|
r = c.post(
|
|
"/v1/chat/completions",
|
|
content = gen(),
|
|
headers = {"content-type": "application/json"},
|
|
)
|
|
assert r.status_code == 413
|
|
assert "too large" in r.json()["detail"].lower()
|
|
|
|
def test_chunked_upload_under_cap_passes(self, main_module):
|
|
app = _make_protected_app(1024, main_module)
|
|
c = TestClient(app)
|
|
|
|
def gen():
|
|
yield b'{"text":"'
|
|
yield b"x" * 50
|
|
yield b'"}'
|
|
|
|
r = c.post(
|
|
"/v1/chat/completions",
|
|
content = gen(),
|
|
headers = {"content-type": "application/json"},
|
|
)
|
|
assert r.status_code == 200
|
|
assert r.json()["n"] == 50
|
|
|
|
def test_get_not_subject_to_cap(self, main_module):
|
|
app = _make_protected_app(1024, main_module)
|
|
c = TestClient(app)
|
|
r = c.get("/api/train/status")
|
|
assert r.status_code == 200
|
|
|
|
def test_upload_passthrough_uses_dedicated_declared_cap(self, main_module):
|
|
app = _make_protected_app(
|
|
128,
|
|
main_module,
|
|
upload_passthrough_prefixes = ("/api/train/upload",),
|
|
upload_passthrough_max_bytes_getter = lambda: 1024,
|
|
)
|
|
c = TestClient(app)
|
|
r = c.post(
|
|
"/api/train/upload",
|
|
content = b"x" * 512,
|
|
headers = {"content-type": "application/octet-stream"},
|
|
)
|
|
assert r.status_code == 200
|
|
assert r.json()["total"] == 512
|
|
|
|
def test_upload_passthrough_rejects_declared_body_over_dedicated_cap(self, main_module):
|
|
app = _make_protected_app(
|
|
128,
|
|
main_module,
|
|
upload_passthrough_prefixes = ("/api/train/upload",),
|
|
upload_passthrough_max_bytes_getter = lambda: 256,
|
|
)
|
|
c = TestClient(app)
|
|
r = c.post(
|
|
"/api/train/upload",
|
|
content = b"x" * 512,
|
|
headers = {"content-type": "application/octet-stream"},
|
|
)
|
|
assert r.status_code == 413
|
|
assert "256" in r.json()["detail"]
|
|
|
|
def test_upload_passthrough_requires_content_length(self, main_module):
|
|
app = _make_protected_app(
|
|
128,
|
|
main_module,
|
|
upload_passthrough_prefixes = ("/api/train/upload",),
|
|
upload_passthrough_max_bytes_getter = lambda: 1024,
|
|
)
|
|
c = TestClient(app)
|
|
|
|
def gen():
|
|
yield b"x" * 64
|
|
yield b"y" * 64
|
|
|
|
r = c.post(
|
|
"/api/train/upload",
|
|
content = gen(),
|
|
headers = {"content-type": "application/octet-stream"},
|
|
)
|
|
assert r.status_code == 411
|
|
assert "Content-Length" in r.json()["detail"]
|
|
|
|
|
|
# SecurityHeadersMiddleware / CSP
|
|
|
|
|
|
def _make_csp_app(main_module, attach_nonce: str | None = None):
|
|
app = FastAPI()
|
|
app.add_middleware(main_module.SecurityHeadersMiddleware)
|
|
|
|
@app.get("/plain")
|
|
async def plain():
|
|
return {"ok": True}
|
|
|
|
@app.get("/with-nonce")
|
|
async def with_nonce():
|
|
headers = {}
|
|
if attach_nonce:
|
|
headers[main_module._CSP_SCRIPT_NONCE_HEADER] = attach_nonce
|
|
return Response(
|
|
content = b"<html></html>",
|
|
media_type = "text/html",
|
|
headers = headers,
|
|
)
|
|
|
|
return app
|
|
|
|
|
|
class TestSecurityHeadersMiddleware:
|
|
def test_csp_has_no_unsafe_inline_for_script_src(self, main_module):
|
|
app = _make_csp_app(main_module)
|
|
c = TestClient(app)
|
|
r = c.get("/plain")
|
|
assert r.status_code == 200
|
|
csp = r.headers["content-security-policy"]
|
|
# Parse per-directive so style-src unsafe-inline does not false-match.
|
|
directives = {
|
|
chunk.strip().split(" ", 1)[0]: chunk.strip()
|
|
for chunk in csp.split(";")
|
|
if chunk.strip()
|
|
}
|
|
assert "script-src" in directives
|
|
assert "'unsafe-inline'" not in directives["script-src"]
|
|
# style-src keeps unsafe-inline for Vite-injected styles.
|
|
assert "'unsafe-inline'" in directives["style-src"]
|
|
|
|
def test_default_security_headers_present(self, main_module):
|
|
app = _make_csp_app(main_module)
|
|
c = TestClient(app)
|
|
r = c.get("/plain")
|
|
assert r.headers["x-frame-options"] == "DENY"
|
|
assert r.headers["x-content-type-options"] == "nosniff"
|
|
assert r.headers["referrer-policy"] == "no-referrer"
|
|
permissions_policy = r.headers["permissions-policy"]
|
|
assert "camera=()" in permissions_policy
|
|
assert "microphone=(self)" in permissions_policy
|
|
assert "geolocation=()" in permissions_policy
|
|
assert r.headers["server"] == "unsloth-studio"
|
|
|
|
def test_internal_nonce_header_is_spliced_into_csp_and_stripped(self, main_module):
|
|
nonce = "test-nonce-abc"
|
|
app = _make_csp_app(main_module, attach_nonce = nonce)
|
|
c = TestClient(app)
|
|
r = c.get("/with-nonce")
|
|
csp = r.headers["content-security-policy"]
|
|
assert f"'nonce-{nonce}'" in csp
|
|
# Internal handoff header must not leak to clients.
|
|
assert main_module._CSP_SCRIPT_NONCE_HEADER not in {k.lower() for k in r.headers.keys()}
|
|
|
|
def test_build_csp_helper_shape(self, main_module):
|
|
plain = main_module._build_csp()
|
|
assert "script-src 'self';" in plain
|
|
assert "'unsafe-inline'" not in plain.split("script-src", 1)[1].split(";", 1)[0]
|
|
nonced = main_module._build_csp("XYZ")
|
|
assert "script-src 'self' 'nonce-XYZ';" in nonced
|
|
|
|
def test_img_and_media_allow_https_sources(self, main_module):
|
|
# Model-card READMEs and citation favicons pull images/media from many
|
|
# https origins (HF LFS/XET CDNs, shields/badge hosts, GitHub-hosted
|
|
# assets, audio/video samples). img-src/media-src allow any https source
|
|
# so they render; this mirrors the desktop CSP in tauri.conf.json.
|
|
csp = main_module._build_csp()
|
|
directives = {
|
|
chunk.strip().split()[0]: chunk.strip().split()
|
|
for chunk in csp.split(";")
|
|
if chunk.strip()
|
|
}
|
|
for name in ("img-src", "media-src"):
|
|
assert name in directives, f"missing {name} directive"
|
|
# Tokenise and compare with `==` so CodeQL's URL-substring rule does
|
|
# not read directive-string `in` membership as URL sanitisation.
|
|
assert any(src == "https:" for src in directives[name])
|
|
|
|
def test_headers_applied_to_streaming_response(self, main_module):
|
|
# The ASGI middleware must set headers on streaming responses too.
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
app = FastAPI()
|
|
app.add_middleware(main_module.SecurityHeadersMiddleware)
|
|
|
|
@app.get("/stream")
|
|
async def stream():
|
|
async def gen():
|
|
yield b"a"
|
|
yield b"b"
|
|
|
|
return StreamingResponse(gen(), media_type = "text/plain")
|
|
|
|
r = TestClient(app).get("/stream")
|
|
assert r.status_code == 200
|
|
assert r.text == "ab"
|
|
assert r.headers["x-content-type-options"] == "nosniff"
|
|
assert r.headers["server"] == "unsloth-studio"
|
|
assert "content-security-policy" in r.headers
|
|
|
|
def test_artifact_preview_frame_omits_x_frame_options(self, main_module):
|
|
app = FastAPI()
|
|
app.add_middleware(main_module.SecurityHeadersMiddleware)
|
|
|
|
@app.get(main_module._ARTIFACT_PREVIEW_FRAME_PATH)
|
|
async def frame():
|
|
return Response(content = b"<html></html>", media_type = "text/html")
|
|
|
|
r = TestClient(app).get(main_module._ARTIFACT_PREVIEW_FRAME_PATH)
|
|
assert r.status_code == 200
|
|
assert "x-frame-options" not in {k.lower() for k in r.headers.keys()}
|
|
assert r.headers["referrer-policy"] == "no-referrer"
|
|
|
|
def test_response_start_with_tuple_headers_is_hardened(self, main_module):
|
|
# An ASGI server may emit tuple-valued raw headers; the middleware must
|
|
# coerce to a list and still inject security headers without crashing.
|
|
import asyncio
|
|
|
|
async def _inner_app(scope, receive, send):
|
|
await send(
|
|
{
|
|
"type": "http.response.start",
|
|
"status": 200,
|
|
"headers": ((b"content-type", b"text/plain"),), # tuple, not list
|
|
}
|
|
)
|
|
await send({"type": "http.response.body", "body": b"ok"})
|
|
|
|
captured = {}
|
|
|
|
async def _send(message):
|
|
if message["type"] == "http.response.start":
|
|
captured["headers"] = dict(message["headers"])
|
|
|
|
async def _receive():
|
|
return {"type": "http.request"}
|
|
|
|
mw = main_module.SecurityHeadersMiddleware(_inner_app)
|
|
asyncio.run(mw({"type": "http", "path": "/plain"}, _receive, _send))
|
|
|
|
hdrs = captured["headers"]
|
|
assert hdrs[b"server"] == b"unsloth-studio"
|
|
assert b"content-security-policy" in hdrs
|
|
assert hdrs[b"x-frame-options"] == b"DENY"
|
|
|
|
|
|
# /api/health auth gate
|
|
|
|
|
|
@pytest.fixture
|
|
def health_app(tmp_path, monkeypatch):
|
|
"""Mount /api/health on a fresh app against an isolated auth db."""
|
|
from auth import storage
|
|
|
|
monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
|
|
monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password")
|
|
monkeypatch.setattr(storage, "_bootstrap_password", None)
|
|
|
|
import main as _main
|
|
|
|
app = FastAPI()
|
|
app.add_api_route("/api/health", _main.health_check, methods = ["GET"])
|
|
|
|
import secrets as _secrets
|
|
|
|
storage.create_initial_user(
|
|
username = storage.DEFAULT_ADMIN_USERNAME,
|
|
password = "human-password-123",
|
|
jwt_secret = _secrets.token_urlsafe(64),
|
|
must_change_password = False,
|
|
)
|
|
return app
|
|
|
|
|
|
class TestHealthAuthGate:
|
|
# Launcher / frontend bootstrap fields are unauth so the Tauri watchdog can
|
|
# re-adopt a sibling backend and the SPA can detect chat-only mode before
|
|
# any token exists. Version / device_type still require a bearer.
|
|
LAUNCHER_BITS = (
|
|
"service",
|
|
"studio_root_id",
|
|
"chat_only",
|
|
"desktop_protocol_version",
|
|
"desktop_manageability_version",
|
|
"supports_desktop_auth",
|
|
"supports_desktop_backend_ownership",
|
|
"native_path_leases_supported",
|
|
)
|
|
FINGERPRINT_FIELDS = ("version", "studio_version", "device_type")
|
|
|
|
def test_no_auth_exposes_launcher_bits(self, health_app):
|
|
c = TestClient(health_app)
|
|
r = c.get("/api/health")
|
|
assert r.status_code == 200
|
|
body = r.json()
|
|
assert body["status"] == "healthy"
|
|
assert "timestamp" in body
|
|
for field in self.LAUNCHER_BITS:
|
|
assert field in body, f"missing launcher bit: {field}"
|
|
assert body["service"] == "Unsloth UI Backend"
|
|
for forbidden in self.FINGERPRINT_FIELDS:
|
|
assert forbidden not in body
|
|
|
|
def test_invalid_bearer_returns_launcher_bits_only(self, health_app):
|
|
# Regression: calling the async dep without await let any Bearer header pass.
|
|
c = TestClient(health_app)
|
|
r = c.get(
|
|
"/api/health",
|
|
headers = {"Authorization": "Bearer not-a-real-token"},
|
|
)
|
|
assert r.status_code == 200
|
|
body = r.json()
|
|
assert body["status"] == "healthy"
|
|
for field in self.LAUNCHER_BITS:
|
|
assert field in body
|
|
for forbidden in self.FINGERPRINT_FIELDS:
|
|
assert forbidden not in body
|
|
|
|
def test_valid_bearer_returns_full_payload(self, health_app):
|
|
from auth import storage
|
|
from auth.authentication import create_access_token
|
|
|
|
token = create_access_token(storage.DEFAULT_ADMIN_USERNAME)
|
|
c = TestClient(health_app)
|
|
r = c.get(
|
|
"/api/health",
|
|
headers = {"Authorization": f"Bearer {token}"},
|
|
)
|
|
assert r.status_code == 200
|
|
body = r.json()
|
|
assert body["status"] == "healthy"
|
|
for field in self.LAUNCHER_BITS + self.FINGERPRINT_FIELDS:
|
|
assert field in body, f"missing: {field}"
|