feat: Validate tiled ServiceAccount config at startup (#1548)

Ensure that the tiled auth configuration is valid and fail fast instead
of when the first data is written.

Author: tpoliaw

helm/blueapi/config_schema.json

@@ -345,8 +345,15 @@
                     "default": "account",
                     "title": "Audience",
                     "type": "string"
+                },
+                "tiled_service_account_check": {
+                    "title": "Tiled Service Account Check",
+                    "type": "string"
                 }
             },
+            "required": [
+                "tiled_service_account_check"
+            ],
             "title": "OpaConfig",
             "type": "object",
             "$id": "OpaConfig"

helm/blueapi/values.schema.json

@@ -755,6 +755,9 @@
             "$id": "OpaConfig",
             "title": "OpaConfig",
             "type": "object",
+            "required": [
+                "tiled_service_account_check"
+            ],
             "properties": {
                 "audience": {
                     "title": "Audience",
@@ -768,6 +771,10 @@
                     "format": "uri",
                     "maxLength": 2083,
                     "minLength": 1
+                },
+                "tiled_service_account_check": {
+                    "title": "Tiled Service Account Check",
+                    "type": "string"
                 }
             },
             "additionalProperties": false

src/blueapi/config.py

@@ -299,6 +299,7 @@ class Tag(StrEnum):
 class OpaConfig(BlueapiBaseModel):
     root: HttpUrl = HttpUrl("http://localhost:8181")
     audience: str = "account"
+    tiled_service_account_check: str
 
 
 class ApplicationConfig(BlueapiBaseModel):

src/blueapi/service/authorization.py

@@ -5,7 +5,8 @@
 
 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__)
 
@@ -45,3 +46,29 @@ def for_config(
             return aclosing(cls(instrument, 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._config.tiled_service_account_check,
+            {"token": token, "beamline": self._instrument},
+        ):
+            raise ValueError(
+                f"Tiled service account is not valid for '{self._instrument}'"
+            )
+
+
+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())

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()
 

tests/unit_tests/service/test_authorization.py

@@ -1,11 +1,13 @@
-from unittest.mock import AsyncMock, 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):
@@ -60,3 +103,49 @@ async def test_opa_adds_input_fields(session: MagicMock, opa_config: OpaConfig):
         "foo/bar",
         json={"input": {"beamline": "p45", "audience": "account", "foo": "bar"}},
     )
+
+
+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()

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/",
                 "audience": "account",
+                "tiled_service_account_check": "v1/tiled_service_account",
             },
         },
         {