mirror of
https://github.com/pyMC-dev/pyMC_Repeater.git
synced 2026-11-03 11:52:22 +01:00
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:
1 parent
4705c99c38
commit
56e2fd6417
2 files changed
+75
-12
No files matched your search
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in new issue
Block a user