feat: Add submit task permission checks

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",

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"

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

src/blueapi/service/authorization.py

@@ -1,15 +1,20 @@
 import logging
+import re
 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
+from blueapi.service.model import TaskRequest
 
 LOGGER = logging.getLogger(__name__)
+INSTRUMENT_SESSION_RE = re.compile(r"^[a-z]{2}(?P<proposal>\d+)-(?P<visit>\d+)$")
 
 
 class OpaClient:
@@ -62,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
@@ -71,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
@@ -87,3 +110,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 OpaUserClient(opa, token)
+    return None

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,
 )
 
@@ -22,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",
     )
 
@@ -149,3 +152,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

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",
             },
         },
         {