From b3a0f69440d5027af469adc79c86ee5d1fff49a4 Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Fri, 13 Nov 2026 19:16:49 +0000 Subject: [PATCH] wip: fix login logout --- docs/how-to/1.2 authn-authz.md | 4 +- src/blueapi/cli/cli.py | 36 +-- src/blueapi/client/client.py | 50 +++- src/blueapi/client/rest.py | 129 ++++++--- src/blueapi/client/session.py | 263 +++++++++++++---- src/blueapi/config.py | 2 +- tests/conftest.py | 3 +- tests/system_tests/test_auth_token_cache.py | 186 +++++++++++++ tests/unit_tests/cli/test_cli.py | 84 ++++++ tests/unit_tests/client/test_client.py | 25 +- tests/unit_tests/client/test_rest.py | 181 +++++++++++- tests/unit_tests/client/test_session.py | 294 +++++++++++++++++++- 12 files changed, 1129 insertions(+), 128 deletions(-) create mode 100644 tests/system_tests/test_auth_token_cache.py diff --git a/docs/how-to/1.2 authn-authz.md b/docs/how-to/1.2 authn-authz.md index 5671d14ec..faa18cfbe 100644 --- a/docs/how-to/1.2 authn-authz.md +++ b/docs/how-to/1.2 authn-authz.md @@ -14,7 +14,9 @@ Here is an example configuration for authenticating to p46-blueapi: :language: yaml ``` -- **auth_token_path**: (Optional) Specify where to save the token. If omitted, the default is `~/.cache/blueapi_cache` or `$XDG_CACHE_HOME/blueapi_cache` if `XDG_CACHE_HOME` is set. +- **auth_token_path**: (Optional) Specify where to save the token. If omitted, a separate cache is kept for each server, named after its host (and port, if given): e.g. `~/.cache/blueapi_cache.i22-blueapi.diamond.ac.uk` or `~/.cache/blueapi_cache.localhost_8000`. `$XDG_CACHE_HOME` is used instead of `~/.cache` if set. + + Access tokens are refreshed shortly before they expire. If the authentication server cannot be reached the login is kept and an error explains what happened; the cache is only removed once the authentication server rejects the refresh token, at which point you need to log in again. --- diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index d78a88d01..33c403376 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -31,7 +31,11 @@ UnauthorisedAccessError, UnknownPlanError, ) -from blueapi.client.session import SessionCacheManager, SessionManager +from blueapi.client.session import ( + AuthenticationServerUnavailableError, + SessionManager, + default_cache_path, +) from blueapi.config import ( ApplicationConfig, ConfigLoader, @@ -291,6 +295,8 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: raise ClickException( "Access denied. Please check your login status and try again." ) from e + except AuthenticationServerUnavailableError as e: + raise ClickException(str(e)) from e return wrapper @@ -604,23 +610,11 @@ def login(obj: dict) -> None: Authenticate with the blueapi using the OIDC (OpenID Connect) flow. """ config: ApplicationConfig = obj["config"] + client = BlueapiClient.from_config(config) try: - auth: SessionManager = SessionManager.from_cache(config.auth_token_path) - access_token = auth.get_valid_access_token() - assert access_token - print("Logged in") - except Exception: - client = BlueapiClient.from_config(config) - try: - if oidc := client.oidc_config: - auth = SessionManager( - oidc, cache_manager=SessionCacheManager(config.auth_token_path) - ) - auth.start_device_flow() - else: - print("Server is not configured to use authentication!") - except NonJsonResponseError as e: - print(str(e)) + client.login(config.auth_token_path) + except NonJsonResponseError as e: + print(str(e)) @main.command(name="logout") @@ -631,7 +625,13 @@ def logout(obj: dict) -> None: """ config: ApplicationConfig = obj["config"] try: - auth: SessionManager = SessionManager.from_cache(config.auth_token_path) + if config.auth_token_path is None and (legacy := default_cache_path()).exists(): + # Login cached by an older blueapi at the shared default path; log it + # out too so its refresh token is revoked rather than orphaned + SessionManager.from_cache(legacy).logout() + auth: SessionManager = SessionManager.from_cache( + config.auth_token_path, str(config.api.url) + ) auth.logout() except FileNotFoundError: print("Logged out") diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 2105ac64e..dbfc2dca2 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -13,7 +13,11 @@ from observability_utils.tracing import get_tracer, start_as_current_span from pydantic import HttpUrl -from blueapi.client.session import SessionCacheManager, SessionManager +from blueapi.client.session import ( + SessionCacheManager, + SessionManager, + migrate_legacy_cache, +) from blueapi.config import ( ApplicationConfig, ConfigLoader, @@ -90,7 +94,9 @@ def from_config_file(cls, config_file: str) -> Self: def from_config(cls, config: ApplicationConfig) -> Self: session_manager: SessionManager | None = None try: - session_manager = SessionManager.from_cache(config.auth_token_path) + session_manager = SessionManager.from_cache( + config.auth_token_path, str(config.api.url) + ) except Exception: ... # Swallow exceptions rest = BlueapiRestClient(config.api, session_manager=session_manager) @@ -657,20 +663,40 @@ def get_python_env( return self._rest.get_python_environment(name=name, source=source) def login(self, token_path: Path | None = None): + """ + Log in to the server's OIDC provider, reusing a cached login if it is + still valid and was made with the same OIDC config the server uses now. + + Raises: + AuthenticationServerUnavailableError: If the cached login needs + refreshing but the OIDC provider cannot be reached. + """ + oidc = self.oidc_config + if oidc is None: + print("Server is not configured to use authentication!") + return + server_url = str(self._rest.url) + if token_path is None: + migrate_legacy_cache(server_url, oidc) try: - auth: SessionManager = SessionManager.from_cache(token_path) - access_token = auth.get_valid_access_token() - assert access_token - print("Logged in") + auth = SessionManager.from_cache(token_path, server_url) except Exception: - if oidc := self.oidc_config: - auth = SessionManager( - oidc, cache_manager=SessionCacheManager(token_path) + auth = None + if auth is not None: + if not auth.is_for(oidc): + print( + "The server's authentication config has changed since " + "you logged in, logging in again" ) - auth.start_device_flow() + elif auth.get_valid_access_token(): + print("Logged in") self._rest.session_manager = auth - else: - print("Server is not configured to use authentication!") + return + auth = SessionManager( + oidc, cache_manager=SessionCacheManager(token_path, server_url) + ) + auth.start_device_flow() + self._rest.session_manager = auth def logout(self): if sm := self._rest.session_manager: diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py index 7bd242a01..4ba495c6f 100644 --- a/src/blueapi/client/rest.py +++ b/src/blueapi/client/rest.py @@ -10,14 +10,14 @@ get_tracer, start_as_current_span, ) -from pydantic import BaseModel, TypeAdapter, ValidationError +from pydantic import BaseModel, HttpUrl, TypeAdapter, ValidationError from pydantic_core import PydanticSerializationError from websockets.exceptions import InvalidStatus -from websockets.sync.client import connect +from websockets.sync.client import ClientConnection, connect from blueapi import __version__ from blueapi.client.protocols import ClientObjectRef -from blueapi.client.session import JWTAuth, SessionManager +from blueapi.client.session import SessionManager from blueapi.config import RestConfig from blueapi.core.bluesky_types import DataEvent from blueapi.service.model import ( @@ -201,6 +201,24 @@ def _response_json(response: requests.Response) -> Any: ) from exc +def _access_token(session_manager: SessionManager | None) -> str: + return session_manager.get_valid_access_token() if session_manager else "" + + +def _replacement_token(session_manager: SessionManager | None, rejected: str) -> str: + """A different token to retry with after `rejected` was refused, or """ "" + if not (session_manager and rejected): + return "" + token = session_manager.get_valid_access_token(rejected_token=rejected) + return token if token != rejected else "" + + +def _with_bearer(headers: Mapping[str, str], token: str) -> dict[str, str]: + if not token: + return dict(headers) + return {**headers, "Authorization": f"Bearer {token}"} + + class BlueapiRestClient: _config: RestConfig session_manager: SessionManager | None @@ -215,6 +233,10 @@ def __init__( self.session_manager = session_manager self._pool = requests.Session() + @property + def url(self) -> HttpUrl: + return self._config.url + def get_plans(self) -> PlanResponse: return self._request_and_deserialize("/plans", PlanResponse) @@ -298,7 +320,11 @@ def delete_environment(self) -> EnvironmentResponse: def get_oidc_config(self) -> OIDCConfig | None: try: - return self._request_and_deserialize("/config/oidc", OIDCConfig) + # Unauthenticated: a stale cached login (possibly for an OIDC provider + # the server no longer uses) must not stop us finding the new config + return self._request_and_deserialize( + "/config/oidc", OIDCConfig, authenticated=False + ) except NoContentError: # Server is not using authentication return None @@ -321,20 +347,31 @@ def _request_and_deserialize( method="GET", get_exception: Callable[[requests.Response], Exception | None] = _exception, params: Mapping[str, Any] | None = None, + authenticated: bool = True, ) -> T: url = self._config.url.unicode_string().removesuffix("/") + suffix # Get the trace context to propagate to the REST API headers = get_context_propagator() headers["User-Agent"] = USER_AGENT - try: - response = self._pool.request( + session_manager = self.session_manager if authenticated else None + + def send(token: str) -> requests.Response: + return self._pool.request( method, url, json=data, params=params, - headers=headers, - auth=JWTAuth(self.session_manager), + headers=_with_bearer(headers, token), ) + + try: + token = _access_token(session_manager) + response = send(token) + if response.status_code == HTTPStatus.UNAUTHORIZED and ( + new_token := _replacement_token(session_manager, token) + ): + # Token rejected (revoked, clock skew...): refresh and retry once + response = send(new_token) except requests.exceptions.ConnectionError as ce: raise ServiceUnavailableError() from ce exception = get_exception(response) @@ -362,39 +399,55 @@ def run_blocking( self, req: TaskRequest ) -> Iterable[DataEvent | WorkerEvent | ProgressEvent]: url = self._config.ws_address.unicode_string().rstrip("/") + "/api/v2/run_plan" + with self._connect_websocket(url) as ws: + ws.send(Submit(task=req).model_dump_json(fallback=_task_model_fallback)) + for message in ws: + event = ControlResponse.validate_json(message) + match event: + case Update(data=data): + yield data + case InvalidArgs(errors=errors): + raise InvalidParametersError( + [ + ParameterError( + loc=e.loc, msg=e.msg, type=e.type, input=e.input + ) + for e in errors + ] + ) + case PlanNotFound(plan_name=name): + raise UnknownPlanError(message=name) + case ServerBusy(): + raise BlueskyRemoteControlError(409, "Server is busy") + case Unauthorized(): + raise UnauthorisedAccessError( + 403, "Not authorized to submit task" + ) + + def _connect_websocket(self, url: str) -> ClientConnection: + """ + Open the websocket, refreshing the token and retrying once if the + handshake is rejected as unauthenticated. + """ headers = get_context_propagator() - if self.session_manager: - auth = self.session_manager.get_valid_access_token() - headers["Authorization"] = f"Bearer {auth}" - try: - with connect( + + def open_with(token: str) -> ClientConnection: + return connect( url, - additional_headers=headers, + additional_headers=_with_bearer(headers, token), user_agent_header=USER_AGENT, - ) as ws: - ws.send(Submit(task=req).model_dump_json(fallback=_task_model_fallback)) - for message in ws: - event = ControlResponse.validate_json(message) - match event: - case Update(data=data): - yield data - case InvalidArgs(errors=errors): - raise InvalidParametersError( - [ - ParameterError( - loc=e.loc, msg=e.msg, type=e.type, input=e.input - ) - for e in errors - ] - ) - case PlanNotFound(plan_name=name): - raise UnknownPlanError(message=name) - case ServerBusy(): - raise BlueskyRemoteControlError(409, "Server is busy") - case Unauthorized(): - raise UnauthorisedAccessError( - 403, "Not authorized to submit task" - ) + ) + + token = _access_token(self.session_manager) + try: + try: + return open_with(token) + except InvalidStatus as istat: + if istat.response.status_code == HTTPStatus.UNAUTHORIZED and ( + new_token := _replacement_token(self.session_manager, token) + ): + return open_with(new_token) + raise except InvalidStatus as istat: match istat.response.status_code: case 401 | 403: diff --git a/src/blueapi/client/session.py b/src/blueapi/client/session.py index 07ca6a7ae..4e4cab7e2 100644 --- a/src/blueapi/client/session.py +++ b/src/blueapi/client/session.py @@ -1,7 +1,9 @@ from __future__ import annotations import base64 +import logging import os +import tempfile import time import webbrowser from abc import ABC, abstractmethod @@ -9,6 +11,7 @@ from http import HTTPStatus from pathlib import Path from typing import Any, cast +from urllib.parse import urlsplit import jwt import requests @@ -18,8 +21,42 @@ from blueapi.config import OIDCConfig from blueapi.service.model import Cache +LOGGER = logging.getLogger(__name__) + DEFAULT_CACHE_DIR = "~/.cache/" +CACHE_FILE_PREFIX = "blueapi_cache" SCOPES = "openid offline_access" +# Refresh access tokens that expire within this many seconds, so that they do +# not expire between being checked here and being validated by the server +REFRESH_MARGIN_SECONDS = 30 +REQUEST_TIMEOUT_SECONDS = 10 + + +class AuthenticationServerUnavailableError(Exception): + """ + The OIDC provider could not be reached (or failed) while refreshing tokens. + + The cached login is kept, as the refresh token may still be valid once the + provider is reachable again. + """ + + +def default_cache_path(server_url: str | None = None) -> Path: + """ + Default token cache file, one per blueapi server so that logging in to one + server does not replace (or get confused with) the login for another. + + e.g. ~/.cache/blueapi_cache.i22-blueapi.diamond.ac.uk + ~/.cache/blueapi_cache.localhost_8000 + """ + cache_dir = Path(os.environ.get("XDG_CACHE_HOME", DEFAULT_CACHE_DIR)).expanduser() + if not server_url: + return cache_dir / CACHE_FILE_PREFIX + url = urlsplit(str(server_url)) + name = url.hostname or "default" + if url.port is not None: + name += f"_{url.port}" + return cache_dir / f"{CACHE_FILE_PREFIX}.{name}" class CacheManager(ABC): @@ -34,9 +71,9 @@ def delete_cache(self) -> None: ... class SessionCacheManager(CacheManager): - def __init__(self, token_path: Path | None) -> None: + def __init__(self, token_path: Path | None, server_url: str | None = None) -> None: self._token_path: Path = ( - token_path if token_path else self._default_token_cache_path() + Path(token_path) if token_path else default_cache_path(server_url) ) @cached_property @@ -44,11 +81,23 @@ def _file_path(self) -> str: return os.path.expanduser(self._token_path) def save_cache(self, cache: Cache) -> None: - self.delete_cache() + """ + Write the cache atomically so that another process reading it never + sees a partially written (or missing) file. + """ self._create_parent_folder_if_necessary() - with open(self._file_path, "xb") as token_file: - token_file.write(base64.b64encode(cache.model_dump_json().encode("utf-8"))) - os.chmod(self._file_path, 0o600) + directory = os.path.dirname(self._file_path) + # mkstemp creates the file with 0o600 permissions + fd, tmp_path = tempfile.mkstemp(dir=directory, prefix=".blueapi_cache.") + try: + with os.fdopen(fd, "wb") as token_file: + token_file.write( + base64.b64encode(cache.model_dump_json().encode("utf-8")) + ) + os.replace(tmp_path, self._file_path) + except BaseException: + Path(tmp_path).unlink(missing_ok=True) + raise def load_cache(self) -> Cache: self._create_parent_folder_if_necessary() @@ -60,14 +109,6 @@ def load_cache(self) -> Cache: def delete_cache(self) -> None: Path(self._file_path).unlink(missing_ok=True) - @staticmethod - def _default_token_cache_path() -> Path: - """ - Return the default cache file path. - """ - cache_path = os.environ.get("XDG_CACHE_HOME", DEFAULT_CACHE_DIR) - return Path(cache_path).expanduser() / "blueapi_cache" - def can_access_cache(self) -> bool: assert self._token_path try: @@ -91,49 +132,92 @@ def __init__(self, server_config: OIDCConfig, cache_manager: CacheManager) -> No self._cache_manager: CacheManager = cache_manager @classmethod - def from_cache(cls, auth_token_path: Path | None) -> SessionManager: - cache_manager = SessionCacheManager(auth_token_path) + def from_cache( + cls, auth_token_path: Path | None, server_url: str | None = None + ) -> SessionManager: + cache_manager = SessionCacheManager(auth_token_path, server_url) cache = cache_manager.load_cache() return SessionManager( server_config=cache.oidc_config, cache_manager=cache_manager ) + @property + def server_config(self) -> OIDCConfig: + return self._server_config + + def is_for(self, server_config: OIDCConfig) -> bool: + """ + Whether this session was created for the given OIDC config. If a server + changes its OIDC config, tokens cached for the old config are useless. + """ + return self._server_config.model_dump() == server_config.model_dump() + def delete_cache(self) -> None: self._cache_manager.delete_cache() - def get_valid_access_token(self) -> str: + def get_valid_access_token(self, rejected_token: str | None = None) -> str: """ - Retrieves a valid access token. + Retrieves a valid access token, refreshing it if it is about to expire. + + Args: + rejected_token: An access token the server has rejected. If the + cache still holds this token it is refreshed even if it has + not expired. If another process has already replaced it, the + newer token is used instead. Returns: str: A valid access token if successful. - "": If the operation fails (no valid token could be fetched or refreshed) + "": If there is no usable login. The cache is only deleted if it + is unreadable, or the OIDC provider rejects the refresh token. + + Raises: + AuthenticationServerUnavailableError: If the token needs refreshing + but the OIDC provider cannot be reached. The cache is kept. """ try: cache = self._cache_manager.load_cache() - self.decode_jwt(cache.access_token) - return cache.access_token - except jwt.ExpiredSignatureError: - cache = self._cache_manager.load_cache() - return self._refresh_auth_token(cache.refresh_token) - except Exception: + except FileNotFoundError: + return "" + except Exception as e: + LOGGER.warning("Discarding unreadable login cache: %s", e) self.delete_cache() return "" - @cached_property - def client(self): - return jwt.PyJWKClient(self._server_config.jwks_uri) - - def decode_jwt(self, json_web_token: str): - signing_key = self.client.get_signing_key_from_jwt(json_web_token) - return jwt.decode( - json_web_token, - signing_key.key, - algorithms=self._server_config.id_token_signing_alg_values_supported, - verify=True, + try: + expires_in = self._seconds_until_expiry(cache.access_token) + except jwt.InvalidTokenError as e: + # Not issued for this server's audience/issuer, refreshing cannot + # change that so the user must log in again + LOGGER.warning("Cached access token is not usable: %s", e) + self.delete_cache() + return "" + + if cache.access_token != rejected_token and ( + expires_in > REFRESH_MARGIN_SECONDS + ): + return cache.access_token + return self._refresh_auth_token(cache) + + def _seconds_until_expiry(self, access_token: str) -> float: + """ + Check the token belongs to this server's OIDC config and return how long + it has left. The signature is not verified: the server does that, and + verifying here would mean every request needs the OIDC provider to be + reachable (to fetch its signing keys). + """ + claims = jwt.decode( + access_token, + options={ + "verify_signature": False, + "verify_aud": True, + "verify_iss": True, + "verify_exp": False, + "require": ["exp"], + }, audience=self._server_config.client_audience, issuer=self._server_config.issuer, ) + return claims["exp"] - time.time() def logout(self) -> None: cache = self._cache_manager.load_cache() @@ -145,6 +229,7 @@ def logout(self) -> None: "id_token_hint": cache.id_token, "client_id": self._server_config.client_id, }, + timeout=REQUEST_TIMEOUT_SECONDS, ) response.raise_for_status() print("Logged out") @@ -154,30 +239,69 @@ def logout(self) -> None: f"to log out from the server.{e}" ) - def _refresh_auth_token(self, refresh_token: str) -> str: - response = requests.post( - self._server_config.token_endpoint, - data={ - "client_id": self._server_config.client_id, - "grant_type": "refresh_token", - "refresh_token": refresh_token, - }, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - ) + def _refresh_auth_token(self, cache: Cache) -> str: + """ + Exchange the refresh token for new tokens. + + The existing cache is left untouched until the refresh succeeds, and is + only deleted if the provider says the refresh token is no longer valid. + """ + issuer = self._server_config.issuer + try: + response = requests.post( + self._server_config.token_endpoint, + data={ + "client_id": self._server_config.client_id, + "grant_type": "refresh_token", + "refresh_token": cache.refresh_token, + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + timeout=REQUEST_TIMEOUT_SECONDS, + ) + except requests.RequestException as e: + raise AuthenticationServerUnavailableError( + "Your access token needs refreshing but the authentication server " + f"({issuer}) could not be reached: {e}. You are still logged in, " + "please try again once it is reachable." + ) from e + if response.status_code == HTTPStatus.OK: - token = response.json() + try: + token = response.json() + access_token = token["access_token"] + except (ValueError, KeyError, TypeError) as e: + raise AuthenticationServerUnavailableError( + "Your access token needs refreshing but the authentication " + f"server ({issuer}) returned an unexpected response. You are " + "still logged in, please try again later." + ) from e self._cache_manager.save_cache( Cache( oidc_config=self._server_config, - refresh_token=token["refresh_token"], - id_token=token["id_token"], - access_token=token["access_token"], + refresh_token=token.get("refresh_token", cache.refresh_token), + id_token=token.get("id_token", cache.id_token), + access_token=access_token, ), ) - return token["access_token"] - else: + return access_token + error, description = _oauth_error(response) + if error == "invalid_grant": + # The refresh token has expired, been revoked, or was issued to a + # different client. Only now is the login really gone. + LOGGER.warning( + "Your login session has expired or been revoked (%s), " + "please log in again with 'blueapi login'", + description, + ) self.delete_cache() return "" + # Anything else (5xx, invalid_client, ...) is a problem with the + # authentication server or its configuration, not with this login + raise AuthenticationServerUnavailableError( + "Your access token needs refreshing but the authentication server " + f"({issuer}) returned an error: {response.status_code} {description}. " + "You are still logged in, please try again later." + ) def poll_for_token( self, device_code: str, polling_interval: float, expires_in: float @@ -192,6 +316,7 @@ def poll_for_token( "client_id": self._server_config.client_id, }, headers={"Content-Type": "application/x-www-form-urlencoded"}, + timeout=REQUEST_TIMEOUT_SECONDS, ) if response.status_code == HTTPStatus.OK: return response.json() @@ -209,6 +334,7 @@ def start_device_flow(self): "scope": SCOPES, }, headers={"Content-Type": "application/x-www-form-urlencoded"}, + timeout=REQUEST_TIMEOUT_SECONDS, ) response.raise_for_status() @@ -235,7 +361,40 @@ def start_device_flow(self): print("Logged in and cached new token") +def _oauth_error(response: requests.Response) -> tuple[str | None, str]: + """The OAuth2 error code (RFC 6749 5.2) and a description of it""" + try: + body = response.json() + except ValueError: + return None, response.text or response.reason + if not isinstance(body, dict): + return None, response.text + error = body.get("error") + return error, body.get("error_description") or error or response.text + + +def migrate_legacy_cache(server_url: str, server_config: OIDCConfig) -> None: + """ + Move a login cached at the old shared default path to this server's cache + path, if it was made with this server's OIDC config. + """ + new_path = default_cache_path(server_url) + legacy_path = default_cache_path(None) + if new_path.exists() or not legacy_path.exists(): + return + try: + cache = SessionCacheManager(legacy_path).load_cache() + except Exception: + return + if SessionManager(cache.oidc_config, SessionCacheManager(legacy_path)).is_for( + server_config + ): + os.replace(legacy_path, new_path) + + class JWTAuth(AuthBase): + """Adds the bearer token (if any) to requests""" + def __init__(self, session_manager: SessionManager | None): self.token: str = ( session_manager.get_valid_access_token() if session_manager else "" diff --git a/src/blueapi/config.py b/src/blueapi/config.py index f65dd65f0..792789b3a 100644 --- a/src/blueapi/config.py +++ b/src/blueapi/config.py @@ -273,7 +273,7 @@ def _well_known_url(self) -> str: @cached_property def _config_from_oidc_url(self) -> dict[str, Any]: - response = requests.get(self._well_known_url) + response = requests.get(self._well_known_url, timeout=10) response.raise_for_status() return response.json() diff --git a/tests/conftest.py b/tests/conftest.py index 52c99ee13..87761361e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -240,7 +240,7 @@ def valid_token_with_jwt_invalid_audience(rsa_private_key: str) -> dict[str, Any @pytest.fixture def new_token(rsa_private_key: str) -> dict[str, Any]: - return _make_token("new_token", -100, +1700, rsa_private_key) + return _make_token("new_token", -100, +1700, rsa_private_key, jwt_access_token=True) @pytest.fixture @@ -309,6 +309,7 @@ def mock_authn_server( requests_mock.post( oidc_well_known["token_endpoint"], status=400, + json={"error": "invalid_grant", "error_description": "Token is not active"}, match=[ responses.matchers.urlencoded_params_matcher( { diff --git a/tests/system_tests/test_auth_token_cache.py b/tests/system_tests/test_auth_token_cache.py new file mode 100644 index 000000000..bd106d9fd --- /dev/null +++ b/tests/system_tests/test_auth_token_cache.py @@ -0,0 +1,186 @@ +""" +System tests for how the client caches and reuses OIDC tokens. + +The token cache stores the OIDC config of whichever server the user logged in +to. If the client is then used against a server with a different OIDC config +(e.g. the server's config changed, or another beamline's server is used) the +cached config is still trusted: `login()` reports "Logged in" because the +cached token validates against the *cached* config, but the server rejects it. + +These tests start their own blueapi server on a different port, using the +already running stack (keycloak etc.) for everything else. +""" + +import os +import subprocess +import sys +import time +from collections.abc import Generator +from pathlib import Path +from unittest.mock import patch + +import jwt +import pytest +import requests +import yaml +from tests.system_tests.test_blueapi_client import browser_login +from tests.system_tests.test_blueapi_system import _DATA_PATH, User + +from blueapi.client import BlueapiClient +from blueapi.client.session import AuthenticationServerUnavailableError, SessionManager + +# Port 8000 is the stack's blueapi, 8001 is taken by another stack service +STACK_SERVER_URL = "http://localhost:8000" +OTHER_SERVER_PORT = 8010 +OTHER_SERVER_URL = f"http://localhost:{OTHER_SERVER_PORT}" +KEYCLOAK_ISSUER = "http://localhost:8081/realms/master" + + +@pytest.fixture(scope="module") +def other_server(tmp_path_factory) -> Generator[str]: + """ + A second blueapi server, sharing the same keycloak, but using a different + OIDC client/audience to the stack's server on port 8000. + """ + config = { + "api": {"url": f"http://0.0.0.0:{OTHER_SERVER_PORT}"}, + "env": {"sources": [{"kind": "planFunctions", "module": "dodal.plans"}]}, + "oidc": { + "issuer": KEYCLOAK_ISSUER, + "client_id": "tiled-cli", + "client_audience": "tiled", + }, + } + config_file = tmp_path_factory.mktemp("other_server") / "config.yaml" + config_file.write_text(yaml.dump(config)) + proc = subprocess.Popen( + [sys.executable, "-m", "blueapi", "-c", str(config_file), "serve"] + ) + try: + for _ in range(120): + try: + if requests.get(f"{OTHER_SERVER_URL}/healthz").ok: + break + except requests.ConnectionError: + pass + if proc.poll() is not None: + raise RuntimeError("Second blueapi server exited during startup") + time.sleep(0.5) + else: + raise TimeoutError("Second blueapi server did not start") + yield OTHER_SERVER_URL + finally: + proc.terminate() + proc.wait(timeout=10) + + +@pytest.fixture +def cache_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """Use the default token cache location, but inside a temp dir""" + # Playwright finds its browsers via XDG_CACHE_HOME, so pin it first + original_cache = Path(os.environ.get("XDG_CACHE_HOME", "~/.cache")).expanduser() + monkeypatch.setenv( + "PLAYWRIGHT_BROWSERS_PATH", + os.environ.get( + "PLAYWRIGHT_BROWSERS_PATH", str(original_cache / "ms-playwright") + ), + ) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + return tmp_path / "cache" + + +def client_config(tmp_path: Path, api_url: str) -> Path: + """Write a CLI config for the given server, using the default token path""" + config = yaml.safe_load((_DATA_PATH / "config-cli.yaml").read_text()) + config["api"] = {"url": api_url} + config_file = tmp_path / f"config-{api_url.rsplit(':', 1)[-1]}.yaml" + config_file.write_text(yaml.dump(config)) + return config_file + + +def login(config_file: Path, user: User) -> BlueapiClient: + with patch( + "blueapi.client.session.webbrowser.open_new_tab", + side_effect=lambda url: browser_login(url, user), + ): + BlueapiClient.from_config_file(str(config_file)).login() + return BlueapiClient.from_config_file(str(config_file)) + + +def test_login_to_server_with_different_oidc_config( + tmp_path: Path, cache_home: Path, other_server: str +): + """ + Logged in to server A, then use server B (with a different OIDC config). + + Currently login() to B says "Logged in" using A's cached token/config, and + every request to B is then rejected. + """ + client_a = login(client_config(tmp_path, STACK_SERVER_URL), User.alice) + assert client_a.get_state() # sanity check: authenticated against A + + client_b = login(client_config(tmp_path, other_server), User.alice) + + # Fails with UnauthorisedAccessError: the token sent is for 'ixx-blueapi' + # but server B expects audience 'tiled' + assert client_b.get_state() + + # Logging in to B should not have logged us out of A + client_a = BlueapiClient.from_config_file( + str(client_config(tmp_path, STACK_SERVER_URL)) + ) + assert client_a.get_state() + + +def test_cache_kept_when_keycloak_unreachable( + tmp_path: Path, cache_home: Path, monkeypatch: pytest.MonkeyPatch +): + """ + With an offline_access refresh token the session is long lived, so failing + to reach keycloak should not throw the cache away. + """ + login(client_config(tmp_path, STACK_SERVER_URL), User.alice) + cached_files = list(cache_home.iterdir()) + assert cached_files + + # Route all HTTP traffic (requests and urllib) via a dead proxy so that + # keycloak (JWKS, token endpoint) cannot be reached + for var in ("no_proxy", "NO_PROXY"): + monkeypatch.delenv(var, raising=False) + for var in ("http_proxy", "HTTP_PROXY"): + monkeypatch.setenv(var, "http://127.0.0.1:9") + + # A token that is not about to expire is usable without keycloak + session_manager = SessionManager.from_cache(None, STACK_SERVER_URL) + token = session_manager.get_valid_access_token() + assert token + + # Needing a refresh while keycloak is unreachable is a helpful error... + with pytest.raises( + AuthenticationServerUnavailableError, match="You are still logged in" + ): + session_manager.get_valid_access_token(rejected_token=token) + + # ...and not a silent logout + assert all(f.exists() for f in cached_files) + + +def test_access_token_refreshed_before_it_expires(tmp_path: Path, cache_home: Path): + """ + A token that is about to expire should be refreshed, otherwise it can + expire between the client checking it and the server validating it. + """ + login(client_config(tmp_path, STACK_SERVER_URL), User.alice) + session_manager = SessionManager.from_cache(None, STACK_SERVER_URL) + token = session_manager.get_valid_access_token() + expiry = jwt.decode(token, options={"verify_signature": False})["exp"] + + # Wait until the token has less than 30 seconds left (keycloak gives 60s) + time.sleep(max(0.0, expiry - time.time() - 25)) + + new_token = session_manager.get_valid_access_token() + assert new_token != token + assert ( + jwt.decode(new_token, options={"verify_signature": False})["exp"] - time.time() + > 30 + ) diff --git a/tests/unit_tests/cli/test_cli.py b/tests/unit_tests/cli/test_cli.py index f858e6ac0..feff5d639 100644 --- a/tests/unit_tests/cli/test_cli.py +++ b/tests/unit_tests/cli/test_cli.py @@ -36,6 +36,7 @@ ) from blueapi.config import ( ApplicationConfig, + OIDCConfig, ScratchConfig, ScratchRepository, generate_config_schema, @@ -1518,3 +1519,86 @@ def test_run_ws_runs_blocking_plan(mock_client: Mock, runner: CliRunner): TaskRequest(name="name", params={}, instrument_session="cm12345-1"), ) assert res.exit_code == 0 + + +def test_login_when_oidc_provider_unreachable_keeps_cache( + runner: CliRunner, + config_with_auth: str, + oidc_config: OIDCConfig, + cached_valid_refresh: Path, +): + with responses.RequestsMock(assert_all_requests_are_fired=False) as mock: + mock.get("http://localhost:8000/config/oidc", json=oidc_config.model_dump()) + mock.get( + "https://auth.example.com/realms/master/.well-known/openid-configuration", + body=ConnectionError("down"), + ) + result = runner.invoke(main, ["-c", config_with_auth, "login"]) + + assert result.exit_code == 1 + assert "could not be reached" in result.output + assert "You are still logged in" in result.output + assert cached_valid_refresh.exists() + + +def test_login_again_when_server_oidc_config_changed( + runner: CliRunner, + config_with_auth: str, + mock_authn_server: responses.RequestsMock, + oidc_config: OIDCConfig, + cached_valid_token: Path, +): + mock_authn_server.replace( + responses.GET, + "http://localhost:8000/config/oidc", + json=oidc_config.model_copy( + update={"client_audience": "new-audience"} + ).model_dump(), + ) + with patch("webbrowser.open_new_tab", return_value=False): + result = runner.invoke(main, ["-c", config_with_auth, "login"]) + assert "authentication config has changed" in result.output + assert "Please login from this URL" in result.output + + +def test_command_when_oidc_provider_unreachable_is_clean_error( + runner: CliRunner, + config_with_auth: str, + cached_valid_refresh: Path, +): + with responses.RequestsMock(assert_all_requests_are_fired=False) as mock: + mock.get( + "https://auth.example.com/realms/master/.well-known/openid-configuration", + body=ConnectionError("down"), + ) + result = runner.invoke(main, ["-c", config_with_auth, "controller", "plans"]) + + assert result.exit_code == 1 + assert result.exception is None or isinstance(result.exception, SystemExit) + assert result.output.startswith("Error: Your access token needs refreshing") + assert cached_valid_refresh.exists() + + +def test_login_when_blueapi_server_down_is_clean_error( + runner: CliRunner, config_with_auth: str, cached_valid_token: Path +): + with responses.RequestsMock(assert_all_requests_are_fired=False) as mock: + mock.get("http://localhost:8000/config/oidc", body=ConnectionError("refused")) + result = runner.invoke(main, ["-c", config_with_auth, "login"]) + assert result.exit_code == 1 + assert result.output == "Error: Failed to establish connection to blueapi server.\n" + + +def test_logout_also_logs_out_legacy_cache( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + mock_authn_server: responses.RequestsMock, + cached_valid_token: Path, +): + # cached_valid_token is at tmp_path/blueapi_cache, i.e. the legacy default + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + result = runner.invoke(main, ["logout"]) + assert result.exit_code == 0 + assert "Logged out" in result.output + assert not cached_valid_token.exists() diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py index e7e31f930..b8e1df91b 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -767,12 +767,13 @@ def test_client_login_existing_login(mock_session_manager: Mock, client: Blueapi mock_session_manager.from_cache.assert_called_once() mock_session_manager.from_cache().get_valid_access_token.assert_called_once() + mock_session_manager.assert_not_called() @patch("blueapi.client.client.SessionManager") def test_client_new_login(mock_session_manager: Mock, client: BlueapiClient): manager = Mock() - manager.get_valid_access_token.side_effect = ValueError("No existing token") + manager.get_valid_access_token.return_value = "" mock_session_manager.from_cache.return_value = manager @@ -782,6 +783,28 @@ def test_client_new_login(mock_session_manager: Mock, client: BlueapiClient): mock_session_manager.return_value.start_device_flow.assert_called_once() +@patch("blueapi.client.client.SessionManager") +def test_client_login_without_cache(mock_session_manager: Mock, client: BlueapiClient): + mock_session_manager.from_cache.side_effect = FileNotFoundError() + + client.login() + + mock_session_manager.return_value.start_device_flow.assert_called_once() + + +@patch("blueapi.client.client.SessionManager") +def test_client_login_again_if_server_oidc_config_changed( + mock_session_manager: Mock, client: BlueapiClient +): + manager = mock_session_manager.from_cache.return_value + manager.is_for.return_value = False + + client.login() + + manager.get_valid_access_token.assert_not_called() + mock_session_manager.return_value.start_device_flow.assert_called_once() + + @patch("blueapi.client.client.SessionManager") def test_client_login_no_oidc( mock_session_manager: Mock, mock_rest: Mock, client: BlueapiClient diff --git a/tests/unit_tests/client/test_rest.py b/tests/unit_tests/client/test_rest.py index a36ee71a3..68d4d2ec1 100644 --- a/tests/unit_tests/client/test_rest.py +++ b/tests/unit_tests/client/test_rest.py @@ -28,7 +28,11 @@ _create_task_exceptions, _exception, ) -from blueapi.client.session import SessionCacheManager, SessionManager +from blueapi.client.session import ( + AuthenticationServerUnavailableError, + SessionCacheManager, + SessionManager, +) from blueapi.config import OIDCConfig from blueapi.service.model import ( DeviceModel, @@ -230,6 +234,7 @@ def test_refresh_if_signature_expired( rest_with_auth: BlueapiRestClient, mock_authn_server: responses.RequestsMock, cached_valid_refresh: Path, + new_token: dict[str, Any], ): environment_id = uuid.uuid4() mock_authn_server.stop() # Cannot use multiple RequestsMock context manager @@ -248,7 +253,10 @@ def test_refresh_if_signature_expired( ) calls = mock_get_env.calls assert len(calls) == 1 - assert calls[0].request.headers["Authorization"] == "Bearer new_token" + assert ( + calls[0].request.headers["Authorization"] + == f"Bearer {new_token['access_token']}" + ) def test_parameter_error_field(): @@ -590,3 +598,172 @@ def test_run_blocking_connect_error( conn = rest.run_blocking(TASK_REQUEST) with pytest.raises(exp_err): next(iter(conn)) + + +def test_rejected_token_is_refreshed_and_request_retried( + rest_with_auth: BlueapiRestClient, + mock_authn_server: responses.RequestsMock, + cached_valid_token: Path, + cached_valid_token_value: str, + new_token: dict[str, Any], +): + environment_id = uuid.uuid4() + mock_authn_server.stop() + unauthorised = mock_authn_server.get( + "http://localhost:8000/environment", + status=401, + match=[ + responses.matchers.header_matcher( + {"Authorization": f"Bearer {cached_valid_token_value}"} + ) + ], + ) + authorised = mock_authn_server.get( + "http://localhost:8000/environment", + json=EnvironmentResponse( + environment_id=environment_id, initialized=True + ).model_dump(mode="json"), + match=[ + responses.matchers.header_matcher( + {"Authorization": f"Bearer {new_token['access_token']}"} + ) + ], + ) + with mock_authn_server: + result = rest_with_auth.get_environment() + assert result.environment_id == environment_id + assert unauthorised.call_count == 1 + assert authorised.call_count == 1 + + +def test_request_only_retried_once( + rest_with_auth: BlueapiRestClient, + mock_authn_server: responses.RequestsMock, + cached_valid_token: Path, +): + mock_authn_server.stop() + unauthorised = mock_authn_server.get( + "http://localhost:8000/environment", status=401 + ) + with mock_authn_server, pytest.raises(UnauthorisedAccessError): + rest_with_auth.get_environment() + assert unauthorised.call_count == 2 + + +def _ws_with_one_update() -> MagicMock: + ws = MagicMock() + ws.__enter__.return_value.__iter__.return_value = iter( + ['{"kind": "update", "data": {"name": "start", "doc":{}, "task_id":"t_uid"}}'] + ) + return ws + + +@patch("blueapi.client.rest.connect") +def test_run_blocking_no_empty_bearer_when_logged_out( + mock_connect: Mock, oidc_config: OIDCConfig, token_cache_file: Path +): + # Session manager exists but has no cached login + rest = BlueapiRestClient( + session_manager=SessionManager( + oidc_config, SessionCacheManager(token_cache_file) + ) + ) + mock_connect.return_value = _ws_with_one_update() + next(iter(rest.run_blocking(TASK_REQUEST))) + assert "Authorization" not in mock_connect.call_args.kwargs["additional_headers"] + + +@patch("blueapi.client.rest.connect") +def test_run_blocking_retries_once_with_refreshed_token( + mock_connect: Mock, + rest_with_auth: BlueapiRestClient, + mock_authn_server: responses.RequestsMock, + cached_valid_token: Path, + cached_valid_token_value: str, + new_token: dict[str, Any], +): + mock_connect.side_effect = [ + InvalidStatus( + response=Response(status_code=401, reason_phrase="", headers=Headers()) + ), + _ws_with_one_update(), + ] + next(iter(rest_with_auth.run_blocking(TASK_REQUEST))) + sent = [ + c.kwargs["additional_headers"]["Authorization"] + for c in mock_connect.call_args_list + ] + assert sent == [ + f"Bearer {cached_valid_token_value}", + f"Bearer {new_token['access_token']}", + ] + + +@patch("blueapi.client.rest.connect") +def test_run_blocking_auth_server_unavailable_is_raised( + mock_connect: Mock, rest: BlueapiRestClient +): + rest.session_manager = Mock() + rest.session_manager.get_valid_access_token.side_effect = ( + AuthenticationServerUnavailableError("down") + ) + with pytest.raises(AuthenticationServerUnavailableError): + next(iter(rest.run_blocking(TASK_REQUEST))) + mock_connect.assert_not_called() + + +def test_oidc_config_fetched_without_authentication( + rest: BlueapiRestClient, oidc_config: OIDCConfig +): + # e.g. a stale login whose OIDC provider is gone must not block this + rest.session_manager = Mock() + rest.session_manager.get_valid_access_token.side_effect = ( + AuthenticationServerUnavailableError("old provider is gone") + ) + with responses.RequestsMock() as mock: + call = mock.get( + "http://localhost:8000/config/oidc", json=oidc_config.model_dump() + ) + assert rest.get_oidc_config() == oidc_config + assert "Authorization" not in call.calls[0].request.headers + rest.session_manager.get_valid_access_token.assert_not_called() + + +def test_retry_uses_session_and_follows_redirects( + rest_with_auth: BlueapiRestClient, + mock_authn_server: responses.RequestsMock, + cached_valid_token: Path, + cached_valid_token_value: str, + new_token: dict[str, Any], +): + environment_id = uuid.uuid4() + mock_authn_server.stop() + mock_authn_server.get( + "http://localhost:8000/environment", + status=401, + match=[ + responses.matchers.header_matcher( + {"Authorization": f"Bearer {cached_valid_token_value}"} + ) + ], + ) + mock_authn_server.add( + responses.GET, + "http://localhost:8000/environment", + status=307, + headers={"Location": "http://localhost:8000/environment/"}, + match=[ + responses.matchers.header_matcher( + {"Authorization": f"Bearer {new_token['access_token']}"} + ) + ], + ) + mock_authn_server.get( + "http://localhost:8000/environment/", + json=EnvironmentResponse( + environment_id=environment_id, initialized=True + ).model_dump(mode="json"), + ) + with mock_authn_server: + result = rest_with_auth.get_environment() + assert result.environment_id == environment_id diff --git a/tests/unit_tests/client/test_session.py b/tests/unit_tests/client/test_session.py index f54334fdc..7b2c99935 100644 --- a/tests/unit_tests/client/test_session.py +++ b/tests/unit_tests/client/test_session.py @@ -1,14 +1,24 @@ import os +import stat from pathlib import Path from typing import Any from unittest.mock import patch import pytest +import requests import responses from starlette.status import HTTP_403_FORBIDDEN +from tests.conftest import ISSUER, _make_token -from blueapi.client.session import SessionCacheManager, SessionManager +from blueapi.client.session import ( + AuthenticationServerUnavailableError, + SessionCacheManager, + SessionManager, + default_cache_path, + migrate_legacy_cache, +) from blueapi.config import OIDCConfig +from blueapi.service.model import Cache @pytest.fixture @@ -42,9 +52,10 @@ def test_refresh_auth_token( mock_authn_server: responses.RequestsMock, session_manager: SessionManager, cached_valid_refresh: Path, + new_token: dict[str, Any], ): token = session_manager.get_valid_access_token() - assert token == "new_token" + assert token == new_token["access_token"] def test_get_empty_token_if_no_cache(session_manager: SessionManager): @@ -118,3 +129,282 @@ def test_session_cache_manager_returns_writable_file_path(tmp_path): Path(cache._file_path).touch() assert os.path.isfile(cache._file_path) assert cache._file_path == f"{tmp_path}/blueapi_cache" + + +@pytest.mark.parametrize( + "server_url,file_name", + [ + (None, "blueapi_cache"), + ( + "https://i22-blueapi.diamond.ac.uk", + "blueapi_cache.i22-blueapi.diamond.ac.uk", + ), + ("http://localhost:8000/", "blueapi_cache.localhost_8000"), + ("http://localhost:8010", "blueapi_cache.localhost_8010"), + ], +) +def test_default_cache_path_is_per_server( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + server_url: str | None, + file_name: str, +): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + assert default_cache_path(server_url) == tmp_path / file_name + assert SessionCacheManager(None, server_url)._token_path == tmp_path / file_name + + +def test_explicit_token_path_ignores_server_url(tmp_path: Path): + path = tmp_path / "my_token" + assert SessionCacheManager(path, "http://localhost:8000")._token_path == path + + +def test_save_cache_is_private_and_leaves_no_temp_files( + session_manager: SessionManager, cached_valid_token: Path, tmp_path: Path +): + cache_manager = SessionCacheManager(cached_valid_token) + cache = cache_manager.load_cache() + cache_manager.save_cache(cache.model_copy(update={"access_token": "updated"})) + + assert cache_manager.load_cache().access_token == "updated" + assert stat.S_IMODE(cached_valid_token.stat().st_mode) == 0o600 + assert list(tmp_path.iterdir()) == [cached_valid_token] + + +def test_valid_token_does_not_need_oidc_provider( + oidc_config: OIDCConfig, cached_valid_token: Path, cached_valid_token_value: str +): + # No mock_authn_server: any request to the OIDC provider would fail + session_manager = SessionManager( + oidc_config, SessionCacheManager(cached_valid_token) + ) + assert session_manager.get_valid_access_token() == cached_valid_token_value + + +def _cache_token( + path: Path, oidc_config: OIDCConfig, rsa_private_key: str, expires_in: float +) -> str: + token = _make_token("t", -60, expires_in, rsa_private_key, jwt_access_token=True) + SessionCacheManager(path).save_cache( + Cache( + oidc_config=oidc_config, + access_token=token["access_token"], + refresh_token=token["refresh_token"], + id_token=token["id_token"], + ) + ) + return token["access_token"] + + +def test_token_about_to_expire_is_refreshed( + session_manager: SessionManager, + auth_token_path: Path, + oidc_config: OIDCConfig, + rsa_private_key: str, + new_token: dict[str, Any], +): + _cache_token(auth_token_path, oidc_config, rsa_private_key, expires_in=20) + assert session_manager.get_valid_access_token() == new_token["access_token"] + + +def test_token_not_about_to_expire_is_not_refreshed( + session_manager: SessionManager, + auth_token_path: Path, + oidc_config: OIDCConfig, + rsa_private_key: str, +): + token = _cache_token(auth_token_path, oidc_config, rsa_private_key, expires_in=60) + assert session_manager.get_valid_access_token() == token + + +def test_rejected_token_is_refreshed( + session_manager: SessionManager, + cached_valid_token: Path, + cached_valid_token_value: str, + new_token: dict[str, Any], +): + token = session_manager.get_valid_access_token( + rejected_token=cached_valid_token_value + ) + assert token == new_token["access_token"] + + +def test_rejected_token_already_replaced_by_another_process( + oidc_config: OIDCConfig, + cached_valid_token: Path, + cached_valid_token_value: str, +): + # No mock_authn_server: a refresh would fail + session_manager = SessionManager( + oidc_config, SessionCacheManager(cached_valid_token) + ) + token = session_manager.get_valid_access_token(rejected_token="an_older_token") + assert token == cached_valid_token_value + + +def test_cache_kept_if_oidc_provider_unreachable( + oidc_config: OIDCConfig, cached_valid_refresh: Path +): + session_manager = SessionManager( + oidc_config, SessionCacheManager(cached_valid_refresh) + ) + with ( + responses.RequestsMock() as mock, + pytest.raises( + AuthenticationServerUnavailableError, match="could not be reached" + ), + ): + mock.get( + ISSUER + "/.well-known/openid-configuration", + body=requests.ConnectionError("down"), + ) + session_manager.get_valid_access_token() + assert cached_valid_refresh.exists() + + +def test_cache_kept_if_oidc_provider_errors( + oidc_config: OIDCConfig, + oidc_well_known: dict[str, Any], + cached_valid_refresh: Path, +): + session_manager = SessionManager( + oidc_config, SessionCacheManager(cached_valid_refresh) + ) + with ( + responses.RequestsMock() as mock, + pytest.raises(AuthenticationServerUnavailableError, match="503"), + ): + mock.get(ISSUER + "/.well-known/openid-configuration", json=oidc_well_known) + mock.post(oidc_well_known["token_endpoint"], status=503) + session_manager.get_valid_access_token() + assert cached_valid_refresh.exists() + + +def test_cache_deleted_if_token_for_different_audience( + session_manager: SessionManager, cache_with_invalid_audience: Path +): + assert session_manager.get_valid_access_token() == "" + assert not cache_with_invalid_audience.exists() + + +def test_is_for(oidc_config: OIDCConfig, session_manager: SessionManager): + assert session_manager.is_for(oidc_config) + assert not session_manager.is_for( + oidc_config.model_copy(update={"client_audience": "other"}) + ) + + +def test_cache_kept_if_refresh_returns_unexpected_body( + oidc_config: OIDCConfig, + oidc_well_known: dict[str, Any], + cached_valid_refresh: Path, +): + session_manager = SessionManager( + oidc_config, SessionCacheManager(cached_valid_refresh) + ) + with ( + responses.RequestsMock() as mock, + pytest.raises(AuthenticationServerUnavailableError, match="unexpected"), + ): + mock.get(ISSUER + "/.well-known/openid-configuration", json=oidc_well_known) + mock.post(oidc_well_known["token_endpoint"], body="proxy login") + session_manager.get_valid_access_token() + assert cached_valid_refresh.exists() + + +@pytest.mark.parametrize( + "status,error", [(400, "invalid_client"), (401, "unauthorized_client"), (400, None)] +) +def test_cache_kept_if_refresh_rejected_for_other_reasons( + oidc_config: OIDCConfig, + oidc_well_known: dict[str, Any], + cached_valid_refresh: Path, + status: int, + error: str | None, +): + session_manager = SessionManager( + oidc_config, SessionCacheManager(cached_valid_refresh) + ) + with ( + responses.RequestsMock() as mock, + pytest.raises(AuthenticationServerUnavailableError, match=str(status)), + ): + mock.get(ISSUER + "/.well-known/openid-configuration", json=oidc_well_known) + mock.post( + oidc_well_known["token_endpoint"], + status=status, + json={"error": error} if error else None, + ) + session_manager.get_valid_access_token() + assert cached_valid_refresh.exists() + + +def test_requests_to_oidc_provider_have_timeouts( + session_manager: SessionManager, + mock_authn_server: responses.RequestsMock, + cached_valid_refresh: Path, + device_code: str, +): + session_manager.get_valid_access_token() # refresh + session_manager.poll_for_token(device_code, 1, 2) + session_manager.logout() + assert mock_authn_server.calls + for call in mock_authn_server.calls: + timeout = call.request.req_kwargs["timeout"] # type: ignore[attr-defined] + assert timeout, call.request.url + + +def _write_cache(path: Path, oidc_config: OIDCConfig) -> None: + SessionCacheManager(path).save_cache( + Cache( + oidc_config=oidc_config, + access_token="a", + refresh_token="r", + id_token="i", + ) + ) + + +def test_legacy_cache_migrated_if_for_this_server( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, oidc_config: OIDCConfig +): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + _write_cache(tmp_path / "blueapi_cache", oidc_config) + + migrate_legacy_cache("http://localhost:8000", oidc_config) + + assert not (tmp_path / "blueapi_cache").exists() + migrated = SessionCacheManager(tmp_path / "blueapi_cache.localhost_8000") + assert migrated.load_cache().refresh_token == "r" + + +def test_legacy_cache_not_migrated_if_for_another_server( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, oidc_config: OIDCConfig +): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + _write_cache(tmp_path / "blueapi_cache", oidc_config) + + migrate_legacy_cache( + "http://localhost:8000", + oidc_config.model_copy(update={"client_audience": "other"}), + ) + + assert (tmp_path / "blueapi_cache").exists() + assert not (tmp_path / "blueapi_cache.localhost_8000").exists() + + +def test_legacy_cache_does_not_replace_existing_login( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, oidc_config: OIDCConfig +): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + _write_cache(tmp_path / "blueapi_cache", oidc_config) + SessionCacheManager(tmp_path / "blueapi_cache.localhost_8000").save_cache( + Cache( + oidc_config=oidc_config, access_token="a", refresh_token="new", id_token="i" + ) + ) + + migrate_legacy_cache("http://localhost:8000", oidc_config) + + current = SessionCacheManager(tmp_path / "blueapi_cache.localhost_8000") + assert current.load_cache().refresh_token == "new"