diff --git a/studio/backend/main.py b/studio/backend/main.py index 08c7bc6826..aaeb7917d8 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -492,8 +492,7 @@ app.add_middleware(LoggingMiddleware) # img/media-src allow any https origin so HF model-card assets render (mirrors # tauri.conf.json); scripts/frames/connect-src stay same-origin + HF. -from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402 -from starlette.requests import Request as _StarletteRequest # noqa: E402 +from starlette.datastructures import MutableHeaders # noqa: E402 _CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce" @@ -549,28 +548,51 @@ def _build_csp(script_nonce: "str | None" = None) -> str: ) -class SecurityHeadersMiddleware(BaseHTTPMiddleware): - """Set baseline security headers; splice per-response inline-script nonces into CSP.""" +class SecurityHeadersMiddleware: + """Set baseline security headers; splice per-response inline-script nonces into CSP. - async def dispatch(self, request: _StarletteRequest, call_next): - response = await call_next(request) - # Strip the internal nonce hand-off header so it never reaches the client - nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER) - if nonce is not None: - del response.headers[_CSP_SCRIPT_NONCE_HEADER] - response.headers.setdefault("Content-Security-Policy", _build_csp(nonce)) - # Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and - # DENY would block serve_kernel_port_as_iframe regardless of CSP. - if not _IS_COLAB and request.url.path != _ARTIFACT_PREVIEW_FRAME_PATH: - response.headers.setdefault("X-Frame-Options", "DENY") - response.headers.setdefault("X-Content-Type-Options", "nosniff") - response.headers.setdefault("Referrer-Policy", "no-referrer") - response.headers.setdefault( - "Permissions-Policy", - "camera=(), microphone=(self), geolocation=()", - ) - response.headers["server"] = "unsloth-studio" - return response + Pure ASGI (not BaseHTTPMiddleware) so streaming responses are not wrapped in + an anyio stream. Header logic mirrors the prior version exactly via + MutableHeaders on the response-start message. + """ + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + path = scope.get("path", "") + + async def send_wrapper(message): + if message["type"] == "http.response.start": + # ASGI headers are an iterable; coerce to a list so MutableHeaders + # can mutate in place even if a server sends a tuple or omits it. + raw = message.setdefault("headers", []) + if not isinstance(raw, list): + raw = list(raw) + message["headers"] = raw + headers = MutableHeaders(raw = raw) + # Strip the internal nonce hand-off header so it never reaches the client + nonce = headers.get(_CSP_SCRIPT_NONCE_HEADER) + if nonce is not None: + del headers[_CSP_SCRIPT_NONCE_HEADER] + headers.setdefault("Content-Security-Policy", _build_csp(nonce)) + # Omit X-Frame-Options in Colab: CSP frame-ancestors handles it, and + # DENY would block serve_kernel_port_as_iframe regardless of CSP. + if not _IS_COLAB and path != _ARTIFACT_PREVIEW_FRAME_PATH: + headers.setdefault("X-Frame-Options", "DENY") + headers.setdefault("X-Content-Type-Options", "nosniff") + headers.setdefault("Referrer-Policy", "no-referrer") + headers.setdefault( + "Permissions-Policy", + "camera=(), microphone=(self), geolocation=()", + ) + headers["server"] = "unsloth-studio" + await send(message) + + await self.app(scope, receive, send_wrapper) app.add_middleware(SecurityHeadersMiddleware) diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 1005431926..6379fc88e0 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -293,6 +293,73 @@ class TestSecurityHeadersMiddleware: # 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"", 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