fix(security): realtime dropped the org header and put the API key in the URL

tests/test_realtime_auth.py

@@ -0,0 +1,108 @@
+"""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
+
+import pytest
+
+from wave_sdk.client import WaveClient
+from wave_sdk.realtime import RealtimeAPI
+
+
+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_sdk.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

wave_sdk/realtime.py

@@ -7,13 +7,19 @@
 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
 
 import contextlib
 import json
 from collections.abc import Iterator
 from typing import Any, Callable
+from urllib.parse import quote, urlencode
 
 import httpx
 
@@ -30,6 +36,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.
 
@@ -40,19 +55,40 @@ 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
             raise ImportError(
                 "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]:
@@ -97,34 +133,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},
         )