diff --git a/backend/app/api/routes/status.py b/backend/app/api/routes/status.py index 106af8c..4f1d012 100644 --- a/backend/app/api/routes/status.py +++ b/backend/app/api/routes/status.py @@ -11,11 +11,23 @@ _connections: list[WebSocket] = [] @router.websocket("/ws/status") -async def ws_status(websocket: WebSocket, token: str | None = None) -> None: - if not token or not decode_token(token): - await websocket.close(code=1008) # Policy Violation - return +async def ws_status(websocket: WebSocket) -> None: + # Accept first so we can send a close frame with a reason code await websocket.accept() + try: + # Expect the first message to be a JSON auth payload: {"token": ""} + raw = await websocket.receive_text() + try: + payload = json.loads(raw) + token = payload.get("token", "") + except (json.JSONDecodeError, AttributeError): + token = "" + if not token or not decode_token(token): + await websocket.close(code=1008) # Policy Violation + return + except WebSocketDisconnect: + return + _connections.append(websocket) try: while True: diff --git a/backend/tests/test_status.py b/backend/tests/test_status.py index aa0d148..b9dfcba 100644 --- a/backend/tests/test_status.py +++ b/backend/tests/test_status.py @@ -22,24 +22,33 @@ def _make_token() -> str: # --------------------------------------------------------------------------- def test_websocket_rejected_without_token(): - """Connection with no token must be closed before being accepted.""" - with TestClient(app) as client, pytest.raises(WebSocketDisconnect), client.websocket_connect("/api/v1/status/ws/status"): - pass + """Connection that sends no token field must be closed with 1008.""" + with TestClient(app) as client, pytest.raises(WebSocketDisconnect), client.websocket_connect("/api/v1/status/ws/status") as ws: + ws.send_text(json.dumps({})) # missing token field + ws.receive_text() # triggers WebSocketDisconnect from server close def test_websocket_rejected_with_invalid_token(): - """Connection with a garbage token must be closed.""" - with TestClient(app) as client, pytest.raises(WebSocketDisconnect), client.websocket_connect("/api/v1/status/ws/status?token=not-a-valid-jwt"): - pass + """Connection that sends a garbage token must be closed.""" + with TestClient(app) as client, pytest.raises(WebSocketDisconnect), client.websocket_connect("/api/v1/status/ws/status") as ws: + ws.send_text(json.dumps({"token": "not-a-valid-jwt"})) + ws.receive_text() + + +def test_websocket_rejected_with_malformed_json(): + """Connection that sends non-JSON as auth must be closed.""" + with TestClient(app) as client, pytest.raises(WebSocketDisconnect), client.websocket_connect("/api/v1/status/ws/status") as ws: + ws.send_text("not-json") + ws.receive_text() def test_websocket_accepted_with_valid_token(): - """Connection with a valid JWT must be accepted and kept open.""" + """Connection that sends a valid JWT as first message must be accepted.""" token = _make_token() - with TestClient(app) as client, client.websocket_connect(f"/api/v1/status/ws/status?token={token}") as ws: - # Connection is open — we can send a ping and it should not raise + with TestClient(app) as client, client.websocket_connect("/api/v1/status/ws/status") as ws: + ws.send_text(json.dumps({"token": token})) + # Connection is open — subsequent messages should not raise ws.send_text("ping") - # Server keeps the connection open (no disconnect expected) # --------------------------------------------------------------------------- diff --git a/frontend/src/hooks/useStatusPolling.ts b/frontend/src/hooks/useStatusPolling.ts index 5d5cbf7..192a005 100644 --- a/frontend/src/hooks/useStatusPolling.ts +++ b/frontend/src/hooks/useStatusPolling.ts @@ -24,11 +24,16 @@ export function useStatusPolling() { const protocol = window.location.protocol === 'https:' ? 'wss' : 'ws' const host = window.location.host // includes port when non-standard - const url = `${protocol}://${host}/api/v1/status/ws/status?token=${encodeURIComponent(token)}` + const url = `${protocol}://${host}/api/v1/status/ws/status` const ws = new WebSocket(url) wsRef.current = ws + // Send token as first message (not in URL to avoid log/history exposure) + ws.onopen = () => { + ws.send(JSON.stringify({ token })) + } + ws.onmessage = (event) => { try { const msg: StatusMessage = JSON.parse(event.data)