diff --git a/finbot/core/auth/session.py b/finbot/core/auth/session.py index ce08f070..a636f315 100644 --- a/finbot/core/auth/session.py +++ b/finbot/core/auth/session.py @@ -539,6 +539,44 @@ def get_session( if own_db: db.close() + def get_session_identity(self, session_id: str) -> tuple[str, str] | None: + """Resolve a session cookie to (namespace, user_id) without side effects. + + For callers that can't set a cookie on the response (WebSocket + handshakes). Unlike get_session, this never rotates, deletes or + touches the session, so the browser's cookie stays valid. + """ + if not session_id: + return None + db = SessionLocal() + try: + session = ( + db.query(UserSession) + .filter(UserSession.session_id == session_id) + .first() + ) + if not session or session.is_expired(): + return None + if not self._verify_session_signature( + session.session_data, session.signature + ): + return None + max_age = ( + settings.MAX_TEMP_SESSION_AGE + if session.is_temporary + else settings.MAX_PERM_SESSION_AGE + ) + created_at = ( + session.created_at + if session.created_at.tzinfo + else session.created_at.replace(tzinfo=UTC) + ) + if (datetime.now(UTC) - created_at).total_seconds() > max_age: + return None + return session.namespace, session.user_id + finally: + db.close() + def _rotate_session( self, old_context: SessionContext, db: Session ) -> SessionContext: diff --git a/finbot/core/websocket/routes.py b/finbot/core/websocket/routes.py index 4510c020..a8775f3c 100644 --- a/finbot/core/websocket/routes.py +++ b/finbot/core/websocket/routes.py @@ -3,10 +3,11 @@ import json import logging -from fastapi import APIRouter, HTTPException, Query, WebSocket, WebSocketDisconnect +from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect, status from pydantic import BaseModel from finbot.config import settings +from finbot.core.auth.session import session_manager from finbot.core.websocket.events import WSEvent, WSEventType from finbot.core.websocket.manager import get_ws_manager @@ -55,23 +56,27 @@ async def push_test_event(req: TestEventRequest): @router.websocket("/connect") -async def websocket_endpoint( - websocket: WebSocket, - namespace: str = Query(...), - user_id: str = Query(...), -): +async def websocket_endpoint(websocket: WebSocket): """ WebSocket endpoint for real-time updates. - Query params: - - namespace: User's namespace - - user_id: User's ID + The user is identified from the session cookie sent on the handshake. + Connections without a valid session are closed with 1008. Message format (JSON): - - {"action": "subscribe", "topic": "..."} + - {"action": "subscribe", "topic": "..."} (own activity topic only) - {"action": "unsubscribe", "topic": "..."} - {"action": "ping"} """ + identity = session_manager.get_session_identity( + websocket.cookies.get(settings.SESSION_COOKIE_NAME) + ) + if not identity: + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + return + namespace, user_id = identity + own_topic = f"activity:{namespace}:{user_id}" + manager = get_ws_manager() connection_id = await manager.connect(websocket, user_id, namespace) @@ -86,8 +91,16 @@ async def websocket_endpoint( if action == "subscribe": topic = message.get("topic") - if topic: + if topic == own_topic: await manager.subscribe(connection_id, topic) + elif topic: + await manager.send_to_connection( + connection_id, + WSEvent( + type=WSEventType.ERROR, + data={"message": "Not allowed to subscribe to this topic"}, + ), + ) elif action == "unsubscribe": topic = message.get("topic") diff --git a/finbot/static/js/vendor/ctf_sidecar.js b/finbot/static/js/vendor/ctf_sidecar.js index a525f699..2ff97fe1 100644 --- a/finbot/static/js/vendor/ctf_sidecar.js +++ b/finbot/static/js/vendor/ctf_sidecar.js @@ -276,12 +276,10 @@ class CTFSidecar { // WebSocket for live updates connectWebSocket() { - // Get namespace and user_id from session (you'll need to expose these) - const namespace = window.CTF_NAMESPACE || 'default'; - const userId = window.CTF_USER_ID || 'anonymous'; - + // The server identifies the user from the session cookie and + // auto-subscribes the connection to its activity topic. const wsProtocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; - const wsUrl = `${wsProtocol}//${window.location.host}/ws/connect?namespace=${namespace}&user_id=${userId}`; + const wsUrl = `${wsProtocol}//${window.location.host}/ws/connect`; try { this.ws = new WebSocket(wsUrl); @@ -289,12 +287,6 @@ class CTFSidecar { this.ws.onopen = () => { console.log('CTF Sidecar: WebSocket connected'); this.updateWsStatus('Connected', true); - - // Subscribe to activity topic (server auto-subscribes, but explicit is fine) - this.ws.send(JSON.stringify({ - action: 'subscribe', - topic: `activity:${namespace}:${userId}` - })); }; this.ws.onmessage = (event) => { diff --git a/tests/unit/auth/test_websocket_auth.py b/tests/unit/auth/test_websocket_auth.py new file mode 100644 index 00000000..8b28c500 --- /dev/null +++ b/tests/unit/auth/test_websocket_auth.py @@ -0,0 +1,127 @@ +""" +WebSocket /ws/connect authentication tests + +The handshake must identify the user from the session cookie, not from +client-supplied query params, and a connection may only subscribe to its +own activity topic. +""" + +from datetime import UTC, datetime, timedelta + +import pytest +from fastapi.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +from finbot.config import settings +from finbot.core.auth.session import session_manager +from finbot.core.data.models import UserSession +from finbot.core.websocket.manager import get_ws_manager + + +def _cookie_header(session_id: str) -> dict: + return {"cookie": f"{settings.SESSION_COOKIE_NAME}={session_id}"} + + +VICTIM_URL = "/ws/connect?namespace=ns_victim&user_id=user_victim" + + +def _assert_rejected(client: TestClient, url: str = VICTIM_URL, headers=None): + with pytest.raises(WebSocketDisconnect) as exc_info: + with client.websocket_connect(url, headers=headers or {}) as ws: + ws.receive_json() + assert exc_info.value.code == 1008 + + +@pytest.mark.unit +def test_ws_rejects_connection_without_session_cookie(client, db): + """No cookie: the handshake is refused even if IDs are passed in the URL.""" + _assert_rejected(client) + + +@pytest.mark.unit +def test_ws_rejects_unknown_session_cookie(client, db): + _assert_rejected(client, headers=_cookie_header("not-a-real-session")) + + +@pytest.mark.unit +def test_ws_rejects_expired_session(client, db): + ctx = session_manager.create_session(email="ws_expired@example.com") + row = db.query(UserSession).filter(UserSession.session_id == ctx.session_id).first() + row.expires_at = datetime.now(UTC) - timedelta(minutes=1) + db.commit() + + _assert_rejected(client, headers=_cookie_header(ctx.session_id)) + + +@pytest.mark.unit +def test_ws_rejects_tampered_session(client, db): + ctx = session_manager.create_session(email="ws_tampered@example.com") + row = db.query(UserSession).filter(UserSession.session_id == ctx.session_id).first() + row.signature = "0" * 64 + db.commit() + + _assert_rejected(client, headers=_cookie_header(ctx.session_id)) + + +@pytest.mark.unit +def test_ws_identity_comes_from_session_not_query_params(client, db): + """Query params pointing at another user are ignored.""" + ctx = session_manager.create_session(email="ws_owner@example.com") + manager = get_ws_manager() + + with client.websocket_connect(VICTIM_URL, headers=_cookie_header(ctx.session_id)) as ws: + assert ws.receive_json()["type"] == "connected" + assert manager.get_user_connection_count(ctx.namespace, ctx.user_id) == 1 + assert manager.get_user_connection_count("ns_victim", "user_victim") == 0 + + +@pytest.mark.unit +def test_ws_allows_subscribing_to_own_activity_topic(client, db): + ctx = session_manager.create_session(email="ws_own_topic@example.com") + own_topic = f"activity:{ctx.namespace}:{ctx.user_id}" + + with client.websocket_connect("/ws/connect", headers=_cookie_header(ctx.session_id)) as ws: + ws.receive_json() # connected + ws.send_json({"action": "subscribe", "topic": own_topic}) + msg = ws.receive_json() + assert msg["type"] == "subscribed" + assert msg["data"]["topic"] == own_topic + + +@pytest.mark.unit +def test_ws_rejects_subscribing_to_another_users_topic(client, db): + ctx = session_manager.create_session(email="ws_snoop@example.com") + other_topic = "activity:ns_someone_else:user_someone_else" + manager = get_ws_manager() + + with client.websocket_connect("/ws/connect", headers=_cookie_header(ctx.session_id)) as ws: + connection_id = ws.receive_json()["data"]["connection_id"] + ws.send_json({"action": "subscribe", "topic": other_topic}) + msg = ws.receive_json() + assert msg["type"] == "error" + assert other_topic not in manager._connections[connection_id].subscriptions + + +@pytest.mark.unit +def test_get_session_identity_does_not_rotate_or_touch_session(db): + """A session due for rotation must keep its ID, since a WebSocket + handshake can't hand the browser a new cookie.""" + ctx = session_manager.create_session(email="ws_rotation@example.com") + row = db.query(UserSession).filter(UserSession.session_id == ctx.session_id).first() + row.last_rotation = datetime.now(UTC) - timedelta(days=1) + db.commit() + expires_before = row.expires_at + + identity = session_manager.get_session_identity(ctx.session_id) + + assert identity == (ctx.namespace, ctx.user_id) + db.expire_all() + rows = db.query(UserSession).filter(UserSession.user_id == ctx.user_id).all() + assert [r.session_id for r in rows] == [ctx.session_id] + assert rows[0].expires_at == expires_before + + +@pytest.mark.unit +def test_get_session_identity_handles_missing_cookie(db): + assert session_manager.get_session_identity(None) is None + assert session_manager.get_session_identity("") is None