feat: Token exchange for tiled insertion

.github/workflows/_system_test.yml

@@ -31,6 +31,9 @@ jobs:
           EPICS_CA_NAME_SERVERS: 127.0.0.1:9064
           EPICS_PVA_NAME_SERVERS: 127.0.0.1:9075
         run: uv run blueapi -c ${{ github.workspace }}/tests/system_tests/config.yaml serve &
+
+      - name: Install playwright browsers
+        run: uv run playwright install --with-deps 
 
       - name: Run tests
         run: uv run --locked tox -e system-test

.gitleaksignore

@@ -0,0 +1,2 @@
+tests/system_tests/services/blueapi-oauth2-proxy/oauth2-proxy.cfg:generic-api-key:16
+tests/system_tests/services/tiled-oauth2-proxy/oauth2-proxy.cfg:generic-api-key:14

helm/blueapi/config_schema.json

@@ -486,17 +486,22 @@
                     "title": "Url",
                     "type": "string"
                 },
-                "api_key": {
-                    "anyOf": [
-                        {
-                            "type": "string"
-                        },
-                        {
-                            "type": "null"
-                        }
-                    ],
-                    "default": null,
-                    "title": "Api Key"
+                "token_exchange_secret": {
+                    "default": "",
+                    "description": "Token exchange client secret",
+                    "title": "Token Exchange Secret",
+                    "type": "string"
+                },
+                "token_url": {
+                    "default": "",
+                    "title": "Token Url",
+                    "type": "string"
+                },
+                "token_exchange_client_id": {
+                    "default": "",
+                    "description": "Token exchange Client ID",
+                    "title": "Token Exchange Client Id",
+                    "type": "string"
                 }
             },
             "title": "TiledConfig",

helm/blueapi/values.schema.json

@@ -893,23 +893,29 @@
             "title": "TiledConfig",
             "type": "object",
             "properties": {
-                "api_key": {
-                    "title": "Api Key",
-                    "anyOf": [
-                        {
-                            "type": "string"
-                        },
-                        {
-                            "type": "null"
-                        }
-                    ]
-                },
                 "enabled": {
                     "title": "Enabled",
                     "description": "True if blueapi should forward data to a Tiled instance",
                     "default": false,
                     "type": "boolean"
                 },
+                "token_exchange_client_id": {
+                    "title": "Token Exchange Client Id",
+                    "description": "Token exchange Client ID",
+                    "default": "",
+                    "type": "string"
+                },
+                "token_exchange_secret": {
+                    "title": "Token Exchange Secret",
+                    "description": "Token exchange client secret",
+                    "default": "",
+                    "type": "string"
+                },
+                "token_url": {
+                    "title": "Token Url",
+                    "default": "",
+                    "type": "string"
+                },
                 "url": {
                     "title": "Url",
                     "default": "http://localhost:8407/",

pyproject.toml

@@ -13,7 +13,7 @@ classifiers = [
 ]
 description = "Lightweight bluesky-as-a-service wrapper application. Also usable as a library."
 dependencies = [
-    "tiled[client]>=0.2.3",
+    "tiled[client] @ git+https://github.com/bluesky/tiled.git@add-auth-for-client",
     "bluesky[plotting]>=1.14.0",                     # plotting includes matplotlib, required for BestEffortCallback in run plans
     "ophyd-async>=0.13.5",
     "aioca",
@@ -69,7 +69,10 @@ dev = [
     "mock",
     "jwcrypto",
     "deepdiff",
-    "tiled[minimal-server]>=0.2.3", # For system-test of dls.py
+    "tiled[minimal-server] @ git+https://github.com/bluesky/tiled.git@add-auth-for-client", # For system-test of dls.py
+    "playwright",
+    "pytest-playwright",
+    "respx",
 ]
 
 [project.scripts]

src/blueapi/config.py

@@ -112,7 +112,13 @@ class TiledConfig(BlueapiBaseModel):
         default=False,
     )
     url: HttpUrl = HttpUrl("http://localhost:8407")
-    api_key: str | None = os.environ.get("TILED_SINGLE_USER_API_KEY", None)
+    token_exchange_secret: str = Field(
+        description="Token exchange client secret", default=""
+    )
+    token_url: str = Field(default="")
+    token_exchange_client_id: str = Field(
+        description="Token exchange Client ID", default=""
+    )
 
 
 class WorkerEventConfig(BlueapiBaseModel):

src/blueapi/core/context.py

@@ -181,6 +181,17 @@ def _update_scan_num(md: dict[str, Any]) -> int:
                     "Tiled has been configured but `instrument` metadata is not set - "
                     "this field is required to make authorization decisions."
                 )
+            if configuration.oidc is None:
+                raise InvalidConfigError(
+                    "Tiled has been configured but oidc configuration is missing "
+                    "this field is required to make authorization decisions."
+                )
+            if tiled_conf.token_exchange_secret == "":
+                raise InvalidConfigError(
+                    "Tiled has been enabled but Token exchange secret has not been set "
+                    "this field is required to enable tiled insertion."
+                )
+            tiled_conf.token_url = configuration.oidc.token_endpoint
             self.tiled_conf = tiled_conf
 
     def find_device(self, addr: str | list[str]) -> Device | None:

src/blueapi/service/authentication.py

@@ -2,24 +2,27 @@
 
 import base64
 import os
+import threading
 import time
 import webbrowser
 from abc import ABC, abstractmethod
+from enum import Enum
 from functools import cached_property
 from http import HTTPStatus
 from pathlib import Path
 from typing import Any, cast
 
+import httpx
 import jwt
 import requests
-from pydantic import TypeAdapter
+from pydantic import BaseModel, TypeAdapter, computed_field
 from requests.auth import AuthBase
 
-from blueapi.config import OIDCConfig
+from blueapi.config import OIDCConfig, TiledConfig
 from blueapi.service.model import Cache
 
 DEFAULT_CACHE_DIR = "~/.cache/"
-SCOPES = "openid offline_access"
+SCOPES = "openid"
 
 
 class CacheManager(ABC):
@@ -239,3 +242,104 @@ def __call__(self, request):
         if self.token:
             request.headers["Authorization"] = f"Bearer {self.token}"
         return request
+
+
+class TokenType(str, Enum):
+    refresh_token = "refresh_token"
+    access_token = "access_token"
+
+
+class Token(BaseModel):
+    token: str
+    expires_at: float | None
+
+    @computed_field
+    @property
+    def expired(self) -> bool:
+        if self.expires_at is None:
+            # Assume token is valid
+            return False
+        return time.time() > self.expires_at
+
+    def _get_token_expires_at(
+        self, token_dict: dict[str, Any], token_type: TokenType
+    ) -> int | None:
+        expires_at = None
+        if token_type == TokenType.access_token:
+            if "expires_in" in token_dict:
+                expires_at = int(time.time()) + int(token_dict["expires_in"])
+        elif token_type == TokenType.refresh_token:
+            if "refresh_expires_in" in token_dict:
+                expires_at = int(time.time()) + int(token_dict["refresh_expires_in"])
+        return expires_at
+
+    def __init__(self, token_dict: dict[str, Any], token_type: TokenType):
+        token = token_dict.get(token_type)
+        if token is None:
+            raise ValueError(f"Not able to find {token_type} in response")
+        super().__init__(
+            token=token, expires_at=self._get_token_expires_at(token_dict, token_type)
+        )
+
+    def __str__(self) -> str:
+        return str(self.token)
+
+
+class TiledAuth(httpx.Auth):
+    def __init__(self, tiled_config: TiledConfig, blueapi_jwt_token: str):
+        self._tiled_config = tiled_config
+        self._blueapi_jwt_token = blueapi_jwt_token
+        self._sync_lock = threading.RLock()
+        self._access_token: Token | None = None
+        self._refresh_token: Token | None = None
+
+    def exchange_access_token(self):
+        request_data = {
+            "client_id": self._tiled_config.token_exchange_client_id,
+            "client_secret": self._tiled_config.token_exchange_secret,
+            "subject_token": self._blueapi_jwt_token,
+            "subject_token_type": "urn:ietf:params:oauth:token-type:access_token",
+            "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+            "requested_token_type": "urn:ietf:params:oauth:token-type:refresh_token",
+        }
+        with self._sync_lock:
+            response = httpx.post(
+                self._tiled_config.token_url,
+                data=request_data,
+            )
+            response.raise_for_status()
+            self.sync_tokens(response.json())
+
+    def refresh_token(self):
+        if self._refresh_token is None:
+            raise Exception("Cannot refresh session as no refresh token available")
+        with self._sync_lock:
+            response = httpx.post(
+                self._tiled_config.token_url,
+                data={
+                    "client_id": self._tiled_config.token_exchange_client_id,
+                    "client_secret": self._tiled_config.token_exchange_secret,
+                    "grant_type": "refresh_token",
+                    "refresh_token": self._refresh_token,
+                },
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+            response.raise_for_status()
+            self.sync_tokens(response.json())
+
+    def sync_tokens(self, response):
+        self._access_token = Token(response, TokenType.access_token)
+        self._refresh_token = Token(response, TokenType.refresh_token)
+
+    def sync_auth_flow(self, request):
+        if self._access_token is not None and self._access_token.expired is not True:
+            request.headers["Authorization"] = f"Bearer {self._access_token}"
+            yield request
+        elif self._access_token is None:
+            self.exchange_access_token()
+            request.headers["Authorization"] = f"Bearer {self._access_token}"
+            yield request
+        else:
+            self.refresh_token()
+            request.headers["Authorization"] = f"Bearer {self._access_token}"
+            yield request

src/blueapi/service/constants.py

@@ -0,0 +1,4 @@
+CONTEXT_HEADER = "traceparent"
+VENDOR_CONTEXT_HEADER = "tracestate"
+AUTHORIZAITON_HEADER = "authorization"
+PROPAGATED_HEADERS = {CONTEXT_HEADER, VENDOR_CONTEXT_HEADER, AUTHORIZAITON_HEADER}

src/blueapi/service/interface.py

@@ -12,6 +12,8 @@
 from blueapi.core.context import BlueskyContext
 from blueapi.core.event import EventStream
 from blueapi.log import set_up_logging
+from blueapi.service.authentication import TiledAuth
+from blueapi.service.constants import AUTHORIZAITON_HEADER
 from blueapi.service.model import (
     DeviceModel,
     PlanModel,
@@ -184,10 +186,27 @@ def begin_task(
 
     if tiled_config := active_context.tiled_conf:
         # Tiled queries the root node, so must create an authorized client
+        blueapi_jwt_token = ""
+        if pass_through_headers is None:
+            raise ValueError(
+                "Tiled config is enabled but no "
+                f"{AUTHORIZAITON_HEADER} header in request"
+            )
+        authorization_header_value = pass_through_headers.get(AUTHORIZAITON_HEADER)
+        from fastapi.security.utils import get_authorization_scheme_param
+
+        _, blueapi_jwt_token = get_authorization_scheme_param(
+            authorization_header_value
+        )
+
+        if blueapi_jwt_token == "":
+            raise KeyError("Tiled config is enabled but no Bearer Token in request")
         tiled_client = from_uri(
             str(tiled_config.url),
-            api_key=tiled_config.api_key,
-            headers=pass_through_headers,
+            auth=TiledAuth(
+                tiled_config,
+                blueapi_jwt_token=blueapi_jwt_token,
+            ),
         )
         tiled_writer_token = active_context.run_engine.subscribe(
             TiledWriter(tiled_client, batch_size=1)

src/blueapi/service/main.py

@@ -36,6 +36,11 @@
 
 from blueapi.config import ApplicationConfig, OIDCConfig
 from blueapi.service import interface
+from blueapi.service.constants import (
+    CONTEXT_HEADER,
+    PROPAGATED_HEADERS,
+    VENDOR_CONTEXT_HEADER,
+)
 from blueapi.worker import TrackableTask, WorkerState
 from blueapi.worker.event import TaskStatusEnum
 
@@ -67,10 +72,7 @@
 RUNNER: WorkerDispatcher | None = None
 
 LOGGER = logging.getLogger(__name__)
-CONTEXT_HEADER = "traceparent"
-VENDOR_CONTEXT_HEADER = "tracestate"
-AUTHORIZAITON_HEADER = "authorization"
-PROPAGATED_HEADERS = {CONTEXT_HEADER, VENDOR_CONTEXT_HEADER, AUTHORIZAITON_HEADER}
+
 DOCS_ENDPOINT = "/docs"