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

simonrosenberg · Debug Agent · claude · web-flow · commit fc9afbbf8f0b · 2026-04-23T16:44:59.000+01:00
Co-authored-by: Debug Agent &lt;debug@example.com&gt;
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,18 +85,102 @@ def _resolve_websocket_session_api_key(
     return None
 
 
+# Give clients 10 seconds to send auth frame after connection opens.
+# This balances security (don't hold connections indefinitely) with
+# accommodating slow networks and client startup time.
+_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: we must accept() before reading frames because the
+    # WebSocket protocol requires the handshake to complete first.  The legacy
+    # path above can reject *before* accepting (close on an un-accepted socket
+    # sends an HTTP 403-style response), but here we need to read a frame.
     await websocket.accept()
+    try:
+        raw = await asyncio.wait_for(
+            websocket.receive_text(),
+            timeout=_FIRST_MESSAGE_AUTH_TIMEOUT_SECONDS,
+        )
+        data = json.loads(raw)
+    except TimeoutError:
+        logger.warning(
+            "WebSocket first-message auth failed: timeout waiting for auth frame"
+        )
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+    except json.JSONDecodeError:
+        logger.warning("WebSocket first-message auth failed: malformed JSON")
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+    except WebSocketDisconnect:
+        logger.warning("WebSocket first-message auth failed: client disconnected")
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+
+    if not isinstance(data, dict):
+        logger.warning(
+            "WebSocket first-message auth failed: payload is not a JSON object"
+        )
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+    if data.get("type") != "auth":
+        logger.warning("WebSocket first-message auth failed: wrong message type")
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+    if data.get("session_api_key") not in config.session_api_keys:
+        logger.warning("WebSocket first-message auth failed: invalid API key")
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+
+    logger.info("WebSocket authenticated via first-message auth")
     return True
 
 
@@ -329,9 +420,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_agent_server_wsproto.py b/tests/agent_server/test_agent_server_wsproto.py
@@ -7,10 +7,12 @@
 import socket
 import sys
 import time
+from uuid import uuid4
 
 import pytest
 import requests
 import websockets
+import websockets.exceptions
 
 
 def find_free_port():
@@ -155,3 +157,96 @@ async def test_agent_server_websocket_with_wsproto_header_auth(agent_server):
                 {"role": "user", "content": "Hello from wsproto header auth test"}
             )
         )
+
+
+@pytest.mark.asyncio
+async def test_agent_server_websocket_first_message_auth_accepted(agent_server):
+    """First-message auth: connect with no query/header key, auth via first frame.
+
+    Exercises the real WebSocket protocol transition (handshake → consume first
+    frame as auth → continue normal message flow) that mock-only tests can't
+    cover. See PR review feedback on test coverage gaps.
+    """
+    port = agent_server["port"]
+    api_key = agent_server["api_key"]
+
+    response = requests.post(
+        f"http://127.0.0.1:{port}/api/conversations",
+        headers={"X-Session-API-Key": api_key},
+        json={
+            "agent": {
+                "kind": "Agent",
+                "llm": {
+                    "usage_id": "test-llm",
+                    "model": "test-provider/test-model",
+                    "api_key": "test-key",
+                },
+                "tools": [],
+            },
+            "workspace": {"working_dir": "/tmp/test-workspace"},
+        },
+    )
+    assert response.status_code in [200, 201]
+    conversation_id = response.json()["id"]
+
+    # No session_api_key in URL or header — must authenticate via first frame.
+    ws_url = f"ws://127.0.0.1:{port}/sockets/events/{conversation_id}?resend_all=true"
+
+    async with websockets.connect(ws_url, open_timeout=5) as ws:
+        # Send the auth frame as the very first message after handshake.
+        await ws.send(json.dumps({"type": "auth", "session_api_key": api_key}))
+
+        # Connection must remain usable: try to receive (resend_all may produce
+        # nothing for an empty conversation, so a timeout here is fine).
+        try:
+            response = await asyncio.wait_for(ws.recv(), timeout=2)
+            assert response is not None
+        except TimeoutError:
+            pass
+
+        # Subsequent message must be processed as a Message (not auth) — proves
+        # the auth frame was consumed by the auth handler, not the main loop.
+        await ws.send(
+            json.dumps({"role": "user", "content": "Hello after first-message auth"})
+        )
+
+
+@pytest.mark.asyncio
+async def test_agent_server_websocket_first_message_auth_rejected(agent_server):
+    """First-message auth: invalid key triggers WebSocket close with code 4001."""
+    port = agent_server["port"]
+
+    # No conversation needed — auth rejection happens before conversation lookup.
+    ws_url = f"ws://127.0.0.1:{port}/sockets/events/{uuid4()}"
+
+    async with websockets.connect(ws_url, open_timeout=5) as ws:
+        # Send an invalid first-message auth frame.
+        await ws.send(
+            json.dumps({"type": "auth", "session_api_key": "definitely-wrong-key"})
+        )
+
+        # Server must close the connection with code 4001 ("Authentication
+        # failed"). Receiving on a closed socket raises ConnectionClosed.
+        with pytest.raises(websockets.exceptions.ConnectionClosed) as exc_info:
+            await asyncio.wait_for(ws.recv(), timeout=5)
+
+    assert exc_info.value.rcvd is not None
+    assert exc_info.value.rcvd.code == 4001
+
+
+@pytest.mark.asyncio
+async def test_agent_server_websocket_first_message_auth_malformed(agent_server):
+    """First-message auth: malformed JSON triggers close with code 4001."""
+    port = agent_server["port"]
+
+    ws_url = f"ws://127.0.0.1:{port}/sockets/events/{uuid4()}"
+
+    async with websockets.connect(ws_url, open_timeout=5) as ws:
+        # Send invalid JSON as the first frame.
+        await ws.send("this is not json")
+
+        with pytest.raises(websockets.exceptions.ConnectionClosed) as exc_info:
+            await asyncio.wait_for(ws.recv(), timeout=5)
+
+    assert exc_info.value.rcvd is not None
+    assert exc_info.value.rcvd.code == 4001
diff --git a/tests/agent_server/test_websocket_first_message_auth.py b/tests/agent_server/test_websocket_first_message_auth.py
