fix(websocket): harden responses websocket origin handling (#1481)

## Description

Validate browser WebSocket origins before accepting WS sessions.


## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [x] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- validate Responses WebSocket `Origin` before routing the session
upstream
- keep native clients that omit `Origin` working
- allow loopback origins by default and support explicit origins via
`HEADROOM_WS_ORIGINS`


## Testing


- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
.venv/bin/python -m pytest tests/test_openai_codex_routing.py
Result: 19 passed
.venv/bin/python -m ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_routing.py
Result: All checks passed
```

## Real Behavior Proof

- Environment: macOS
- Exact command / steps: `venv/bin/python -m pytest
tests/test_openai_codex_routing.py``.venv/bin/python -m ruff check
headroom/proxy/handlers/openai.py tests/test_openai_codex_routing.py`
- Observed result:`19 passed` `All checks passed`
- Not tested: NA

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
This commit is contained in:
Devanshi Vyas 2026-06-26 17:06:07 -07:00 committed by GitHub
parent 9f772378d3
commit c632023cc1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 147 additions and 1 deletions

View file

@ -29,6 +29,7 @@ from headroom.proxy.helpers import (
extract_tags, extract_tags,
jitter_delay_ms, jitter_delay_ms,
) )
from headroom.proxy.loopback_guard import is_loopback_host
from headroom.proxy.stage_timer import StageTimer, emit_stage_timings_log from headroom.proxy.stage_timer import StageTimer, emit_stage_timings_log
from headroom.proxy.ws_session_registry import ( from headroom.proxy.ws_session_registry import (
TerminationCause, TerminationCause,
@ -65,6 +66,81 @@ _OPENAI_RESPONSES_UNIT_PARALLELISM_MAX = 16
_OPENAI_RESPONSES_UNIT_CACHE_INIT_LOCK = threading.RLock() _OPENAI_RESPONSES_UNIT_CACHE_INIT_LOCK = threading.RLock()
_OPENAI_RESPONSES_UNIT_EXECUTOR_LOCK = threading.RLock() _OPENAI_RESPONSES_UNIT_EXECUTOR_LOCK = threading.RLock()
_OPENAI_RESPONSES_UNIT_EXECUTOR: ThreadPoolExecutor | None = None _OPENAI_RESPONSES_UNIT_EXECUTOR: ThreadPoolExecutor | None = None
_WS_ALLOWED_ORIGINS_ENV = "HEADROOM_WS_ORIGINS"
_CORS_ALLOWED_ORIGINS_ENV = "HEADROOM_CORS_ORIGINS"
def _header_get(headers: dict[str, str], name: str) -> str | None:
"""Case-insensitive header lookup for plain dicts."""
lowered = name.lower()
for key, value in headers.items():
if key.lower() == lowered:
return value
return None
def _normalize_origin(origin: str) -> str | None:
parsed = urlparse(origin.strip())
if not parsed.scheme or not parsed.hostname:
return None
scheme = parsed.scheme.lower()
hostname = parsed.hostname.lower()
if scheme not in {"http", "https", "ws", "wss"}:
return None
port = parsed.port
default_port = (scheme in {"http", "ws"} and port == 80) or (
scheme in {"https", "wss"} and port == 443
)
port_part = "" if port is None or default_port else f":{port}"
return f"{scheme}://{hostname}{port_part}"
def _allowed_ws_origins_from_env() -> list[str] | None:
raw = os.environ.get(_WS_ALLOWED_ORIGINS_ENV)
if raw is None or not raw.strip():
raw = os.environ.get(_CORS_ALLOWED_ORIGINS_ENV)
if raw is None or not raw.strip():
return None
return [origin.strip() for origin in raw.split(",") if origin.strip()]
def _is_loopback_ws_origin(origin: str) -> bool:
parsed = urlparse(origin.strip())
if parsed.scheme.lower() not in {"http", "https", "ws", "wss"}:
return False
if parsed.hostname is None:
return False
return is_loopback_host(parsed.hostname)
def _is_allowed_websocket_origin(headers: dict[str, str]) -> bool:
"""Return True when the WebSocket Origin matches the configured policy.
Native clients commonly omit Origin, so absence is allowed. When Origin is
present, default to loopback-only and allow explicit configured origins via
HEADROOM_WS_ORIGINS or HEADROOM_CORS_ORIGINS.
"""
origin = _header_get(headers, "origin")
if not origin:
return True
allowed_origins = _allowed_ws_origins_from_env()
if allowed_origins is None:
return _is_loopback_ws_origin(origin)
if "*" in allowed_origins:
return True
normalized_origin = _normalize_origin(origin)
if normalized_origin is None:
return False
normalized_allowed = {
normalized
for allowed in allowed_origins
for normalized in (_normalize_origin(allowed),)
if normalized is not None
}
return normalized_origin in normalized_allowed
def _usage_int(value: Any) -> int: def _usage_int(value: Any) -> int:
@ -3577,6 +3653,16 @@ class OpenAIHandlerMixin:
_ws_path = getattr(_ws_url_obj, "path", "") if _ws_url_obj is not None else "" _ws_path = getattr(_ws_url_obj, "path", "") if _ws_url_obj is not None else ""
if not _ws_path: if not _ws_path:
_ws_path = "/v1/responses" _ws_path = "/v1/responses"
if not _is_allowed_websocket_origin(ws_headers):
logger.warning(
"event=websocket_origin_not_allowed request_id=%s session_id=%s path=%s origin=%r",
request_id,
session_id,
_ws_path,
_header_get(ws_headers, "origin"),
)
await websocket.close(code=1008, reason="origin not allowed")
return
# WS sessions bypass the HTTP middleware that stamps X-Client: codex on # WS sessions bypass the HTTP middleware that stamps X-Client: codex on
# the Responses endpoint, so apply the same path-based stamp here before # the Responses endpoint, so apply the same path-based stamp here before
# classify_client runs (parallels server.py / should_stamp_codex_client). # classify_client runs (parallels server.py / should_stamp_codex_client).

View file

@ -11,6 +11,7 @@ from fastapi import Request
from headroom.proxy.handlers.openai import ( from headroom.proxy.handlers.openai import (
OpenAIHandlerMixin, OpenAIHandlerMixin,
_is_allowed_websocket_origin,
_openai_responses_unit_cache_key, _openai_responses_unit_cache_key,
_resolve_codex_routing_headers, _resolve_codex_routing_headers,
) )
@ -443,10 +444,49 @@ class _DummyWebSocket:
def __init__(self, headers: dict[str, str]): def __init__(self, headers: dict[str, str]):
self.headers = headers self.headers = headers
self.accepted_subprotocol = None self.accepted_subprotocol = None
self.closed = False
self.close_code = None
self.close_reason = None
async def accept(self, subprotocol=None): async def accept(self, subprotocol=None, headers=None):
self.accepted_subprotocol = subprotocol self.accepted_subprotocol = subprotocol
async def close(self, code=1000, reason=None):
self.closed = True
self.close_code = code
self.close_reason = reason
def test_websocket_origin_policy_allows_native_clients_without_origin(monkeypatch):
monkeypatch.delenv("HEADROOM_WS_ORIGINS", raising=False)
monkeypatch.delenv("HEADROOM_CORS_ORIGINS", raising=False)
assert _is_allowed_websocket_origin({"authorization": "Bearer token"}) is True
def test_websocket_origin_policy_allows_loopback_origins_by_default(monkeypatch):
monkeypatch.delenv("HEADROOM_WS_ORIGINS", raising=False)
monkeypatch.delenv("HEADROOM_CORS_ORIGINS", raising=False)
assert _is_allowed_websocket_origin({"origin": "http://localhost:3000"}) is True
assert _is_allowed_websocket_origin({"origin": "https://127.0.0.1:8787"}) is True
def test_websocket_origin_policy_requires_config_for_remote_origins(monkeypatch):
monkeypatch.delenv("HEADROOM_WS_ORIGINS", raising=False)
monkeypatch.delenv("HEADROOM_CORS_ORIGINS", raising=False)
assert _is_allowed_websocket_origin({"origin": "https://remote.example"}) is False
assert _is_allowed_websocket_origin({"origin": "http://"}) is False
def test_websocket_origin_policy_can_be_pinned_with_env(monkeypatch):
monkeypatch.setenv("HEADROOM_WS_ORIGINS", "https://dash.example.com")
monkeypatch.delenv("HEADROOM_CORS_ORIGINS", raising=False)
assert _is_allowed_websocket_origin({"origin": "https://dash.example.com"}) is True
assert _is_allowed_websocket_origin({"origin": "http://localhost:3000"}) is False
def test_handle_openai_responses_ws_resolves_codex_routing_headers(): def test_handle_openai_responses_ws_resolves_codex_routing_headers():
class SentinelError(RuntimeError): class SentinelError(RuntimeError):
@ -462,3 +502,23 @@ def test_handle_openai_responses_ws_resolves_codex_routing_headers():
): ):
with pytest.raises(SentinelError, match="resolved"): with pytest.raises(SentinelError, match="resolved"):
anyio.run(handler.handle_openai_responses_ws, websocket) anyio.run(handler.handle_openai_responses_ws, websocket)
def test_handle_openai_responses_ws_closes_unconfigured_origin(monkeypatch):
handler = _DummyOpenAIHandler()
websocket = _DummyWebSocket({"origin": "https://remote.example"})
monkeypatch.delenv("HEADROOM_WS_ORIGINS", raising=False)
monkeypatch.delenv("HEADROOM_CORS_ORIGINS", raising=False)
with patch.dict(sys.modules, {"websockets": MagicMock()}):
with patch(
"headroom.proxy.handlers.openai._resolve_codex_routing_headers",
side_effect=AssertionError("routing should not run"),
):
anyio.run(handler.handle_openai_responses_ws, websocket)
assert websocket.closed is True
assert websocket.close_code == 1008
assert websocket.close_reason == "origin not allowed"
assert websocket.accepted_subprotocol is None