From a74ba71c07691bf6a5548ca54f7da093972cc13d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 17 Jun 2026 22:37:21 -0700 Subject: [PATCH] Studio: convert SecurityHeadersMiddleware to pure ASGI (#6394) * 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> --- studio/backend/main.py | 68 ++++++++++++++++--------- studio/backend/tests/test_middleware.py | 67 ++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 23 deletions(-) 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