diff --git a/tests/test_realtime_auth.py b/tests/test_realtime_auth.py new file mode 100644 index 0000000..50dab79 --- /dev/null +++ b/tests/test_realtime_auth.py @@ -0,0 +1,107 @@ +"""Realtime auth-transport tests. + +Covers two defects that were live on main: + +1. ``RealtimeAPI`` built its own header dict and omitted ``X-Organization-Id``, so every realtime + operation ran unscoped while the rest of the SDK carried the tenant. +2. The API key travelled in the WebSocket URL query string, where every hop that logs a request + line records it. +""" +from __future__ import annotations + +import sys +import types +from wave.client import WaveClient +from wave.realtime import RealtimeAPI + +import pytest + + +class _FakeSocket: + def __init__(self, url: str, header: list[str] | None = None, **_: object) -> None: + self.url = url + self.header = header or [] + + +@pytest.fixture +def captured_ws(monkeypatch): + """Stub the optional ``websocket-client`` dep and capture the upgrade it would perform.""" + calls: list[_FakeSocket] = [] + + def create_connection(url: str, header: list[str] | None = None, **kwargs: object) -> _FakeSocket: + sock = _FakeSocket(url, header, **kwargs) + calls.append(sock) + return sock + + module = types.ModuleType("websocket") + module.create_connection = create_connection # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "websocket", module) + return calls + + +def _api(**client_kwargs) -> RealtimeAPI: + client = WaveClient(api_key="sk-test-key", **client_kwargs) + return RealtimeAPI(client) + + +# --- Finding 1: multi-tenant isolation ------------------------------------------------------- + + +def test_rest_headers_carry_the_organization(): + headers = _api(organization_id="org_123")._headers() + assert headers["X-Organization-Id"] == "org_123" + assert headers["Authorization"] == "Bearer sk-test-key" + + +def test_rest_headers_omit_the_organization_when_unset(): + assert "X-Organization-Id" not in _api()._headers() + + +def test_ws_upgrade_carries_the_organization(captured_ws): + _api(organization_id="org_123").connect("stream:abc") + assert "X-Organization-Id: org_123" in captured_ws[0].header + + +def test_rest_path_encodes_the_channel(): + from wave.realtime import _channel_path + + assert _channel_path("stream:abc") == "stream:abc" + assert _channel_path("a/../b") == "a%2F..%2Fb" + + +# --- Finding 2: credential in the URL -------------------------------------------------------- + + +def test_api_key_is_not_in_the_connect_url(captured_ws): + _api(organization_id="org_123").connect("stream:abc") + assert "sk-test-key" not in captured_ws[0].url + assert "access_token" not in captured_ws[0].url + + +def test_api_key_travels_in_the_upgrade_header(captured_ws): + _api().connect("stream:abc") + assert "Authorization: Bearer sk-test-key" in captured_ws[0].header + + +def test_legacy_query_token_is_opt_in(captured_ws): + client = WaveClient(api_key="sk-test-key") + RealtimeAPI(client, token_in_query=True).connect("stream:abc") + assert "access_token=sk-test-key" in captured_ws[0].url + # The header is still sent — opting into the legacy param does not disable the correct path. + assert "Authorization: Bearer sk-test-key" in captured_ws[0].header + + +# --- Finding 3: query-parameter injection ---------------------------------------------------- + + +def test_channel_cannot_inject_a_query_parameter(captured_ws): + _api().connect("stream:abc&as=victim") + url = captured_ws[0].url + assert "&as=victim" not in url + assert "as%3Dvictim" in url + + +def test_as_parameter_is_encoded(captured_ws): + _api().connect("stream:abc", as_="user&admin=1") + url = captured_ws[0].url + assert "&admin=1" not in url diff --git a/wave/realtime.py b/wave/realtime.py index 3f3de2d..152d630 100644 --- a/wave/realtime.py +++ b/wave/realtime.py @@ -7,6 +7,11 @@ WebSocket support uses the optional ``websocket-client`` package: ``pip install 'wave-sdk[realtime]'``. Auth, scope, entitlement, and metering are enforced server-side (the gateway, via realtime's /v1/verify federation) — the SDK only forwards your API key. + +Credentials travel in the ``Authorization`` header on both the REST calls and the WebSocket upgrade. +``websocket-client`` sets arbitrary upgrade headers, so the browser constraint that forces a +``?access_token=`` query parameter does not apply to this client. See ``token_in_query`` on +:class:`RealtimeAPI` for the legacy escape hatch. """ from __future__ import annotations @@ -14,6 +19,7 @@ import json from collections.abc import Iterator from typing import Any, Callable +from urllib.parse import quote, urlencode from wave.client import WaveClient import httpx @@ -29,6 +35,15 @@ def _http_origin(ws_url: str) -> str: return base +def _channel_path(channel: str) -> str: + """Percent-encode a channel for use as a single REST path segment. + + ``:`` stays literal because WAVE channel names are ``stream:abc`` shaped; everything else that + could leave the segment (``/``, ``?``, ``#``, ``&``) is encoded. + """ + return quote(channel, safe=":") + + class RealtimeChannel: """One subscribed channel over a WebSocket. @@ -39,7 +54,15 @@ class RealtimeChannel: ch.run() # blocks, dispatching frames """ - def __init__(self, channel: str, api_key: str, ws_base: str = _DEFAULT_WS, as_: str | None = None): + def __init__( + self, + channel: str, + api_key: str, + ws_base: str = _DEFAULT_WS, + as_: str | None = None, + organization_id: str | None = None, + token_in_query: bool = False, + ): try: import websocket # websocket-client (optional dep) except ImportError as e: # pragma: no cover - import guard @@ -47,11 +70,24 @@ def __init__(self, channel: str, api_key: str, ws_base: str = _DEFAULT_WS, as_: "WAVE realtime requires the 'websocket-client' package: pip install 'wave-sdk[realtime]'" ) from e self.channel = channel - # Browser/SDK clients can't set headers on the WS upgrade → key travels as a query param (wss). - url = f"{ws_base.rstrip('/')}/v1/connect?channel={channel}&access_token={api_key}" + + # Every value is urlencoded: a channel containing '&' or '#' would otherwise inject or + # truncate query parameters on the upgrade. + params: dict[str, str] = {"channel": channel} if as_: - url += f"&as={as_}" - self._ws = websocket.create_connection(url) + params["as"] = as_ + + headers = [f"Authorization: Bearer {api_key}"] + if organization_id: + headers.append(f"X-Organization-Id: {organization_id}") + + if token_in_query: + # Legacy form for deployments that cannot read the upgrade header. The key lands in + # proxy logs, edge access logs, and shell history — opt in deliberately or not at all. + params["access_token"] = api_key + + url = f"{ws_base.rstrip('/')}/v1/connect?{urlencode(params)}" + self._ws = websocket.create_connection(url, header=headers) self._handlers: dict[str, list[Callable[[Any], None]]] = {} def __iter__(self) -> Iterator[dict]: @@ -96,34 +132,56 @@ def close(self) -> None: class RealtimeAPI: """Realtime entry point. ``wave.realtime.connect('stream:abc')`` for WS; ``publish/presence/history`` - are one-shot REST calls for producers that don't hold a socket.""" + are one-shot REST calls for producers that don't hold a socket. + + ``token_in_query`` re-enables the legacy ``?access_token=`` upgrade parameter for a deployment + that cannot read the ``Authorization`` header. It is off by default because a credential in a URL + is recorded by every hop that logs the request line. + """ - def __init__(self, client: WaveClient, url: str = _DEFAULT_WS): + def __init__(self, client: WaveClient, url: str = _DEFAULT_WS, token_in_query: bool = False): self._api_key = client.api_key + # Multi-tenant isolation: WaveClient stamps X-Organization-Id on every other surface, so + # realtime carries it too — on the REST calls and on the WS upgrade. + self._organization_id = client.organization_id self._ws_base = url.rstrip("/") self._http_base = _http_origin(self._ws_base) + self._token_in_query = token_in_query def connect(self, channel: str, as_: str | None = None) -> RealtimeChannel: - return RealtimeChannel(channel, self._api_key, self._ws_base, as_) + return RealtimeChannel( + channel, + self._api_key, + self._ws_base, + as_, + organization_id=self._organization_id, + token_in_query=self._token_in_query, + ) def _headers(self) -> dict[str, str]: - return {"Authorization": f"Bearer {self._api_key}", "content-type": "application/json"} + headers = {"Authorization": f"Bearer {self._api_key}", "content-type": "application/json"} + if self._organization_id: + headers["X-Organization-Id"] = self._organization_id + return headers def publish(self, channel: str, event: str, data: Any = None) -> dict: r = httpx.post( - f"{self._http_base}/v1/channels/{channel}/publish", + f"{self._http_base}/v1/channels/{_channel_path(channel)}/publish", headers=self._headers(), json={"event": event, "data": data}, ) return r.json() def presence(self, channel: str) -> dict: - r = httpx.get(f"{self._http_base}/v1/channels/{channel}/presence", headers=self._headers()) + r = httpx.get( + f"{self._http_base}/v1/channels/{_channel_path(channel)}/presence", + headers=self._headers(), + ) return r.json() def history(self, channel: str, limit: int = 50) -> dict: r = httpx.get( - f"{self._http_base}/v1/channels/{channel}/history", + f"{self._http_base}/v1/channels/{_channel_path(channel)}/history", headers=self._headers(), params={"limit": limit}, )