From 87e978ef207b86fe5351cbe2d7a29a1823d9b22a Mon Sep 17 00:00:00 2001 From: krypticmouse Date: Mon, 25 May 2026 23:51:29 +0000 Subject: [PATCH] security(server): authenticate WebSocket handshakes and A2A requests (#217) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `AuthMiddleware` is a BaseHTTPMiddleware and never intercepts WebSocket upgrade requests, so `/v1/chat/stream` and `/v1/agents/events` accepted any connection — leaking all agent events/message content and allowing unauthenticated inference even when an API key was configured for HTTP. The A2A JSON-RPC server likewise dispatched every request without auth. - Add `websocket_authorized(websocket, expected_key)` (constant-time compare) and check it in both WS handlers BEFORE `accept()`, closing with code 1008 on failure. Token is read from `?token=` (browsers can't set WS headers) or an `Authorization: Bearer` header. `create_app` now exposes the key via `app.state.api_key`; when empty, auth is disabled, matching the HTTP middleware's local-default behavior (so loopback dev is unchanged). - A2AServer gains an optional `auth_token`: `handle_request(token=...)` rejects with JSON-RPC -32001 before dispatch when configured, advertises `{"schemes": ["bearer"]}` on the agent card, and stays open when unset. Added `A2AConfig.auth_token`. Verified empirically against the real mounted endpoints via TestClient: no token / wrong token are rejected at the handshake (WebSocketDisconnect), correct token streams normally, and no-key configs still connect freely. Closes #217 Co-Authored-By: Claude Opus 4.7 (1M context) --- src/openjarvis/a2a/server.py | 42 ++++++++- src/openjarvis/core/config.py | 3 + src/openjarvis/server/api_routes.py | 7 ++ src/openjarvis/server/app.py | 3 + src/openjarvis/server/auth_middleware.py | 26 +++++ src/openjarvis/server/ws_bridge.py | 7 ++ tests/a2a/test_a2a_auth.py | 51 ++++++++++ tests/server/test_websocket_auth.py | 115 +++++++++++++++++++++++ 8 files changed, 251 insertions(+), 3 deletions(-) create mode 100644 tests/a2a/test_a2a_auth.py create mode 100644 tests/server/test_websocket_auth.py diff --git a/src/openjarvis/a2a/server.py b/src/openjarvis/a2a/server.py index 29ca5eba..212517cb 100644 --- a/src/openjarvis/a2a/server.py +++ b/src/openjarvis/a2a/server.py @@ -2,6 +2,7 @@ from __future__ import annotations +import secrets from typing import Any, Callable, Dict, List, Optional from openjarvis.a2a.protocol import ( @@ -17,6 +18,11 @@ class A2AServer: """A2A server that processes incoming tasks via agent execution. Can be mounted as routes in the FastAPI server. + + When *auth_token* is set, every :meth:`handle_request` call must present a + matching bearer token or it is rejected before any agent runs. The token + is advertised on the agent card's ``authentication`` field. When unset, + the server is unauthenticated — only mount it on a trusted network. """ def __init__( @@ -25,21 +31,51 @@ class A2AServer: *, handler: Optional[Callable[[str], str]] = None, bus: Optional[EventBus] = None, + auth_token: Optional[str] = None, ) -> None: self._card = agent_card self._handler = handler self._bus = bus + self._auth_token = auth_token or None self._tasks: Dict[str, A2ATask] = {} + if self._auth_token: + # Advertise the required scheme on the discovery card. + self._card.authentication = {"schemes": ["bearer"]} @property def agent_card(self) -> AgentCard: return self._card - def handle_request(self, request_data: Dict[str, Any]) -> Dict[str, Any]: - """Process a JSON-RPC 2.0 A2A request.""" + def authenticate(self, token: Optional[str]) -> bool: + """Constant-time check of a presented bearer *token*. + + Returns ``True`` when no ``auth_token`` is configured (auth disabled). + """ + if not self._auth_token: + return True + return bool(token) and secrets.compare_digest(token, self._auth_token) + + def handle_request( + self, + request_data: Dict[str, Any], + *, + token: Optional[str] = None, + ) -> Dict[str, Any]: + """Process a JSON-RPC 2.0 A2A request. + + *token* is the bearer credential extracted by the transport (e.g. the + HTTP ``Authorization`` header). It is validated before dispatch when + the server is configured with an ``auth_token``. + """ + req_id = request_data.get("id", "") + if not self.authenticate(token): + return A2AResponse( + error={"code": -32001, "message": "Unauthorized"}, + request_id=req_id, + ).to_dict() + method = request_data.get("method", "") params = request_data.get("params", {}) - req_id = request_data.get("id", "") if method == "tasks/send": return self._handle_task_send(params, req_id) diff --git a/src/openjarvis/core/config.py b/src/openjarvis/core/config.py index c7f6bb7d..e2c94705 100644 --- a/src/openjarvis/core/config.py +++ b/src/openjarvis/core/config.py @@ -1371,6 +1371,9 @@ class A2AConfig: """Agent-to-Agent protocol settings.""" enabled: bool = False + # Bearer token required for inbound A2A requests. Empty = unauthenticated + # (only safe on a trusted network). See ``A2AServer(auth_token=...)``. + auth_token: str = "" @dataclass(slots=True) diff --git a/src/openjarvis/server/api_routes.py b/src/openjarvis/server/api_routes.py index 956de624..f4e01d6a 100644 --- a/src/openjarvis/server/api_routes.py +++ b/src/openjarvis/server/api_routes.py @@ -568,6 +568,13 @@ async def websocket_chat_stream(websocket: WebSocket): {"type": "done", "content": "..."} -- final assembled response {"type": "error", "detail": "..."} -- on failure """ + from openjarvis.server.auth_middleware import websocket_authorized + + expected_key = getattr(websocket.app.state, "api_key", "") + if not websocket_authorized(websocket, expected_key): + # 1008 = policy violation; reject before accepting the connection. + await websocket.close(code=1008) + return await websocket.accept() try: while True: diff --git a/src/openjarvis/server/app.py b/src/openjarvis/server/app.py index 762abf13..6828eb8e 100644 --- a/src/openjarvis/server/app.py +++ b/src/openjarvis/server/app.py @@ -223,6 +223,9 @@ def create_app( app.state.agent_manager = agent_manager app.state.agent_scheduler = agent_scheduler app.state.session_start = time.time() + # Exposed so WebSocket handlers can authenticate the handshake (the HTTP + # AuthMiddleware never sees WS upgrade requests). Empty = auth disabled. + app.state.api_key = api_key # Wire up trace store if traces are enabled app.state.trace_store = None diff --git a/src/openjarvis/server/auth_middleware.py b/src/openjarvis/server/auth_middleware.py index fb7e05fd..a9390ee5 100644 --- a/src/openjarvis/server/auth_middleware.py +++ b/src/openjarvis/server/auth_middleware.py @@ -73,3 +73,29 @@ def check_bind_safety(host: str, *, api_key: str) -> None: host, ) sys.exit(1) + + +def websocket_authorized(websocket, expected_key: str) -> bool: # noqa: ANN001 + """Return ``True`` if a WebSocket connection presents the expected key. + + ``AuthMiddleware`` is a ``BaseHTTPMiddleware`` and never sees WebSocket + upgrade requests, so streaming endpoints must check the token themselves + in the handshake before calling ``websocket.accept()``. + + When *expected_key* is empty, authentication is disabled (the loopback / + local-only default, matching :class:`AuthMiddleware`) and all connections + are allowed. The token may be supplied either as a ``?token=`` query + parameter — browsers cannot set headers on a WebSocket handshake — or via + an ``Authorization: Bearer `` header for programmatic clients. + """ + if not expected_key: + return True + token = websocket.query_params.get("token", "") + if not token: + auth = websocket.headers.get("authorization", "") + scheme, _, value = auth.partition(" ") + if scheme.lower() == "bearer": + token = value + if not token: + return False + return secrets.compare_digest(token, expected_key) diff --git a/src/openjarvis/server/ws_bridge.py b/src/openjarvis/server/ws_bridge.py index 7edda66e..60842074 100644 --- a/src/openjarvis/server/ws_bridge.py +++ b/src/openjarvis/server/ws_bridge.py @@ -60,6 +60,13 @@ def create_ws_router(event_bus: EventBus) -> Any: @router.websocket("/v1/agents/events") async def agent_events(websocket: WebSocket) -> None: + from openjarvis.server.auth_middleware import websocket_authorized + + expected_key = getattr(websocket.app.state, "api_key", "") + if not websocket_authorized(websocket, expected_key): + # 1008 = policy violation; reject before accepting the connection. + await websocket.close(code=1008) + return await websocket.accept() # Parse agent_id filter from query string agent_id = websocket.query_params.get("agent_id") diff --git a/tests/a2a/test_a2a_auth.py b/tests/a2a/test_a2a_auth.py new file mode 100644 index 00000000..9ef56b66 --- /dev/null +++ b/tests/a2a/test_a2a_auth.py @@ -0,0 +1,51 @@ +"""A2A server authentication tests (issue #217).""" + +from __future__ import annotations + +from openjarvis.a2a.protocol import AgentCard +from openjarvis.a2a.server import A2AServer + + +def _request(): + return { + "jsonrpc": "2.0", + "id": "1", + "method": "tasks/send", + "params": {"input": "ping"}, + } + + +def test_no_auth_token_allows_requests(): + server = A2AServer(AgentCard(name="x"), handler=lambda t: f"echo:{t}") + resp = server.handle_request(_request()) + assert resp.get("error") is None + + +def test_missing_token_rejected_when_required(): + server = A2AServer(AgentCard(name="x"), handler=lambda t: t, auth_token="sek") + resp = server.handle_request(_request()) + assert resp["error"]["code"] == -32001 + + +def test_wrong_token_rejected(): + server = A2AServer(AgentCard(name="x"), handler=lambda t: t, auth_token="sek") + resp = server.handle_request(_request(), token="nope") + assert resp["error"]["code"] == -32001 + + +def test_correct_token_accepted_and_dispatched(): + server = A2AServer( + AgentCard(name="x"), handler=lambda t: f"echo:{t}", auth_token="sek" + ) + resp = server.handle_request(_request(), token="sek") + assert resp.get("error") is None + + +def test_auth_scheme_advertised_on_card(): + server = A2AServer(AgentCard(name="x"), auth_token="sek") + assert server.agent_card.authentication == {"schemes": ["bearer"]} + + +def test_no_token_unauthenticated_card_stays_empty(): + server = A2AServer(AgentCard(name="x")) + assert server.agent_card.authentication == {} diff --git a/tests/server/test_websocket_auth.py b/tests/server/test_websocket_auth.py new file mode 100644 index 00000000..58750187 --- /dev/null +++ b/tests/server/test_websocket_auth.py @@ -0,0 +1,115 @@ +"""WebSocket authentication tests (issue #217). + +The HTTP ``AuthMiddleware`` never sees WebSocket upgrade requests, so the +streaming endpoints (`/v1/chat/stream`, `/v1/agents/events`) must validate the +token themselves in the handshake before accepting the connection. +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest + +fastapi = pytest.importorskip("fastapi") +from fastapi import FastAPI # noqa: E402 +from starlette.testclient import TestClient # noqa: E402 +from starlette.websockets import WebSocketDisconnect # noqa: E402 + +from openjarvis.core.events import EventBus, EventType # noqa: E402 +from openjarvis.server.api_routes import include_all_routes # noqa: E402 +from openjarvis.server.auth_middleware import websocket_authorized # noqa: E402 +from openjarvis.server.ws_bridge import create_ws_router # noqa: E402 + + +def _ws(query=None, headers=None): + stub = MagicMock() + stub.query_params = query or {} + stub.headers = headers or {} + return stub + + +class TestWebsocketAuthorizedHelper: + def test_no_key_allows_all(self): + assert websocket_authorized(_ws(), "") is True + + def test_token_via_query(self): + assert websocket_authorized(_ws(query={"token": "sek"}), "sek") is True + + def test_token_via_bearer_header(self): + ws = _ws(headers={"authorization": "Bearer sek"}) + assert websocket_authorized(ws, "sek") is True + + def test_wrong_token_rejected(self): + assert websocket_authorized(_ws(query={"token": "nope"}), "sek") is False + + def test_missing_token_rejected_when_required(self): + assert websocket_authorized(_ws(), "sek") is False + + +def _make_app(api_key=""): + app = FastAPI() + engine = MagicMock() + engine.engine_id = "mock" + + async def mock_stream(messages, *, model="test-model", **kwargs): + for tok in ["hi"]: + yield tok + + engine.stream = mock_stream + app.state.engine = engine + app.state.model = "test-model" + app.state.api_key = api_key + include_all_routes(app) + return app + + +class TestChatStreamAuth: + def test_rejected_without_token_when_key_set(self): + client = TestClient(_make_app(api_key="secret")) + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/v1/chat/stream") as ws: + ws.receive_text() + + def test_rejected_with_wrong_token(self): + client = TestClient(_make_app(api_key="secret")) + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/v1/chat/stream?token=wrong") as ws: + ws.receive_text() + + def test_accepted_with_correct_token(self): + client = TestClient(_make_app(api_key="secret")) + with client.websocket_connect("/v1/chat/stream?token=secret") as ws: + ws.send_text(json.dumps({"message": "hi"})) + assert ws.receive_json()["type"] in ("chunk", "done", "error") + + def test_allowed_when_no_key_configured(self): + client = TestClient(_make_app(api_key="")) + with client.websocket_connect("/v1/chat/stream") as ws: + ws.send_text(json.dumps({"message": "hi"})) + assert ws.receive_json()["type"] in ("chunk", "done", "error") + + +class TestAgentEventsAuth: + def _app(self, api_key=""): + app = FastAPI() + app.state.api_key = api_key + app.include_router(create_ws_router(EventBus())) + return app + + def test_rejected_without_token_when_key_set(self): + client = TestClient(self._app(api_key="secret")) + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/v1/agents/events") as ws: + ws.receive_text() + + def test_accepted_with_correct_token(self): + bus = EventBus() + app = FastAPI() + app.state.api_key = "secret" + app.include_router(create_ws_router(bus)) + client = TestClient(app) + with client.websocket_connect("/v1/agents/events?token=secret") as ws: + bus.publish(EventType.AGENT_TICK_START, {"agent_id": "a"}) + assert ws.receive_json()["data"]["agent_id"] == "a"