diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 233a0367c..e4c94d70c 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -37,7 +37,7 @@ import sys import threading import time from collections import OrderedDict -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import fields, is_dataclass, replace from datetime import datetime, timezone from pathlib import Path @@ -2605,6 +2605,92 @@ _is_known_websocket_callback_failure = is_known_websocket_callback_failure _tool_schema_saved_from_tags = tool_schema_saved_from_tags +def read_proxy_token(headers: Mapping[str, str]) -> str | None: + """Return the caller-supplied proxy token from request headers, or ``None``. + + Shared by the HTTP security gate and :class:`WebSocketAuthMiddleware` so the + two transports cannot drift on what counts as a credential. Header names are + expected to be lowercase (Starlette's ``Headers`` is case-insensitive; the + WebSocket middleware lowercases the raw ASGI pairs itself). + """ + auth = str(headers.get("authorization") or "") + if auth.lower().startswith("bearer "): + return auth[7:].strip() or None + raw = headers.get("x-headroom-proxy-token") + return str(raw) if raw else None + + +class WebSocketAuthMiddleware: + """Enforce ``HEADROOM_PROXY_TOKEN`` on WebSocket handshakes. + + The HTTP security gate is registered with ``@app.middleware("http")``, which + is a Starlette ``BaseHTTPMiddleware`` — and that class hands any scope whose + type is not ``http`` straight to the wrapped app. WebSocket connections + therefore never reached the gate, so every ``app.websocket(...)`` route + accepted unauthenticated callers even with a token configured. Those routes + are not incidental: ``/v1/responses`` and ``/v1/live`` relay to the upstream + provider using the operator's own credentials, and they are registered + unconditionally. ``/v1/responses`` exists on both transports, so the POST was + authenticated while the upgrade on the very same path was not. + + Written as a raw ASGI middleware rather than folded into the gate because + that is the only layer that sees the ``websocket`` scope at all. + + Loopback callers are exempt, matching the HTTP gate exactly (same trust + boundary as the admin/debug routes). Credentials are read from headers only: + the handshake carries them fine for the programmatic clients these routes + serve, and accepting a token from the query string would put it in access + logs and browser history. + """ + + def __init__(self, app: Any, *, proxy_token: str | None = None) -> None: + self.app = app + self.proxy_token = proxy_token + # Pre-encoded for constant-time comparison, mirroring the HTTP gate: + # compare_digest on str raises TypeError for non-ASCII input, which + # would turn a rejected handshake into a 500. + self.token_bytes = proxy_token.encode("utf-8") if proxy_token else b"" + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + if scope["type"] != "websocket" or not self.proxy_token: + await self.app(scope, receive, send) + return + + client = scope.get("client") + client_host = client[0] if client else None + if is_loopback_host(client_host): + await self.app(scope, receive, send) + return + + # Starlette's own Headers rather than a hand-built dict: on a repeated + # header it returns the FIRST occurrence, which is what the HTTP gate + # sees. Building a dict here instead took the LAST one, so the two + # transports disagreed about which `Authorization` counted — exactly the + # drift the shared reader below exists to prevent. + from starlette.datastructures import Headers + + provided = read_proxy_token(Headers(scope=scope)) + if provided is not None and hmac.compare_digest( + provided.encode("utf-8", "replace"), self.token_bytes + ): + await self.app(scope, receive, send) + return + + logger.warning( + "event=proxy_auth_rejected transport=websocket path=%s client=%s reason=%s", + scope.get("path"), + client_host, + "missing_token" if provided is None else "bad_token", + ) + # Receive the handshake before refusing it: ASGI servers send + # ``websocket.connect`` and wait for the application to answer, and + # answering with ``websocket.close`` *before* an accept is what refuses + # the upgrade on the wire instead of accepting and then dropping it. + message = await receive() + if message["type"] == "websocket.connect": + await send({"type": "websocket.close", "code": 1008}) + + class WebSocketProjectPrefixMiddleware: """Normalize project-prefixed WebSocket paths before route matching.""" @@ -3450,10 +3536,12 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: ) return response - # ── Security gate (registered last → runs outermost) ────────────────── - # Three concerns, kept together because they all wrap every inbound + # ── Security gate (outermost of the HTTP middlewares) ───────────────── + # Three concerns, kept together because they all wrap every inbound HTTP # request: optional inbound auth on the data plane, response security # headers, and an audit trail for state-mutating admin endpoints. + # WebSocket handshakes are covered separately — see the + # WebSocketAuthMiddleware registration just below this block. _proxy_token = config.proxy_token or os.environ.get("HEADROOM_PROXY_TOKEN") or None # Pre-encode once for constant-time comparison (compare_digest on str raises # TypeError for non-ASCII input, which would turn a 401 into a 500). @@ -3483,12 +3571,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: "Strict-Transport-Security", "max-age=31536000; includeSubDomains" ) - def _extract_proxy_token(headers) -> str | None: - auth = str(headers.get("authorization") or "") - if auth.lower().startswith("bearer "): - return auth[7:].strip() or None - raw = headers.get("x-headroom-proxy-token") - return str(raw) if raw else None + # Delegates so the HTTP gate and WebSocketAuthMiddleware read a credential + # by exactly one rule; they guard the same token on two transports. + _extract_proxy_token = read_proxy_token @app.middleware("http") async def _security_gate(request, call_next): @@ -3530,6 +3615,12 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: logger.debug("admin audit emission failed", exc_info=True) return response + # The gate above is http-only (BaseHTTPMiddleware ignores every other + # scope), so the same token rule is applied to the `websocket` scope here. + # Added after it, which makes it the outermost layer — an unauthenticated + # handshake is refused before any project-prefix or routing work happens. + app.add_middleware(WebSocketAuthMiddleware, proxy_token=_proxy_token) + # Third-party proxy extensions (Enterprise, custom plugins). Discovered via # the `headroom.proxy_extension` entry-point group, but **opt-in only**: # only names listed in config.proxy_extensions (CLI: --proxy-extension, diff --git a/tests/test_proxy_hardening.py b/tests/test_proxy_hardening.py index 65d158c02..85ddec4d8 100644 --- a/tests/test_proxy_hardening.py +++ b/tests/test_proxy_hardening.py @@ -19,7 +19,7 @@ from fastapi.testclient import TestClient from headroom.cache.compression_store import reset_compression_store from headroom.offline import apply_offline_env, is_offline from headroom.proxy.audit import is_auditable_path -from headroom.proxy.server import ProxyConfig, create_app +from headroom.proxy.server import ProxyConfig, WebSocketAuthMiddleware, create_app NONLOOPBACK = ("203.0.113.5", 44444) # TEST-NET-3, never loopback LOOPBACK = ("127.0.0.1", 12345) @@ -85,6 +85,217 @@ class TestInboundAuthToken: assert c.get("/readyz").status_code in (200, 503) # ready/not-ready, never 401 +# ──────────────────── 2.1b inbound auth token over WebSocket ────────────── + + +WS_PATHS = ("/v1/responses", "/v1/live") + + +class _SpyApp: + """Downstream ASGI app that records whether it was ever reached.""" + + def __init__(self) -> None: + self.called = False + + async def __call__(self, scope, receive, send) -> None: + self.called = True + + +def _ws_scope(*, client=NONLOOPBACK, headers=(), path="/v1/responses"): + return { + "type": "websocket", + "path": path, + "client": client, + "headers": [(k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in headers], + } + + +async def _drive(middleware, scope): + """Run one connection through the middleware, returning (sent, downstream).""" + inbox = [{"type": "websocket.connect"}] + sent: list[dict] = [] + + async def receive(): + return inbox.pop(0) if inbox else {"type": "websocket.disconnect"} + + async def send(message): + sent.append(message) + + await middleware(scope, receive, send) + return sent + + +def _closed_with_policy_violation(sent) -> bool: + return any(m.get("type") == "websocket.close" and m.get("code") == 1008 for m in sent) + + +class TestWebSocketAuthMiddleware: + """The middleware itself, driven directly over ASGI. + + Asserted at this layer because a pre-accept close surfaces through + ``TestClient`` as a bare ``AttributeError`` — indistinguishable from any + other handshake failure — so an exception-shape assertion would pass for + the wrong reason. + """ + + async def test_rejects_missing_credential(self): + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token="s3cr3t-token") + + sent = await _drive(mw, _ws_scope()) + + assert downstream.called is False + assert _closed_with_policy_violation(sent) + + async def test_rejects_wrong_credential(self): + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token="s3cr3t-token") + + sent = await _drive(mw, _ws_scope(headers=[("authorization", "Bearer wrong")])) + + assert downstream.called is False + assert _closed_with_policy_violation(sent) + + async def test_accepts_correct_bearer(self): + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token="s3cr3t-token") + + sent = await _drive(mw, _ws_scope(headers=[("authorization", "Bearer s3cr3t-token")])) + + assert downstream.called is True + assert not _closed_with_policy_violation(sent) + + async def test_accepts_custom_header(self): + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token="s3cr3t-token") + + sent = await _drive(mw, _ws_scope(headers=[("x-headroom-proxy-token", "s3cr3t-token")])) + + assert downstream.called is True + assert not _closed_with_policy_violation(sent) + + async def test_loopback_is_exempt(self): + """Same trust boundary the HTTP gate already grants loopback.""" + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token="s3cr3t-token") + + sent = await _drive(mw, _ws_scope(client=LOOPBACK)) + + assert downstream.called is True + assert not _closed_with_policy_violation(sent) + + async def test_unknown_client_is_treated_as_loopback(self): + """Mirrors is_loopback_host(None) -> True, as the HTTP gate does.""" + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token="s3cr3t-token") + + sent = await _drive(mw, _ws_scope(client=None)) + + assert downstream.called is True + assert not _closed_with_policy_violation(sent) + + async def test_repeated_header_resolves_like_the_http_gate(self): + """A duplicated Authorization must mean the same thing on both transports. + + Starlette's Headers (what the HTTP gate reads) returns the FIRST + occurrence. A hand-built dict returns the last, which would let the two + paths disagree about which credential counted. + """ + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token="s3cr3t-token") + + sent = await _drive( + mw, + _ws_scope( + headers=[ + ("authorization", "Bearer s3cr3t-token"), + ("authorization", "Bearer wrong"), + ] + ), + ) + + # First header wins → authenticated, same as the HTTP gate. + assert downstream.called is True + assert not _closed_with_policy_violation(sent) + + async def test_no_token_configured_is_a_passthrough(self): + """Default deployment must gain no new challenge.""" + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token=None) + + sent = await _drive(mw, _ws_scope()) + + assert downstream.called is True + assert not _closed_with_policy_violation(sent) + + async def test_http_scope_is_left_to_the_http_gate(self): + downstream = _SpyApp() + mw = WebSocketAuthMiddleware(downstream, proxy_token="s3cr3t-token") + + sent = await _drive(mw, {**_ws_scope(), "type": "http"}) + + assert downstream.called is True + assert not _closed_with_policy_violation(sent) + + +class TestWebSocketRoutesAreGatedInTheApp: + """The middleware is actually wired into ``create_app``. + + Asserts the security property directly — the route handler must never run + for an unauthenticated handshake — rather than inspecting the exception the + client happens to see. + """ + + @pytest.mark.parametrize("path", WS_PATHS) + def test_unauthenticated_handshake_never_reaches_the_handler(self, path, monkeypatch): + app = _make_app(proxy_token="s3cr3t-token") + reached = _record_ws_handler_reached(app, monkeypatch) + + with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c: + try: + with c.websocket_connect(path): + pass + except Exception: # noqa: BLE001 - the refusal shape is asserted above + pass + + assert reached() is False + + @pytest.mark.parametrize("path", WS_PATHS) + def test_authenticated_handshake_reaches_the_handler(self, path, monkeypatch): + app = _make_app(proxy_token="s3cr3t-token") + reached = _record_ws_handler_reached(app, monkeypatch) + + with TestClient(app, base_url="http://testserver", client=NONLOOPBACK) as c: + try: + with c.websocket_connect(path, headers={"X-Headroom-Proxy-Token": "s3cr3t-token"}): + pass + except Exception: # noqa: BLE001 - route may fail with no upstream + pass + + assert reached() is True + + +def _record_ws_handler_reached(app, monkeypatch): + """Spy both WebSocket route families; returns a callable reporting arrival.""" + from headroom.providers import proxy_routes + + seen: list[str] = [] + + # Each spy must terminate the handshake itself: a handler that returns + # without accepting or closing leaves the client waiting forever. + async def _responses_spy(websocket): + seen.append("responses") + await websocket.close(code=1000) + + async def _live_spy(websocket, *args, **kwargs): + seen.append("live") + await websocket.close(code=1000) + + monkeypatch.setattr(app.state.proxy, "handle_openai_responses_ws", _responses_spy) + monkeypatch.setattr(proxy_routes, "handle_codex_live_websocket", _live_spy) + return lambda: bool(seen) + + # ───────────────────────────── 3.1 security headers ───────────────────────