7381958225* studio: cap training dataset uploads * studio: clean up failed dataset uploads * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: raise upload limits to 500MB * studio: make upload limit configurable * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: stream upload routes * studio: split recipe upload caps * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: tighten upload limit handling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: import settings router directly * studio: polish upload cap setting control * studio: cap settings request bodies * studio: stub settings route in desktop auth test --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
396 lines
14 KiB
Python
396 lines
14 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"
|
|
assert "camera=()" in r.headers["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_src_allows_google_favicons(self, main_module):
|
|
# sources.tsx fetches https://www.google.com/s2/favicons?... ; without
|
|
# this allowlist entry citation favicons fall back to gray initials.
|
|
csp = main_module._build_csp()
|
|
img_directive = next(
|
|
chunk.strip()
|
|
for chunk in csp.split(";")
|
|
if chunk.strip().startswith("img-src ")
|
|
)
|
|
# Tokenise and compare with `==` so CodeQL's URL-substring rule does
|
|
# not read directive-string `in` membership as URL sanitisation.
|
|
img_sources = img_directive.split()
|
|
assert any(src == "https://www.google.com" for src in img_sources)
|
|
# Pre-existing favicon CDNs stay allowed.
|
|
for host in (
|
|
"https://t0.gstatic.com",
|
|
"https://t1.gstatic.com",
|
|
"https://t2.gstatic.com",
|
|
"https://t3.gstatic.com",
|
|
):
|
|
assert any(src == host for src in img_sources)
|
|
|
|
|
|
# =====================================================================
|
|
# /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 available 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 made 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}"
|