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>
This commit is contained in:
Daniel Han
GitHub
parent
77acea751c
commit
a74ba71c07
2 changed files with 113 additions and 24 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"<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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue