diff --git a/repeater/web/companion_ws_proxy.py b/repeater/web/companion_ws_proxy.py index e12218e..ddf9c11 100644 --- a/repeater/web/companion_ws_proxy.py +++ b/repeater/web/companion_ws_proxy.py @@ -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}" ) diff --git a/tests/test_companion_ws_proxy.py b/tests/test_companion_ws_proxy.py index 3bfd6db..f021f45 100644 --- a/tests/test_companion_ws_proxy.py +++ b/tests/test_companion_ws_proxy.py @@ -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")