feat: add authz dependency injections

helm/blueapi/config_schema.json

@@ -344,10 +344,20 @@
                 "tiled_service_account_check": {
                     "title": "Tiled Service Account Check",
                     "type": "string"
+                },
+                "submit_task_check": {
+                    "title": "Submit Task Check",
+                    "type": "string"
+                },
+                "admin_check": {
+                    "title": "Admin Check",
+                    "type": "string"
                 }
             },
             "required": [
-                "tiled_service_account_check"
+                "tiled_service_account_check",
+                "submit_task_check",
+                "admin_check"
             ],
             "title": "OpaConfig",
             "type": "object",

helm/blueapi/values.schema.json

@@ -756,9 +756,15 @@
             "title": "OpaConfig",
             "type": "object",
             "required": [
-                "tiled_service_account_check"
+                "tiled_service_account_check",
+                "submit_task_check",
+                "admin_check"
             ],
             "properties": {
+                "admin_check": {
+                    "title": "Admin Check",
+                    "type": "string"
+                },
                 "root": {
                     "title": "Root",
                     "default": "http://localhost:8181/",
@@ -767,6 +773,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,8 @@ class Tag(StrEnum):
 class OpaConfig(BlueapiBaseModel):
     root: HttpUrl = HttpUrl("http://localhost:8181")
     tiled_service_account_check: str
+    submit_task_check: str
+    admin_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 Annotated, Any, Self, cast
 
 import aiohttp
 from aiohttp import ClientSession
+from fastapi import Depends, HTTPException, Request
+from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
 
 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:
@@ -58,6 +63,39 @@ 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=HTTP_403_FORBIDDEN)
+
+    async def is_admin(self, token: str) -> bool:
+        return await self._call_opa(self._conf.admin_check, {"token": token})
+
+
+class OpaUserClient:
+    client: OpaClient
+    token: str
+
+    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 admin(self) -> bool:
+        return await self.client.is_admin(self.token)
+
 
 async def validate_tiled_config(
     tiled: ServiceAccount | str | None, oidc: OIDCConfig | None, opa: OpaClient | None
@@ -74,3 +112,22 @@ 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=HTTP_401_UNAUTHORIZED)
+        return OpaUserClient(opa, token)
+    return None
+
+
+async def submit_permission(
+    opa: Annotated[OpaUserClient | None, Depends(opa)],
+    task_request: TaskRequest,
+):
+    if opa:
+        await opa.can_submit_task(task_request)

src/blueapi/service/main.py

@@ -40,7 +40,13 @@
 from blueapi.worker import TrackableTask, WorkerState
 from blueapi.worker.event import TaskStatusEnum
 
-from .authorization import OpaClient, validate_tiled_config
+from .authorization import (
+    OpaClient,
+    OpaUserClient,
+    opa,
+    submit_permission,
+    validate_tiled_config,
+)
 from .model import (
     DeviceModel,
     DeviceResponse,
@@ -146,6 +152,33 @@ def get_app(config: ApplicationConfig):
     return app
 
 
+def access_task_permission(
+    opa: Annotated[OpaUserClient | None, Depends(opa)],
+    task_id: str,
+    fedid: Fedid,
+    runner: Annotated[WorkerDispatcher, Depends(_runner)],
+):
+    task = runner.run(interface.get_task_by_id, task_id)
+
+    if opa and not opa.admin() and (task and fedid != task.task.metadata.get("user")):
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+
+
+# start_task_permission is used when there is WorkerTask
+def start_task_permission(
+    task: WorkerTask,
+    opa: Annotated[OpaUserClient, Depends(opa)],
+    fedid: Fedid,
+    runner: Annotated[WorkerDispatcher, Depends(_runner)],
+):
+    if not task.task_id:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
+            detail="No task id provided",
+        )
+    access_task_permission(opa, task.task_id, fedid, runner)
+
+
 async def on_key_error_404(_: Request, __: Exception):
     return JSONResponse(
         status_code=status.HTTP_404_NOT_FOUND,
@@ -271,13 +304,13 @@ def submit_task(
     request: Request,
     response: Response,
     task_request: Annotated[TaskRequest, Body(..., examples=[example_task_request])],
+    _: Annotated[None, Depends(submit_permission)],
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
-    user: Fedid,
+    fedid: Fedid,
 ) -> TaskResponse:
     """Submit a task to the worker."""
     try:
-        user = user or "Unknown"
-        task_id: str = runner.run(interface.submit_task, task_request, {"user": user})
+        task_id: str = runner.run(interface.submit_task, task_request, {"user": fedid})
         response.headers["Location"] = f"{request.url}/{task_id}"
         return TaskResponse(task_id=task_id)
     except ValidationError as e:
@@ -309,6 +342,7 @@ def submit_task(
 @start_as_current_span(TRACER, "task_id")
 def delete_submitted_task(
     task_id: str,
+    _: Annotated[None, Depends(access_task_permission)],
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
 ) -> TaskResponse:
     return TaskResponse(task_id=runner.run(interface.clear_task, task_id))
@@ -326,6 +360,7 @@ def validate_task_status(v: str) -> TaskStatusEnum:
 @secure_router.get("/tasks", status_code=status.HTTP_200_OK, tags=[Tag.TASK])
 @start_as_current_span(TRACER)
 def get_tasks(
+    fedid: Fedid,
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
     task_status: str | SkipJsonSchema[None] = None,
 ) -> TasksListResponse:
@@ -346,6 +381,9 @@ def get_tasks(
         tasks = runner.run(interface.get_tasks_by_status, desired_status)
     else:
         tasks = runner.run(interface.get_tasks)
+
+    tasks = [t for t in tasks if t.task.metadata.get("user") == fedid]
+
     return TasksListResponse(tasks=tasks)
 
 
@@ -363,6 +401,7 @@ def get_tasks(
 def set_active_task(
     request: Request,
     task: WorkerTask,
+    _: Annotated[None, Depends(start_task_permission)],
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
 ) -> WorkerTask:
     """Set a task to active status, the worker should begin it as soon as possible.
@@ -393,6 +432,7 @@ def get_passthrough_headers(request: Request) -> dict[str, str]:
 @start_as_current_span(TRACER, "task_id")
 def get_task(
     task_id: str,
+    _: Annotated[None, Depends(access_task_permission)],
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
 ) -> TrackableTask:
     """Retrieve a task"""
@@ -470,6 +510,9 @@ def get_state(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> WorkerSt
 def set_state(
     state_change_request: StateChangeRequest,
     response: Response,
+    fedid: Fedid,
+    opa: Annotated[OpaUserClient, Depends(opa)],
+    # _: Annotated[None, Depends(access_task_permission)],
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
 ) -> WorkerState:
     """
@@ -496,6 +539,16 @@ def set_state(
         current_state in _ALLOWED_TRANSITIONS
         and new_state in _ALLOWED_TRANSITIONS[current_state]
     ):
+        active = runner.run(interface.get_active_task)
+
+        if (
+            opa
+            and not opa.admin()
+            and active
+            and active.task.metadata.get("user") != fedid
+        ):
+            raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
+
         if new_state == WorkerState.PAUSED:
             runner.run(interface.pause_worker, state_change_request.defer)
         elif new_state == WorkerState.RUNNING: