feat: Add OpaClient to wrap auth checks (#1541)

Proof of concept opa client with dependency injection and example check

Author: tpoliaw

helm/blueapi/config_schema.json

@@ -330,6 +330,27 @@
             "type": "object",
             "$id": "OIDCConfig"
         },
+        "OpaConfig": {
+            "additionalProperties": false,
+            "properties": {
+                "root": {
+                    "default": "http://localhost:8181/",
+                    "format": "uri",
+                    "maxLength": 2083,
+                    "minLength": 1,
+                    "title": "Root",
+                    "type": "string"
+                },
+                "audience": {
+                    "default": "account",
+                    "title": "Audience",
+                    "type": "string"
+                }
+            },
+            "title": "OpaConfig",
+            "type": "object",
+            "$id": "OpaConfig"
+        },
         "PlanSource": {
             "additionalProperties": false,
             "properties": {
@@ -612,6 +633,17 @@
                 }
             ],
             "default": null
+        },
+        "opa": {
+            "anyOf": [
+                {
+                    "$ref": "OpaConfig"
+                },
+                {
+                    "type": "null"
+                }
+            ],
+            "default": null
         }
     },
     "title": "ApplicationConfig",

helm/blueapi/values.schema.json

@@ -751,6 +751,27 @@
             },
             "additionalProperties": false
         },
+        "OpaConfig": {
+            "$id": "OpaConfig",
+            "title": "OpaConfig",
+            "type": "object",
+            "properties": {
+                "audience": {
+                    "title": "Audience",
+                    "default": "account",
+                    "type": "string"
+                },
+                "root": {
+                    "title": "Root",
+                    "default": "http://localhost:8181/",
+                    "type": "string",
+                    "format": "uri",
+                    "maxLength": 2083,
+                    "minLength": 1
+                }
+            },
+            "additionalProperties": false
+        },
         "PlanSource": {
             "$id": "PlanSource",
             "title": "PlanSource",
@@ -1011,6 +1032,16 @@
                         }
                     ]
                 },
+                "opa": {
+                    "anyOf": [
+                        {
+                            "$ref": "OpaConfig"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ]
+                },
                 "scratch": {
                     "anyOf": [
                         {

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"

src/blueapi/config.py

@@ -296,6 +296,11 @@ class Tag(StrEnum):
     META = "Meta"
 
 
+class OpaConfig(BlueapiBaseModel):
+    root: HttpUrl = HttpUrl("http://localhost:8181")
+    audience: str = "account"
+
+
 class ApplicationConfig(BlueapiBaseModel):
     """
     Config for the worker application as a whole. Root of
@@ -335,6 +340,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 +349,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
 

src/blueapi/service/authorization.py

@@ -0,0 +1,47 @@
+import logging
+from collections.abc import Mapping
+from contextlib import AbstractAsyncContextManager, aclosing, nullcontext
+from typing import Any, Self
+
+from aiohttp import ClientSession
+
+from blueapi.config import OpaConfig
+
+LOGGER = logging.getLogger(__name__)
+
+
+class OpaClient:
+    def __init__(self, instrument: str, config: OpaConfig):
+        LOGGER.info("Creating OpaClient for %s with config %s", instrument, config)
+        self._instrument = instrument
+        self._config = config
+        self._session = ClientSession(base_url=config.root.encoded_string())
+        self._audience = config.audience
+
+    async def aclose(self):
+        LOGGER.info("Closing OPA session")
+        await self._session.close()
+
+    async def _call_opa(self, endpoint: str, data: Mapping[str, Any]) -> bool:
+        resp = await self._session.post(
+            endpoint,
+            json={
+                "input": {
+                    "beamline": self._instrument,
+                    "audience": self._audience,
+                    **data,
+                }
+            },
+        )
+        return (await resp.json())["result"]
+
+    @classmethod
+    def for_config(
+        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()

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,
@@ -93,8 +94,11 @@ def teardown_runner():
 def lifespan(config: ApplicationConfig):
     @asynccontextmanager
     async def inner(app: FastAPI):
+        meta = config.env.metadata
         setup_runner(config)
-        yield
+        async with OpaClient.for_config(meta and meta.instrument, config.opa) as opa:
+            app.state.authz = opa
+            yield
         teardown_runner()
 
     return inner

tests/unit_tests/service/test_authorization.py

@@ -0,0 +1,62 @@
+from unittest.mock import AsyncMock, 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)
+
+
+@patch_client_session
+async def test_opa_adds_input_fields(session: MagicMock, opa_config: OpaConfig):
+    session.return_value.post = AsyncMock()
+    async with OpaClient.for_config("p45", opa_config) as opa:
+        assert opa is not None
+        await opa._call_opa("foo/bar", data={"foo": "bar"})
+
+    session.assert_called_once()
+    session().post.assert_called_once_with(
+        "foo/bar",
+        json={"input": {"beamline": "p45", "audience": "account", "foo": "bar"}},
+    )

tests/unit_tests/service/test_main.py

@@ -1,5 +1,5 @@
 from unittest import mock
-from unittest.mock import Mock, call
+from unittest.mock import Mock, call, patch
 
 import pytest
 from fastapi import FastAPI, Request
@@ -10,6 +10,7 @@
 from blueapi.service.main import (
     add_version_headers,
     get_passthrough_headers,
+    lifespan,
     log_request_details,
 )
 
@@ -79,3 +80,18 @@ def test_get_passthrough_headers(
     request = Mock(spec=Request)
     request.headers = headers
     assert get_passthrough_headers(request) == expected_headers
+
+
+@patch("blueapi.service.main.teardown_runner")
+@patch("blueapi.service.main.setup_runner")
+async def test_lifespan(setup: Mock, teardown: Mock):
+    conf = ApplicationConfig()
+    lifespan_fn = lifespan(conf)
+
+    app = Mock()
+
+    async with lifespan_fn(app):
+        setup.assert_called_once_with(conf)
+        teardown.assert_not_called()
+
+    teardown.assert_called_once()

tests/unit_tests/test_config.py

@@ -337,6 +337,10 @@ def test_config_yaml_parsed(temp_yaml_config_file):
                     }
                 ],
             },
+            "opa": {
+                "root": "http://opa.example.com/",
+                "audience": "account",
+            },
         },
         {
             "stomp": {
@@ -392,6 +396,7 @@ def test_config_yaml_parsed(temp_yaml_config_file):
                     }
                 ],
             },
+            "opa": None,
         },
     ],
     indirect=True,