From 019be3593f9824d484e3d50135fbac8de5eb9b69 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 21 May 2026 11:42:59 +0100 Subject: [PATCH 01/11] feat: Add OpaConfig for authorization configuration --- helm/blueapi/config_schema.json | 27 +++++++++++++++++++++++++++ helm/blueapi/values.schema.json | 26 ++++++++++++++++++++++++++ pyproject.toml | 1 + src/blueapi/config.py | 6 ++++++ tests/unit_tests/test_config.py | 4 ++++ uv.lock | 2 ++ 6 files changed, 66 insertions(+) diff --git a/helm/blueapi/config_schema.json b/helm/blueapi/config_schema.json index b5d0c9bf3d..1ad4e82bfe 100644 --- a/helm/blueapi/config_schema.json +++ b/helm/blueapi/config_schema.json @@ -330,6 +330,22 @@ "type": "object", "$id": "OIDCConfig" }, + "OpaConfig": { + "additionalProperties": false, + "properties": { + "root": { + "default": "http://localhost:8181/", + "format": "uri", + "maxLength": 2083, + "minLength": 1, + "title": "Root", + "type": "string" + } + }, + "title": "OpaConfig", + "type": "object", + "$id": "OpaConfig" + }, "PlanSource": { "additionalProperties": false, "properties": { @@ -612,6 +628,17 @@ } ], "default": null + }, + "opa": { + "anyOf": [ + { + "$ref": "OpaConfig" + }, + { + "type": "null" + } + ], + "default": null } }, "title": "ApplicationConfig", diff --git a/helm/blueapi/values.schema.json b/helm/blueapi/values.schema.json index 74deedadb2..6de532cc88 100644 --- a/helm/blueapi/values.schema.json +++ b/helm/blueapi/values.schema.json @@ -751,6 +751,22 @@ }, "additionalProperties": false }, + "OpaConfig": { + "$id": "OpaConfig", + "title": "OpaConfig", + "type": "object", + "properties": { + "root": { + "title": "Root", + "default": "http://localhost:8181/", + "type": "string", + "format": "uri", + "maxLength": 2083, + "minLength": 1 + } + }, + "additionalProperties": false + }, "PlanSource": { "$id": "PlanSource", "title": "PlanSource", @@ -1011,6 +1027,16 @@ } ] }, + "opa": { + "anyOf": [ + { + "$ref": "OpaConfig" + }, + { + "type": "null" + } + ] + }, "scratch": { "anyOf": [ { diff --git a/pyproject.toml b/pyproject.toml index 659779994a..9f4231c76c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ "tomlkit", "graypy>=2.1.0", "httpx>=0.28.1", + "aiohttp>=3.13.5", ] dynamic = ["version"] license.file = "LICENSE" diff --git a/src/blueapi/config.py b/src/blueapi/config.py index 83d6d70211..4c2431cf1e 100644 --- a/src/blueapi/config.py +++ b/src/blueapi/config.py @@ -296,6 +296,10 @@ class Tag(StrEnum): META = "Meta" +class OpaConfig(BlueapiBaseModel): + root: HttpUrl = HttpUrl("http://localhost:8181") + + class ApplicationConfig(BlueapiBaseModel): """ Config for the worker application as a whole. Root of @@ -335,6 +339,7 @@ class ApplicationConfig(BlueapiBaseModel): oidc: OIDCConfig | None = None auth_token_path: Path | None = None numtracker: NumtrackerConfig | None = None + opa: OpaConfig | None = None def __eq__(self, other: object) -> bool: if isinstance(other, ApplicationConfig): @@ -343,6 +348,7 @@ def __eq__(self, other: object) -> bool: & (self.env == other.env) & (self.logging == other.logging) & (self.api == other.api) + & (self.opa == other.opa) ) return False diff --git a/tests/unit_tests/test_config.py b/tests/unit_tests/test_config.py index 30fe551c4c..7ce98f6fe6 100644 --- a/tests/unit_tests/test_config.py +++ b/tests/unit_tests/test_config.py @@ -337,6 +337,9 @@ def test_config_yaml_parsed(temp_yaml_config_file): } ], }, + "opa": { + "root": "http://opa.example.com/", + }, }, { "stomp": { @@ -392,6 +395,7 @@ def test_config_yaml_parsed(temp_yaml_config_file): } ], }, + "opa": None, }, ], indirect=True, diff --git a/uv.lock b/uv.lock index b2d10b5bd5..eab94a8670 100644 --- a/uv.lock +++ b/uv.lock @@ -420,6 +420,7 @@ name = "blueapi" source = { editable = "." } dependencies = [ { name = "aioca" }, + { name = "aiohttp" }, { name = "bluesky", extra = ["plotting"] }, { name = "bluesky-stomp" }, { name = "click" }, @@ -481,6 +482,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "aioca" }, + { name = "aiohttp", specifier = ">=3.13.5" }, { name = "bluesky", extras = ["plotting"], specifier = ">=1.14.0" }, { name = "bluesky-stomp", specifier = ">=0.1.6" }, { name = "click", specifier = ">=8.2.0" }, From 1af9183d9ed0fc3ea73087c42ffa861b2425c5c1 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 21 May 2026 11:52:04 +0100 Subject: [PATCH 02/11] Add OpaClient to wrap OPA interactions --- src/blueapi/service/authorization.py | 61 ++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 src/blueapi/service/authorization.py diff --git a/src/blueapi/service/authorization.py b/src/blueapi/service/authorization.py new file mode 100644 index 0000000000..8e260cbd53 --- /dev/null +++ b/src/blueapi/service/authorization.py @@ -0,0 +1,61 @@ +import logging +from collections.abc import Mapping +from contextlib import AbstractAsyncContextManager, aclosing, nullcontext +from typing import Any, Self + +import aiohttp +from aiohttp import ClientSession + +from blueapi.config import OpaConfig + +LOGGER = logging.getLogger(__name__) + + +class OpaClient: + client: aiohttp.ClientSession + + def __init__(self, instrument: str, config: OpaConfig): + LOGGER.info("Creating OpaClient for %s with config %s", instrument, config) + self._instrument = instrument + self._conf = config + self._session = ClientSession(base_url=config.root.encoded_string()) + + async def aclose(self): + LOGGER.info("Closing OPA session") + await self._session.close() + + async def _call_opa(self, endpoint, data: Mapping[str, Any]) -> bool: + try: + resp = await self._session.post( + endpoint, + json={ + "input": { + "beamline": self._instrument, + "audience": "account", + **data, + } + }, + ) + return (await resp.json())["result"] + except Exception: + LOGGER.exception("Failed to run check", exc_info=True) + raise + + @classmethod + def for_config( + cls, instrument: str, config: OpaConfig | None + ) -> AbstractAsyncContextManager[Self | None]: + if config: + return aclosing(cls(instrument, config)) + LOGGER.info("No OPA config provided - not creating OpaClient") + return nullcontext() + + + +class OpaUserClient: + client: OpaClient + token: str + + def __init__(self, client: OpaClient, token: str): + self.client = client + self.token = token From 6453b30b89da2833ea31662feaa189eed1ab968c Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 21 May 2026 13:32:18 +0100 Subject: [PATCH 03/11] Create OpaClient as part of server lifecycle --- src/blueapi/service/main.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index a53c46885a..b6ece56091 100644 --- a/src/blueapi/service/main.py +++ b/src/blueapi/service/main.py @@ -40,6 +40,7 @@ from blueapi.worker import TrackableTask, WorkerState from blueapi.worker.event import TaskStatusEnum +from .authorization import OpaClient from .model import ( DeviceModel, DeviceResponse, @@ -92,8 +93,13 @@ def teardown_runner(): def lifespan(config: ApplicationConfig): @asynccontextmanager async def inner(app: FastAPI): + if not (meta := config.env.metadata): + raise ValueError("Instrument name is required in metadata") + setup_runner(config) - yield + async with OpaClient.for_config(meta.instrument, config.opa) as opa: + app.state.authz = opa + yield teardown_runner() return inner From edcee3dcc18f6ff072814bb32ea626b8182589f1 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Wed, 27 May 2026 14:54:30 +0100 Subject: [PATCH 04/11] Move instrument requirement into OpaClient init --- src/blueapi/service/authorization.py | 4 +++- src/blueapi/service/main.py | 6 ++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/blueapi/service/authorization.py b/src/blueapi/service/authorization.py index 8e260cbd53..b27c304c97 100644 --- a/src/blueapi/service/authorization.py +++ b/src/blueapi/service/authorization.py @@ -43,9 +43,11 @@ async def _call_opa(self, endpoint, data: Mapping[str, Any]) -> bool: @classmethod def for_config( - cls, instrument: str, config: OpaConfig | None + cls, instrument: str | None, config: OpaConfig | None ) -> AbstractAsyncContextManager[Self | None]: if config: + if not instrument: + raise ValueError("Instrument name is required for OPA client") return aclosing(cls(instrument, config)) LOGGER.info("No OPA config provided - not creating OpaClient") return nullcontext() diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index b6ece56091..1b64630601 100644 --- a/src/blueapi/service/main.py +++ b/src/blueapi/service/main.py @@ -93,11 +93,9 @@ def teardown_runner(): def lifespan(config: ApplicationConfig): @asynccontextmanager async def inner(app: FastAPI): - if not (meta := config.env.metadata): - raise ValueError("Instrument name is required in metadata") - + meta = config.env.metadata setup_runner(config) - async with OpaClient.for_config(meta.instrument, config.opa) as opa: + async with OpaClient.for_config(meta and meta.instrument, config.opa) as opa: app.state.authz = opa yield teardown_runner() From 8f83ef8fe47edf754662b2134506582ee644ad43 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 28 May 2026 17:21:00 +0100 Subject: [PATCH 05/11] Add tests for OpaClient --- .../unit_tests/service/test_authorization.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 tests/unit_tests/service/test_authorization.py diff --git a/tests/unit_tests/service/test_authorization.py b/tests/unit_tests/service/test_authorization.py new file mode 100644 index 0000000000..aa4417d500 --- /dev/null +++ b/tests/unit_tests/service/test_authorization.py @@ -0,0 +1,48 @@ +from unittest.mock import MagicMock, patch + +import pytest +from pydantic import HttpUrl + +from blueapi.config import OpaConfig +from blueapi.service.authorization import ( + OpaClient, +) + +# Reusable client patch decorator +patch_client_session = patch( + "blueapi.service.authorization.ClientSession", + name="mock_client_session", + spec=True, +) + + +@pytest.fixture(scope="module") +def opa_config() -> OpaConfig: + return OpaConfig( + root=HttpUrl("http://auth.example.com"), + ) + + +@patch_client_session +async def test_session_closed(session: MagicMock, opa_config: OpaConfig): + async with OpaClient.for_config("p45", opa_config): + pass + session().close.assert_called_once() + + +@patch_client_session +async def test_opa_client_for_config(session: MagicMock, opa_config: OpaConfig): + async with OpaClient.for_config("p45", opa_config) as opa: + assert opa is not None + session.assert_called_once_with(base_url="http://auth.example.com/") + + +@pytest.mark.parametrize("instrument", [None, "p99"]) +async def test_opa_client_without_config(instrument: str | None): + async with OpaClient.for_config(instrument, None) as opa: + assert opa is None + + +async def test_opa_fails_without_instrument(opa_config: OpaConfig): + with pytest.raises(ValueError, match="Instrument name is required"): + OpaClient.for_config(None, opa_config) From 6dbe33e1cc5ac30db1cb5775e6529acf8fa63573 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 21 May 2026 13:46:19 +0100 Subject: [PATCH 06/11] refactor: Move auth extractors into authentication module And split the auth check into two so that other methods can access the raw bearer token if required. --- src/blueapi/service/authentication.py | 56 ++++++++++++++++++- src/blueapi/service/main.py | 31 +--------- .../unit_tests/service/test_authentication.py | 10 ++-- 3 files changed, 63 insertions(+), 34 deletions(-) diff --git a/src/blueapi/service/authentication.py b/src/blueapi/service/authentication.py index b107f7b2b2..051dc27f19 100644 --- a/src/blueapi/service/authentication.py +++ b/src/blueapi/service/authentication.py @@ -9,13 +9,16 @@ from functools import cached_property from http import HTTPStatus from pathlib import Path -from typing import Any, cast +from typing import Annotated, Any, cast import httpx import jwt import requests +from fastapi import Depends, HTTPException, Request +from fastapi.security.utils import get_authorization_scheme_param from pydantic import TypeAdapter from requests.auth import AuthBase +from starlette.status import HTTP_401_UNAUTHORIZED from blueapi.config import OIDCConfig, ServiceAccount from blueapi.service.model import Cache @@ -272,3 +275,54 @@ def get_access_token(self): def sync_auth_flow(self, request): request.headers["Authorization"] = f"Bearer {self.get_access_token()}" yield request + + +def unchecked_bearer_token(req: Request) -> str | None: + """Get bearer token value from authorization header""" + auth = req.headers.get("Authorization") + scheme, param = get_authorization_scheme_param(auth) + if scheme.casefold() != "bearer": + return None + return param.strip() + + +UncheckedBearerToken = Annotated[str | None, Depends(unchecked_bearer_token)] + + +def build_access_token_check(config: OIDCConfig): + """ + Create a function to validate the bearer token of requests + + The returned function should be used via fastAPI's 'Depends' mechanism to + ensure users are authenticated + """ + jwkclient = jwt.PyJWKClient(config.jwks_uri) + + def validate_bearer_token( + request: Request, token: Annotated[str | None, Depends(unchecked_bearer_token)] + ): + """Check that a bearer token is valid and inject into request state""" + if not token: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail="Not authenticated", + headers={"WWW-Authenticate": "Bearer"}, + ) + + signing_key = jwkclient.get_signing_key_from_jwt(token) + decoded: dict[str, Any] = jwt.decode( + token, + signing_key.key, + algorithms=config.id_token_signing_alg_values_supported, + verify=True, + audience=config.client_audience, + issuer=config.issuer, + ) + request.state.decoded_access_token = decoded + + return validate_bearer_token + + +def access_token(request: Request) -> str | None: + """Get the decoded and verified access token of the user making the request""" + return getattr(request.state, "decoded_access_token", None) diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index 1b64630601..d5ce7c0293 100644 --- a/src/blueapi/service/main.py +++ b/src/blueapi/service/main.py @@ -19,7 +19,6 @@ from fastapi.datastructures import Address from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import RedirectResponse, StreamingResponse -from fastapi.security import OAuth2AuthorizationCodeBearer from observability_utils.tracing import ( add_span_attributes, get_tracer, @@ -37,6 +36,7 @@ from blueapi import __version__ from blueapi.config import ApplicationConfig, OIDCConfig, Tag from blueapi.service import interface +from blueapi.service.authentication import build_access_token_check from blueapi.worker import TrackableTask, WorkerState from blueapi.worker.event import TaskStatusEnum @@ -62,6 +62,7 @@ RUNNER: WorkerDispatcher | None = None LOGGER = logging.getLogger(__name__) +TRACER = get_tracer("interface") def _runner() -> WorkerDispatcher: @@ -121,7 +122,7 @@ def get_app(config: ApplicationConfig): ) dependencies = [] if config.oidc: - dependencies.append(Depends(decode_access_token(config.oidc))) + dependencies.append(Depends(build_access_token_check(config.oidc))) app.swagger_ui_init_oauth = { "clientId": "NOT_SUPPORTED", } @@ -144,32 +145,6 @@ def get_app(config: ApplicationConfig): return app -def decode_access_token(config: OIDCConfig): - jwkclient = jwt.PyJWKClient(config.jwks_uri) - oauth_scheme = OAuth2AuthorizationCodeBearer( - authorizationUrl=config.authorization_endpoint, - tokenUrl=config.token_endpoint, - refreshUrl=config.token_endpoint, - ) - - def inner(request: Request, access_token: str = Depends(oauth_scheme)): - signing_key = jwkclient.get_signing_key_from_jwt(access_token) - decoded: dict[str, Any] = jwt.decode( - access_token, - signing_key.key, - algorithms=config.id_token_signing_alg_values_supported, - verify=True, - audience=config.client_audience, - issuer=config.issuer, - ) - request.state.decoded_access_token = decoded - - return inner - - -TRACER = get_tracer("interface") - - async def on_key_error_404(_: Request, __: Exception): return JSONResponse( status_code=status.HTTP_404_NOT_FOUND, diff --git a/tests/unit_tests/service/test_authentication.py b/tests/unit_tests/service/test_authentication.py index 88227706be..2d3d0d168f 100644 --- a/tests/unit_tests/service/test_authentication.py +++ b/tests/unit_tests/service/test_authentication.py @@ -12,7 +12,7 @@ from starlette.status import HTTP_200_OK, HTTP_403_FORBIDDEN from blueapi.config import OIDCConfig, ServiceAccount -from blueapi.service import main +from blueapi.service import authentication from blueapi.service.authentication import ( SessionCacheManager, SessionManager, @@ -124,9 +124,9 @@ def test_poll_for_token_timeout( def test_server_raises_exception_for_invalid_token( oidc_config: OIDCConfig, mock_authn_server: responses.RequestsMock ): - inner = main.decode_access_token(oidc_config) + inner = authentication.build_access_token_check(oidc_config) with pytest.raises(jwt.PyJWTError): - inner(Mock(), access_token="Invalid Token") + inner(Mock(), token="Invalid Token") def test_processes_valid_token( @@ -134,8 +134,8 @@ def test_processes_valid_token( mock_authn_server: responses.RequestsMock, valid_token_with_jwt, ): - inner = main.decode_access_token(oidc_config) - inner(Mock(), access_token=valid_token_with_jwt["access_token"]) + inner = authentication.build_access_token_check(oidc_config) + inner(Mock(), token=valid_token_with_jwt["access_token"]) def test_session_cache_manager_returns_writable_file_path(tmp_path): From a9441c1c66233812f06e1aee00fdb72ce975005c Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 21 May 2026 16:45:44 +0100 Subject: [PATCH 07/11] Validate tiled service account configuration at startup --- helm/blueapi/config_schema.json | 7 +++++++ helm/blueapi/values.schema.json | 7 +++++++ src/blueapi/config.py | 1 + src/blueapi/service/authorization.py | 28 +++++++++++++++++++++++++++- src/blueapi/service/main.py | 3 ++- tests/unit_tests/test_config.py | 1 + 6 files changed, 45 insertions(+), 2 deletions(-) diff --git a/helm/blueapi/config_schema.json b/helm/blueapi/config_schema.json index 1ad4e82bfe..68fc79539a 100644 --- a/helm/blueapi/config_schema.json +++ b/helm/blueapi/config_schema.json @@ -340,8 +340,15 @@ "minLength": 1, "title": "Root", "type": "string" + }, + "tiled_service_account_check": { + "title": "Tiled Service Account Check", + "type": "string" } }, + "required": [ + "tiled_service_account_check" + ], "title": "OpaConfig", "type": "object", "$id": "OpaConfig" diff --git a/helm/blueapi/values.schema.json b/helm/blueapi/values.schema.json index 6de532cc88..7472c37a23 100644 --- a/helm/blueapi/values.schema.json +++ b/helm/blueapi/values.schema.json @@ -755,6 +755,9 @@ "$id": "OpaConfig", "title": "OpaConfig", "type": "object", + "required": [ + "tiled_service_account_check" + ], "properties": { "root": { "title": "Root", @@ -763,6 +766,10 @@ "format": "uri", "maxLength": 2083, "minLength": 1 + }, + "tiled_service_account_check": { + "title": "Tiled Service Account Check", + "type": "string" } }, "additionalProperties": false diff --git a/src/blueapi/config.py b/src/blueapi/config.py index 4c2431cf1e..06c1499559 100644 --- a/src/blueapi/config.py +++ b/src/blueapi/config.py @@ -298,6 +298,7 @@ class Tag(StrEnum): class OpaConfig(BlueapiBaseModel): root: HttpUrl = HttpUrl("http://localhost:8181") + tiled_service_account_check: str class ApplicationConfig(BlueapiBaseModel): diff --git a/src/blueapi/service/authorization.py b/src/blueapi/service/authorization.py index b27c304c97..3be0814d43 100644 --- a/src/blueapi/service/authorization.py +++ b/src/blueapi/service/authorization.py @@ -6,7 +6,8 @@ import aiohttp from aiohttp import ClientSession -from blueapi.config import OpaConfig +from blueapi.config import OIDCConfig, OpaConfig, ServiceAccount +from blueapi.service.authentication import TiledAuth LOGGER = logging.getLogger(__name__) @@ -52,6 +53,14 @@ def for_config( LOGGER.info("No OPA config provided - not creating OpaClient") return nullcontext() + async def require_tiled_service_account(self, token: str): + if not await self._call_opa( + self._conf.tiled_service_account_check, + {"token": token, "beamline": self._instrument}, + ): + raise ValueError( + f"Tiled service account is not valid for '{self._instrument}'" + ) class OpaUserClient: @@ -61,3 +70,20 @@ class OpaUserClient: def __init__(self, client: OpaClient, token: str): self.client = client self.token = token + + +async def validate_tiled_config( + tiled: ServiceAccount | str | None, oidc: OIDCConfig | None, opa: OpaClient | None +): + if not isinstance(tiled, ServiceAccount): + # can't validate an API key + return + + if not opa or not oidc: + LOGGER.info("Missing OPA or OIDC configuration required to validate tiled auth") + return + + LOGGER.info("Validating tiled configuration") + tiled.token_url = oidc.token_endpoint + auth = TiledAuth(tiled) + await opa.require_tiled_service_account(auth.get_access_token()) diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index d5ce7c0293..20f5fb0ec8 100644 --- a/src/blueapi/service/main.py +++ b/src/blueapi/service/main.py @@ -40,7 +40,7 @@ from blueapi.worker import TrackableTask, WorkerState from blueapi.worker.event import TaskStatusEnum -from .authorization import OpaClient +from .authorization import OpaClient, validate_tiled_config from .model import ( DeviceModel, DeviceResponse, @@ -98,6 +98,7 @@ async def inner(app: FastAPI): setup_runner(config) async with OpaClient.for_config(meta and meta.instrument, config.opa) as opa: app.state.authz = opa + await validate_tiled_config(config.tiled.authentication, config.oidc, opa) yield teardown_runner() diff --git a/tests/unit_tests/test_config.py b/tests/unit_tests/test_config.py index 7ce98f6fe6..b6e52c3938 100644 --- a/tests/unit_tests/test_config.py +++ b/tests/unit_tests/test_config.py @@ -339,6 +339,7 @@ def test_config_yaml_parsed(temp_yaml_config_file): }, "opa": { "root": "http://opa.example.com/", + "tiled_service_account_check": "v1/tiled_service_account", }, }, { From ddf28f24a5c03701f757c4c945d702eb99d06506 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 28 May 2026 17:22:29 +0100 Subject: [PATCH 08/11] Add tests for tiled check --- .../unit_tests/service/test_authorization.py | 93 ++++++++++++++++++- 1 file changed, 91 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/service/test_authorization.py b/tests/unit_tests/service/test_authorization.py index aa4417d500..ca3500bbaa 100644 --- a/tests/unit_tests/service/test_authorization.py +++ b/tests/unit_tests/service/test_authorization.py @@ -1,11 +1,13 @@ -from unittest.mock import MagicMock, patch +from contextlib import AbstractContextManager, nullcontext +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest from pydantic import HttpUrl -from blueapi.config import OpaConfig +from blueapi.config import OIDCConfig, OpaConfig, ServiceAccount from blueapi.service.authorization import ( OpaClient, + validate_tiled_config, ) # Reusable client patch decorator @@ -20,9 +22,50 @@ def opa_config() -> OpaConfig: return OpaConfig( root=HttpUrl("http://auth.example.com"), + tiled_service_account_check="/auth/tiled", ) +@patch_client_session +@pytest.mark.parametrize( + "result,context", + [ + (False, pytest.raises(ValueError, match="Tiled service account is not valid ")), + (True, nullcontext()), + ], +) +async def test_tiled_service_account( + session: MagicMock, + opa_config: OpaConfig, + result: bool, + context: AbstractContextManager, +): + session.return_value.post = AsyncMock( + return_value=MagicMock(json=AsyncMock(return_value={"result": result})) + ) + + client = OpaClient(instrument="p99", config=opa_config) + + session.assert_called_once_with(base_url="http://auth.example.com/") + with context: + await client.require_tiled_service_account(token="foo_bar") + session().post.assert_called_once_with( + "/auth/tiled", + json={"input": {"token": "foo_bar", "beamline": "p99", "audience": "account"}}, + ) + + +@patch_client_session +async def test_exception_raised_when_opa_fails( + session: MagicMock, opa_config: OpaConfig +): + session.return_value.post = AsyncMock(side_effect=RuntimeError("Connection failed")) + async with OpaClient.for_config("p45", opa_config) as client: + assert client is not None + with pytest.raises(RuntimeError, match="Connection failed"): + await client.require_tiled_service_account(token="foo_bar") + + @patch_client_session async def test_session_closed(session: MagicMock, opa_config: OpaConfig): async with OpaClient.for_config("p45", opa_config): @@ -46,3 +89,49 @@ async def test_opa_client_without_config(instrument: str | None): async def test_opa_fails_without_instrument(opa_config: OpaConfig): with pytest.raises(ValueError, match="Instrument name is required"): OpaClient.for_config(None, opa_config) + + +async def test_validate_tiled_config(): + opa = MagicMock(spec=OpaClient) + tiled = ServiceAccount() + oidc = Mock(spec=OIDCConfig) + oidc.token_endpoint = "token-endpoint" + with patch("blueapi.service.authorization.TiledAuth") as auth: + auth.return_value.get_access_token.return_value = "tiled-token" + await validate_tiled_config(tiled, oidc, opa) + + auth.assert_called_once_with(tiled) + opa.require_tiled_service_account.assert_called_once_with("tiled-token") + + +@pytest.mark.parametrize( + "tiled_auth,oidc,opa_client", + [ + (None, None, MagicMock(spec=OpaClient)), + ( + None, + OIDCConfig(well_known_url="http://example.com", client_id="test-client"), + MagicMock(spec=OpaClient), + ), + ("api_key", None, MagicMock(spec=OpaClient)), + ( + "api_key", + OIDCConfig(well_known_url="http://example.com", client_id="test-client"), + MagicMock(spec=OpaClient), + ), + (ServiceAccount(), None, MagicMock(spec=OpaClient)), + ( + ServiceAccount(), + OIDCConfig(well_known_url="http://example.com", client_id="test-client"), + None, + ), + ], +) +async def test_validate_tiled_config_with_missing_config( + tiled_auth: ServiceAccount | str | None, + oidc: OIDCConfig | None, + opa_client: MagicMock | None, +): + assert await validate_tiled_config(tiled_auth, oidc, opa_client) is None + if opa_client is not None: + opa_client.require_tiled_service_account.assert_not_called() From 8ebdcb6d322ebdb88a2214b9fa5916a9f296790c Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Tue, 26 May 2026 12:01:54 +0100 Subject: [PATCH 09/11] Add opa dependency function to create OpaUserClient --- src/blueapi/service/authorization.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/src/blueapi/service/authorization.py b/src/blueapi/service/authorization.py index 3be0814d43..4005188ec1 100644 --- a/src/blueapi/service/authorization.py +++ b/src/blueapi/service/authorization.py @@ -1,13 +1,15 @@ import logging from collections.abc import Mapping from contextlib import AbstractAsyncContextManager, aclosing, nullcontext -from typing import Any, Self +from typing import Any, Self, cast import aiohttp from aiohttp import ClientSession +from fastapi import Depends, HTTPException, Request +from starlette import status from blueapi.config import OIDCConfig, OpaConfig, ServiceAccount -from blueapi.service.authentication import TiledAuth +from blueapi.service.authentication import TiledAuth, unchecked_bearer_token LOGGER = logging.getLogger(__name__) @@ -87,3 +89,14 @@ async def validate_tiled_config( tiled.token_url = oidc.token_endpoint auth = TiledAuth(tiled) await opa.require_tiled_service_account(auth.get_access_token()) + + +async def opa( + request: Request, token: str | None = Depends(unchecked_bearer_token) +) -> OpaUserClient | None: + + if opa := cast(OpaClient | None, getattr(request.app.state, "authz", None)): + if not token: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) + return opa.for_token(token) + return None From 3f430d4dc1e85708a2c6e0c47553d9ec489e20cc Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 28 May 2026 17:24:54 +0100 Subject: [PATCH 10/11] test opa dependency function --- src/blueapi/service/authorization.py | 2 +- .../unit_tests/service/test_authorization.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/blueapi/service/authorization.py b/src/blueapi/service/authorization.py index 4005188ec1..5e923a04c7 100644 --- a/src/blueapi/service/authorization.py +++ b/src/blueapi/service/authorization.py @@ -98,5 +98,5 @@ async def opa( if opa := cast(OpaClient | None, getattr(request.app.state, "authz", None)): if not token: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) - return opa.for_token(token) + return OpaUserClient(opa, token) return None diff --git a/tests/unit_tests/service/test_authorization.py b/tests/unit_tests/service/test_authorization.py index ca3500bbaa..aaa02a7a99 100644 --- a/tests/unit_tests/service/test_authorization.py +++ b/tests/unit_tests/service/test_authorization.py @@ -2,11 +2,13 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest +from fastapi import HTTPException from pydantic import HttpUrl from blueapi.config import OIDCConfig, OpaConfig, ServiceAccount from blueapi.service.authorization import ( OpaClient, + opa, validate_tiled_config, ) @@ -135,3 +137,28 @@ async def test_validate_tiled_config_with_missing_config( assert await validate_tiled_config(tiled_auth, oidc, opa_client) is None if opa_client is not None: opa_client.require_tiled_service_account.assert_not_called() + + +async def test_opa_dependency_method(): + request = MagicMock() + + user_client = await opa(request, "foo_bar") + + assert user_client is not None + assert user_client.client == request.app.state.authz + assert user_client.token == "foo_bar" + + +async def test_opa_dependency_without_token(): + request = MagicMock() + + with pytest.raises(HTTPException, match="401"): + await opa(request, None) + + +@pytest.mark.parametrize("token", ["foo_bar", None]) +async def test_opa_dependency_without_authz(token): + request = MagicMock() + del request.app.state.authz + user_client = await opa(request, token) + assert user_client is None From 3617c09d516dc53781f04f349f9497c8fdf8ac11 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Tue, 26 May 2026 12:17:35 +0100 Subject: [PATCH 11/11] Add can_submit_task auth check method and config --- helm/blueapi/config_schema.json | 7 ++++++- helm/blueapi/values.schema.json | 7 ++++++- src/blueapi/config.py | 1 + src/blueapi/service/authorization.py | 21 +++++++++++++++++++ .../unit_tests/service/test_authorization.py | 1 + tests/unit_tests/test_config.py | 1 + 6 files changed, 36 insertions(+), 2 deletions(-) diff --git a/helm/blueapi/config_schema.json b/helm/blueapi/config_schema.json index 68fc79539a..d5203cd562 100644 --- a/helm/blueapi/config_schema.json +++ b/helm/blueapi/config_schema.json @@ -344,10 +344,15 @@ "tiled_service_account_check": { "title": "Tiled Service Account Check", "type": "string" + }, + "submit_task_check": { + "title": "Submit Task Check", + "type": "string" } }, "required": [ - "tiled_service_account_check" + "tiled_service_account_check", + "submit_task_check" ], "title": "OpaConfig", "type": "object", diff --git a/helm/blueapi/values.schema.json b/helm/blueapi/values.schema.json index 7472c37a23..4fa6f4f280 100644 --- a/helm/blueapi/values.schema.json +++ b/helm/blueapi/values.schema.json @@ -756,7 +756,8 @@ "title": "OpaConfig", "type": "object", "required": [ - "tiled_service_account_check" + "tiled_service_account_check", + "submit_task_check" ], "properties": { "root": { @@ -767,6 +768,10 @@ "maxLength": 2083, "minLength": 1 }, + "submit_task_check": { + "title": "Submit Task Check", + "type": "string" + }, "tiled_service_account_check": { "title": "Tiled Service Account Check", "type": "string" diff --git a/src/blueapi/config.py b/src/blueapi/config.py index 06c1499559..49cfa113b5 100644 --- a/src/blueapi/config.py +++ b/src/blueapi/config.py @@ -299,6 +299,7 @@ class Tag(StrEnum): class OpaConfig(BlueapiBaseModel): root: HttpUrl = HttpUrl("http://localhost:8181") tiled_service_account_check: str + submit_task_check: str class ApplicationConfig(BlueapiBaseModel): diff --git a/src/blueapi/service/authorization.py b/src/blueapi/service/authorization.py index 5e923a04c7..ff02243819 100644 --- a/src/blueapi/service/authorization.py +++ b/src/blueapi/service/authorization.py @@ -1,4 +1,5 @@ import logging +import re from collections.abc import Mapping from contextlib import AbstractAsyncContextManager, aclosing, nullcontext from typing import Any, Self, cast @@ -10,8 +11,10 @@ from blueapi.config import OIDCConfig, OpaConfig, ServiceAccount from blueapi.service.authentication import TiledAuth, unchecked_bearer_token +from blueapi.service.model import TaskRequest LOGGER = logging.getLogger(__name__) +INSTRUMENT_SESSION_RE = re.compile(r"^[a-z]{2}(?P\d+)-(?P\d+)$") class OpaClient: @@ -64,6 +67,20 @@ async def require_tiled_service_account(self, token: str): f"Tiled service account is not valid for '{self._instrument}'" ) + async def require_submit_task(self, instrument_session: str, token: str): + if not (match := INSTRUMENT_SESSION_RE.match(instrument_session)): + raise ValueError("Invalid instrument session") + + if not await self._call_opa( + self._conf.submit_task_check, + { + "token": token, + "proposal": int(match["proposal"]), + "visit": int(match["visit"]), + }, + ): + raise HTTPException(status_code=status.HTTP_403_UNORTHORIZED) + class OpaUserClient: client: OpaClient @@ -73,6 +90,10 @@ def __init__(self, client: OpaClient, token: str): self.client = client self.token = token + async def can_submit_task(self, task: TaskRequest): + LOGGER.info("Checking permissions to run task") + await self.client.require_submit_task(task.instrument_session, self.token) + async def validate_tiled_config( tiled: ServiceAccount | str | None, oidc: OIDCConfig | None, opa: OpaClient | None diff --git a/tests/unit_tests/service/test_authorization.py b/tests/unit_tests/service/test_authorization.py index aaa02a7a99..f4ae123c81 100644 --- a/tests/unit_tests/service/test_authorization.py +++ b/tests/unit_tests/service/test_authorization.py @@ -24,6 +24,7 @@ def opa_config() -> OpaConfig: return OpaConfig( root=HttpUrl("http://auth.example.com"), + submit_task_check="/auth/submit", tiled_service_account_check="/auth/tiled", ) diff --git a/tests/unit_tests/test_config.py b/tests/unit_tests/test_config.py index b6e52c3938..0d161b0614 100644 --- a/tests/unit_tests/test_config.py +++ b/tests/unit_tests/test_config.py @@ -340,6 +340,7 @@ def test_config_yaml_parsed(temp_yaml_config_file): "opa": { "root": "http://opa.example.com/", "tiled_service_account_check": "v1/tiled_service_account", + "submit_task_check": "v1/submit_task", }, }, {