mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 05:12:26 +00:00
krypticmouseandClaude Opus 4.7
87e978ef20
`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) <noreply@anthropic.com>
116 lines
4.0 KiB
Python
116 lines
4.0 KiB
Python
"""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"
|