diff --git a/src/bluesky_httpserver/authentication.py b/src/bluesky_httpserver/authentication.py index 3408983..34f93cf 100644 --- a/src/bluesky_httpserver/authentication.py +++ b/src/bluesky_httpserver/authentication.py @@ -253,7 +253,16 @@ def move_api_key(request: Request, api_key: str | None = Depends(get_api_key)): are authenticated. (This approach was inspired by Jupyter notebook.) """ if ("api_key" in request.query_params) and (request.cookies.get(API_KEY_COOKIE_NAME) != api_key): - request.state.cookies_to_set.append({"key": API_KEY_COOKIE_NAME, "value": api_key}) + # The `set_cookies` HTTP middleware initializes `request.state.cookies_to_set`. + # That middleware is not executed for WebSocket connections, so guard here + # to avoid AttributeError when `state.cookies_to_set` is missing. + if hasattr(request.state, "cookies_to_set"): + request.state.cookies_to_set.append({"key": API_KEY_COOKIE_NAME, "value": api_key}) + else: + # No place to store a cookie for this request (likely a WebSocket). + # Silently skip setting the cookie — browsers opening a websocket + # with ?api_key=... should set the cookie via an HTTP request instead. + logger.debug("Skipping moving api_key to cookie: request.state.cookies_to_set missing") def get_current_principal( diff --git a/src/bluesky_httpserver/tests/test_auth_for_websockets.py b/src/bluesky_httpserver/tests/test_auth_for_websockets.py index c860a5a..bfafde9 100644 --- a/src/bluesky_httpserver/tests/test_auth_for_websockets.py +++ b/src/bluesky_httpserver/tests/test_auth_for_websockets.py @@ -7,6 +7,7 @@ import pytest from bluesky_queueserver.manager.tests.common import re_manager, re_manager_cmd, re_manager_factory # noqa F401 from sqlalchemy.orm import sessionmaker +from starlette.datastructures import State from websockets.sync.client import connect from bluesky_httpserver import authentication as _auth @@ -217,6 +218,40 @@ def __init__(self): return _WS() +def _fake_request_with_state(state): + """Minimal stand-in for a request carrying an API key in the query string.""" + + class _Request: + query_params = {"api_key": "some-api-key"} + cookies: dict = {} + + request = _Request() + request.state = state + return request + + +def test_move_api_key_sets_cookie_on_http_request(): + """On the HTTP path the middleware has initialized 'cookies_to_set'.""" + state = State() + state.cookies_to_set = [] + request = _fake_request_with_state(state) + + _auth.move_api_key(request, "some-api-key") + + assert state.cookies_to_set == [{"key": _auth.API_KEY_COOKIE_NAME, "value": "some-api-key"}] + + +def test_move_api_key_skips_cookie_on_websocket_request(): + """The HTTP middleware does not run for websockets, so 'cookies_to_set' + is missing and the API key simply is not moved to a cookie.""" + state = State() + request = _fake_request_with_state(state) + + _auth.move_api_key(request, "some-api-key") + + assert not hasattr(state, "cookies_to_set") + + def test_authenticate_websocket_first_message_rejects_non_auth_frames(): ws = _fake_ws_with_deps(settings=MagicMock()) assert _auth.authenticate_websocket_first_message(ws, {"type": "ping"}) is None