fix(security): first-message WebSocket auth to prevent token leakage

Debug Agent · claude · Debug Agent · commit 1e8a062c338b · 2026-04-09T19:01:38.000-03:00
Clients can now send {"type": "auth", "session_api_key": "..."} as the
first WebSocket frame instead of passing the token as a query parameter.
This keeps sk-oh-* tokens out of reverse-proxy / load-balancer access
logs (Traefik, Datadog, etc.).

Query param and header auth are preserved as deprecated fallbacks for
backwards compatibility.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/openhands-agent-server/openhands/agent_server/sockets.py b/openhands-agent-server/openhands/agent_server/sockets.py
@@ -2,11 +2,18 @@
 WebSocket endpoints for OpenHands SDK.
 
 These endpoints are separate from the main API routes to handle WebSocket-specific
-authentication. Browsers cannot send custom HTTP headers directly with WebSocket
-connections, so we support the `session_api_key` query param. For non-browser
-clients (e.g. Python/Node), we also support authenticating via headers.
+authentication.  Three auth methods are supported (highest to lowest precedence):
+
+1. **First-message auth** (recommended): The client sends
+   ``{"type": "auth", "session_api_key": "..."}`` as the very first WebSocket
+   frame after the connection opens.  This keeps tokens out of URLs and
+   therefore out of reverse-proxy / load-balancer access logs.
+2. Query parameter ``session_api_key`` — deprecated, kept for backwards compat.
+3. ``X-Session-API-Key`` header — for non-browser clients.
 """
 
+import asyncio
+import json
 import logging
 from dataclasses import dataclass
 from datetime import datetime
@@ -78,19 +85,68 @@ def _resolve_websocket_session_api_key(
     return None
 
 
+_FIRST_MESSAGE_AUTH_TIMEOUT_SECONDS = 10
+
+
 async def _accept_authenticated_websocket(
     websocket: WebSocket,
     session_api_key: str | None,
 ) -> bool:
-    """Authenticate and accept the socket, or close with an auth error."""
+    """Authenticate and accept the socket, or close with an auth error.
+
+    Authentication is attempted in the following order:
+
+    1. Query parameter / header (legacy, deprecated).
+    2. First-message auth — the client sends
+       ``{"type": "auth", "session_api_key": "..."}`` as the first frame.
+
+    The WebSocket is always *accepted* before first-message auth is attempted
+    because raw WebSocket requires ``accept()`` before any frames can be read.
+    """
     config = _get_config(websocket)
     resolved_key = _resolve_websocket_session_api_key(websocket, session_api_key)
-    if config.session_api_keys and resolved_key not in config.session_api_keys:
-        logger.warning("WebSocket authentication failed: invalid or missing API key")
+
+    # No auth configured — accept unconditionally.
+    if not config.session_api_keys:
+        await websocket.accept()
+        return True
+
+    # Legacy path: key supplied via query param or header.
+    if resolved_key is not None:
+        if resolved_key in config.session_api_keys:
+            logger.warning(
+                "session_api_key passed via query param or header is deprecated. "
+                "Use first-message auth instead."
+            )
+            await websocket.accept()
+            return True
+        logger.warning("WebSocket authentication failed: invalid API key")
         await websocket.close(code=4001, reason="Authentication failed")
         return False
+
+    # First-message auth: accept the connection, then read the first frame.
     await websocket.accept()
-    return True
+    try:
+        raw = await asyncio.wait_for(
+            websocket.receive_text(),
+            timeout=_FIRST_MESSAGE_AUTH_TIMEOUT_SECONDS,
+        )
+        data = json.loads(raw)
+    except (asyncio.TimeoutError, json.JSONDecodeError, WebSocketDisconnect):
+        logger.warning("WebSocket first-message auth failed: bad or missing payload")
+        await _safe_close_websocket(websocket, code=4001, reason="Authentication failed")
+        return False
+
+    if (
+        isinstance(data, dict)
+        and data.get("type") == "auth"
+        and data.get("session_api_key") in config.session_api_keys
+    ):
+        return True
+
+    logger.warning("WebSocket first-message auth failed: invalid key or payload")
+    await _safe_close_websocket(websocket, code=4001, reason="Authentication failed")
+    return False
 
 
 @sockets_router.websocket("/events/{conversation_id}")
@@ -329,9 +385,13 @@ async def _send_event(event: Event, websocket: WebSocket):
         logger.exception("error_sending_event: %r", event, stack_info=True)
 
 
-async def _safe_close_websocket(websocket: WebSocket):
+async def _safe_close_websocket(
+    websocket: WebSocket,
+    code: int = 1000,
+    reason: str = "Connection closed",
+):
     try:
-        await websocket.close(code=1000, reason="Connection closed")
+        await websocket.close(code=code, reason=reason)
     except Exception:
         # WebSocket may already be closed or in inconsistent state
         logger.debug("WebSocket close failed (may already be closed)")
diff --git a/tests/agent_server/test_websocket_first_message_auth.py b/tests/agent_server/test_websocket_first_message_auth.py
@@ -0,0 +1,251 @@
+"""Tests for first-message WebSocket authentication in sockets.py."""
+
+import asyncio
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+import pytest
+from fastapi import WebSocketDisconnect
+
+from openhands.agent_server.sockets import _accept_authenticated_websocket
+
+
+def _make_mock_websocket(*, headers=None):
+    """Build a mock WebSocket with configurable query params and headers."""
+    ws = MagicMock()
+    ws.accept = AsyncMock()
+    ws.receive_text = AsyncMock()
+    ws.receive_json = AsyncMock()
+    ws.send_json = AsyncMock()
+    ws.close = AsyncMock()
+    ws.headers = headers or {}
+    return ws
+
+
+# -- No auth configured (empty session_api_keys) --
+
+
+@pytest.mark.asyncio
+async def test_no_auth_configured_accepts_immediately():
+    ws = _make_mock_websocket()
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = []
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is True
+    ws.accept.assert_called_once()
+    ws.receive_text.assert_not_called()
+
+
+# -- Legacy query param auth (deprecated) --
+
+
+@pytest.mark.asyncio
+async def test_legacy_query_param_valid_key():
+    ws = _make_mock_websocket()
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key="sk-oh-valid")
+
+    assert result is True
+    ws.accept.assert_called_once()
+    ws.receive_text.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_legacy_query_param_invalid_key():
+    ws = _make_mock_websocket()
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key="sk-oh-wrong")
+
+    assert result is False
+    ws.close.assert_called_once_with(code=4001, reason="Authentication failed")
+    ws.accept.assert_not_called()
+
+
+# -- Legacy header auth (deprecated) --
+
+
+@pytest.mark.asyncio
+async def test_legacy_header_valid_key():
+    ws = _make_mock_websocket(headers={"x-session-api-key": "sk-oh-valid"})
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is True
+    ws.accept.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_legacy_header_invalid_key():
+    ws = _make_mock_websocket(headers={"x-session-api-key": "sk-oh-wrong"})
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is False
+    ws.close.assert_called_once_with(code=4001, reason="Authentication failed")
+
+
+# -- First-message auth --
+
+
+@pytest.mark.asyncio
+async def test_first_message_auth_valid_key():
+    ws = _make_mock_websocket()
+    ws.receive_text.return_value = json.dumps(
+        {"type": "auth", "session_api_key": "sk-oh-valid"}
+    )
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is True
+    ws.accept.assert_called_once()
+    ws.receive_text.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_first_message_auth_invalid_key():
+    ws = _make_mock_websocket()
+    ws.receive_text.return_value = json.dumps(
+        {"type": "auth", "session_api_key": "sk-oh-wrong"}
+    )
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is False
+    ws.accept.assert_called_once()  # accepted before reading first message
+    ws.close.assert_called_once_with(code=4001, reason="Authentication failed")
+
+
+@pytest.mark.asyncio
+async def test_first_message_auth_wrong_type_field():
+    ws = _make_mock_websocket()
+    ws.receive_text.return_value = json.dumps(
+        {"type": "message", "session_api_key": "sk-oh-valid"}
+    )
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is False
+
+
+@pytest.mark.asyncio
+async def test_first_message_auth_missing_key_field():
+    ws = _make_mock_websocket()
+    ws.receive_text.return_value = json.dumps({"type": "auth"})
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is False
+
+
+@pytest.mark.asyncio
+async def test_first_message_auth_malformed_json():
+    ws = _make_mock_websocket()
+    ws.receive_text.return_value = "not json at all"
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is False
+    ws.close.assert_called_once_with(code=4001, reason="Authentication failed")
+
+
+@pytest.mark.asyncio
+async def test_first_message_auth_client_disconnects():
+    ws = _make_mock_websocket()
+    ws.receive_text.side_effect = WebSocketDisconnect()
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is False
+
+
+@pytest.mark.asyncio
+async def test_first_message_auth_timeout():
+    ws = _make_mock_websocket()
+
+    async def slow_receive():
+        await asyncio.sleep(60)
+
+    ws.receive_text.side_effect = slow_receive
+
+    with (
+        patch("openhands.agent_server.sockets.get_default_config") as mock_config,
+        patch(
+            "openhands.agent_server.sockets._FIRST_MESSAGE_AUTH_TIMEOUT_SECONDS", 0.05
+        ),
+    ):
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(ws, session_api_key=None)
+
+    assert result is False
+    ws.close.assert_called_once_with(code=4001, reason="Authentication failed")
+
+
+# -- End-to-end: first-message auth through events_socket --
+
+
+@pytest.mark.asyncio
+async def test_events_socket_first_message_auth_e2e():
+    """First-message auth works end-to-end through the events_socket endpoint."""
+    from openhands.agent_server.event_service import EventService
+    from openhands.agent_server.sockets import events_socket
+
+    ws = _make_mock_websocket()
+    # First call: auth message (read by _accept_authenticated_websocket via receive_text)
+    # Then receive_json calls: disconnect
+    ws.receive_text.return_value = json.dumps(
+        {"type": "auth", "session_api_key": "sk-oh-valid"}
+    )
+    ws.receive_json.side_effect = WebSocketDisconnect()
+
+    mock_event_service = MagicMock(spec=EventService)
+    mock_event_service.subscribe_to_events = AsyncMock(return_value=uuid4())
+    mock_event_service.unsubscribe_from_events = AsyncMock(return_value=True)
+
+    with (
+        patch(
+            "openhands.agent_server.sockets.conversation_service"
+        ) as mock_conv_service,
+        patch("openhands.agent_server.sockets.get_default_config") as mock_config,
+    ):
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        mock_conv_service.get_event_service = AsyncMock(
+            return_value=mock_event_service
+        )
+
+        await events_socket(uuid4(), ws, session_api_key=None)
+
+    ws.accept.assert_called_once()
+    mock_event_service.subscribe_to_events.assert_called_once()
+    mock_event_service.unsubscribe_from_events.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_events_socket_first_message_auth_rejected():
+    """events_socket returns early when first-message auth fails."""
+    from openhands.agent_server.sockets import events_socket
+
+    ws = _make_mock_websocket()
+    ws.receive_text.return_value = json.dumps(
+        {"type": "auth", "session_api_key": "sk-oh-wrong"}
+    )
+
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+
+        await events_socket(uuid4(), ws, session_api_key=None)
+
+    ws.accept.assert_called_once()
+    # Should not proceed to subscribe
+    ws.receive_json.assert_not_called()
