feat(ws): accept API tokens on /ws/companion_frame

The companion frame proxy only accepted an operator JWT in `?token=`,
while `/ws/packets` also accepts a long-lived API token (X-API-Key
header, or the same query parameter). Apply the identical rule to the
frame proxy so a client that already holds a repeater API key can open
the companion link on the HTTP port without keeping a JWT session alive.

No behaviour change for JWT clients. Same fallback order, close codes
and log lines as PacketWebSocket.
This commit is contained in:
Treehouse-00 committed 2026-09-10 10:22:20 -07:00
1 parent 4705c99c38
commit 56e2fd6417
2 files changed
+75 -12

No files matched your search

+26 -11
View File
@@ -27,8 +27,9 @@ def set_daemon(instance):
class CompanionFrameWebSocket(WebSocket):
def opened(self):
"""Authenticate, resolve companion, open TCP socket, start reader."""
# JWT auth — same pattern as PacketWebSocket
# JWT or API-token auth — same pattern as PacketWebSocket
jwt_handler = cherrypy.config.get("jwt_handler")
token_manager = cherrypy.config.get("token_manager")
qs = ""
if hasattr(self, "environ"):
@@ -38,24 +39,39 @@ class CompanionFrameWebSocket(WebSocket):
token = params.get("token", [None])[0]
companion_name = params.get("companion_name", [None])[0]
api_key = self.environ.get("HTTP_X_API_KEY", "") if hasattr(self, "environ") else ""
if not jwt_handler:
logger.warning("Connection rejected: no JWT handler configured")
self.close(code=1011, reason="server configuration error")
return
if not token:
if not token and not api_key:
logger.warning("Connection rejected: missing token")
self.close(code=1008, reason="unauthorized")
return
try:
payload = jwt_handler.verify_jwt(token)
if not payload:
logger.warning("Connection rejected: invalid token")
self.close(code=1008, reason="unauthorized")
return
except Exception as e:
logger.warning(f"Auth error: {e}")
user = None
if token:
try:
payload = jwt_handler.verify_jwt(token)
if payload:
user = payload.get("sub", "unknown")
except Exception as e:
logger.warning(f"Auth error: {e}")
if user is None:
api_token = api_key or token
if api_token and token_manager:
try:
token_info = token_manager.verify_token(api_token)
if token_info:
user = f"api_token:{token_info.get('name', 'unknown')}"
except Exception as e:
logger.warning(f"API key auth error: {e}")
if user is None:
logger.warning("Connection rejected: invalid token")
self.close(code=1008, reason="unauthorized")
return
@@ -93,7 +109,6 @@ class CompanionFrameWebSocket(WebSocket):
)
self._reader.start()
user = payload.get("sub", "unknown")
logger.info(
f"Companion WS opened: user={user}, companion={companion_name}, tcp={tcp_host}:{tcp_port}"
)
+49 -1
View File
@@ -14,9 +14,11 @@ def cp_cfg(monkeypatch):
return cfg
def _ws(query_string):
def _ws(query_string, api_key=None):
ws = object.__new__(proxy.CompanionFrameWebSocket)
ws.environ = {"QUERY_STRING": query_string}
if api_key is not None:
ws.environ["HTTP_X_API_KEY"] = api_key
ws.close = MagicMock()
ws.send = MagicMock()
ws._teardown = MagicMock()
@@ -43,6 +45,52 @@ def test_opened_rejects_invalid_token(cp_cfg):
ws.close.assert_called_once_with(code=1008, reason="unauthorized")
def test_opened_accepts_api_key_header(cp_cfg, monkeypatch):
cp_cfg["jwt_handler"] = SimpleNamespace(verify_jwt=lambda _t: None)
cp_cfg["token_manager"] = SimpleNamespace(
verify_token=lambda t: {"name": "waev"} if t == "k" else None
)
ws = _ws("companion_name=c1", api_key="k")
ws._resolve_tcp_endpoint = MagicMock(return_value=("127.0.0.1", 5000))
fake_socket = MagicMock()
monkeypatch.setattr(proxy.socket, "socket", lambda *_args, **_kwargs: fake_socket)
monkeypatch.setattr(proxy.threading, "Thread", lambda **_kw: MagicMock())
ws.opened()
ws.close.assert_not_called()
assert ws._companion_name == "c1"
def test_opened_accepts_api_token_in_query_like_packet_ws(cp_cfg, monkeypatch):
cp_cfg["jwt_handler"] = SimpleNamespace(verify_jwt=lambda _t: None)
cp_cfg["token_manager"] = SimpleNamespace(verify_token=lambda _t: {"name": "waev"})
ws = _ws("token=k&companion_name=c1")
ws._resolve_tcp_endpoint = MagicMock(return_value=("127.0.0.1", 5000))
fake_socket = MagicMock()
monkeypatch.setattr(proxy.socket, "socket", lambda *_args, **_kwargs: fake_socket)
monkeypatch.setattr(proxy.threading, "Thread", lambda **_kw: MagicMock())
ws.opened()
ws.close.assert_not_called()
def test_opened_rejects_unknown_api_key(cp_cfg):
cp_cfg["jwt_handler"] = SimpleNamespace(verify_jwt=lambda _t: None)
cp_cfg["token_manager"] = SimpleNamespace(verify_token=lambda _t: None)
ws = _ws("companion_name=c1", api_key="bad")
ws.opened()
ws.close.assert_called_once_with(code=1008, reason="unauthorized")
def test_opened_rejects_api_key_without_token_manager(cp_cfg):
cp_cfg["jwt_handler"] = SimpleNamespace(verify_jwt=lambda _t: None)
ws = _ws("companion_name=c1", api_key="k")
ws.opened()
ws.close.assert_called_once_with(code=1008, reason="unauthorized")
def test_opened_rejects_missing_companion_name(cp_cfg):
cp_cfg["jwt_handler"] = SimpleNamespace(verify_jwt=lambda _t: {"sub": "u"})
ws = _ws("token=t")