feat: add Bearer token auth to WebSocket proxy for federation

muxplex/main.py

@@ -607,9 +607,16 @@ async def terminal_ws_proxy(websocket: WebSocket) -> None:
     host = websocket.client.host if websocket.client else ""
     if host not in ("127.0.0.1", "::1"):
         session_cookie = websocket.cookies.get("muxplex_session")
-        if not session_cookie or not verify_session_cookie(
+        cookie_ok = session_cookie and verify_session_cookie(
             _auth_secret, session_cookie, _auth_ttl
-        ):
+        )
+        bearer_ok = False
+        if _federation_key:
+            auth_header = websocket.headers.get("authorization", "")
+            if auth_header.lower().startswith("bearer "):
+                import hmac
+                bearer_ok = hmac.compare_digest(auth_header[7:], _federation_key)
+        if not cookie_ok and not bearer_ok:
             await websocket.close(code=4001)
             return
 

muxplex/tests/test_ws_proxy.py

@@ -247,6 +247,34 @@ def test_ws_auth_rejection_invalid_cookie():
     assert exc_info.value.code == 4001
 
 
+# ---------------------------------------------------------------------------
+# Test: Bearer token auth accepted
+# ---------------------------------------------------------------------------
+
+
+def test_ws_bearer_auth_accepted(monkeypatch):
+    """WebSocket from non-localhost with valid Bearer federation key is NOT rejected with 4001.
+
+    When a valid federation key is provided as 'Authorization: Bearer <key>',
+    the WebSocket connection must be accepted (not closed with code 4001).
+    """
+    fed_key = "test-federation-secret-key"
+    monkeypatch.setattr("muxplex.main._federation_key", fed_key)
+
+    fake_ws = FakeTtydWs(responses=[])
+    monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws)
+
+    # Connect without a session cookie but with a valid Bearer token.
+    # Should NOT raise WebSocketDisconnect with code 4001.
+    with TestClient(app) as c:
+        # If Bearer auth is not implemented, this raises WebSocketDisconnect(code=4001)
+        with c.websocket_connect(
+            "/terminal/ws",
+            headers={"Authorization": f"Bearer {fed_key}"},
+        ) as _ws:
+            pass  # Successfully connected — auth was accepted
+
+
 # ---------------------------------------------------------------------------
 # Tests 4–5: browser → ttyd relay
 # ---------------------------------------------------------------------------