diff --git a/muxplex/main.py b/muxplex/main.py index 77895d5..5de4a29 100644 --- a/muxplex/main.py +++ b/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 diff --git a/muxplex/tests/test_ws_proxy.py b/muxplex/tests/test_ws_proxy.py index d9d6cf1..7e8ee9f 100644 --- a/muxplex/tests/test_ws_proxy.py +++ b/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 ', + 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 # ---------------------------------------------------------------------------