Make auth callbacks async (#1239)

* Make auth callbacks async

* Fix workflow tests

* Update docs

* Reorganize broadcast_invalidate_status_counts

* Run DB calls in separate thread

* Update documentation for async callback

---------

Co-authored-by: Ben Elam <baelam@es.net>
Co-authored-by: Mark90 <mark_moes@live.nl>

Author: 3 people

docs/reference-docs/auth-backend-and-frontend.md

@@ -275,7 +275,8 @@ app.register_graphql_authorization(graphql_authorization_instance)
     Role-based access control for workflows is currently in beta.
     Initial support has been added to the backend, but the feature is not fully communicated through the UI yet.
 
-Certain `orchestrator-core` decorators accept authorization callbacks of type `type Authorizer = Callable[OIDCUserModel, bool]`, which return True when the input user is authorized, otherwise False.
+Certain `orchestrator-core` decorators accept authorization callbacks of type `type Authorizer = Callable[[OIDCUserModel | None], Awaitable[bool]]`, which return True when the input user is authorized, otherwise False.
+In other words, authorization callbacks are async, take a nullable OIDCUserModel (or subclass) as argument, and return a bool.
 
 A table (below) is available for comparing possible configuration states with the policy that will be enforced.
 
@@ -528,8 +529,11 @@ are prioritized in different workflow and inputstep configurations.
 Assume we have the following function that can be used to create callbacks:
 
 ```python
-def allow_roles(*roles) -> Callable[OIDCUserModel|None, bool]:
-    def f(user: OIDCUserModel) -> bool:
+from oauth2_lib.fastapi import OIDCUserModel
+from orchestrator.workflows.utils import Authorizer
+
+def allow_roles(*roles) -> Authorizer:
+    async def f(user: OIDCUserModel) -> bool:
         if is_admin(user):  # Relative to your authorization provider
             return True
         for role in roles:

orchestrator/api/api_v1/endpoints/processes.py

@@ -13,6 +13,7 @@
 
 """Module that implements process related API endpoints."""
 
+import asyncio
 import struct
 import zlib
 from http import HTTPStatus
@@ -62,10 +63,12 @@
 from orchestrator.websocket import (
     WS_CHANNELS,
     broadcast_invalidate_status_counts,
+    broadcast_invalidate_status_counts_async,
     broadcast_process_update_to_websocket,
     websocket_manager,
 )
 from orchestrator.workflow import ProcessStat, ProcessStatus, StepList, Workflow
+from orchestrator.workflows import get_workflow
 from pydantic_forms.types import JSON, State
 
 router = APIRouter()
@@ -175,16 +178,29 @@ def delete(process_id: UUID) -> None:
     status_code=HTTPStatus.CREATED,
     dependencies=[Depends(check_global_lock, use_cache=False)],
 )
-def new_process(
+async def new_process(
     workflow_key: str,
     request: Request,
     json_data: list[dict[str, Any]] | None = Body(...),
     user: str = Depends(user_name),
     user_model: OIDCUserModel | None = Depends(authenticate),
 ) -> dict[str, UUID]:
     broadcast_func = api_broadcast_process_data(request)
-    process_id = start_process(
-        workflow_key, user_inputs=json_data, user_model=user_model, user=user, broadcast_func=broadcast_func
+
+    workflow = get_workflow(workflow_key)
+    if not workflow:
+        raise_status(HTTPStatus.NOT_FOUND, "Workflow does not exist")
+
+    if not await workflow.authorize_callback(user_model):
+        raise_status(HTTPStatus.FORBIDDEN, f"User is not authorized to execute '{workflow_key}' workflow")
+
+    process_id = await asyncio.to_thread(
+        start_process,
+        workflow_key,
+        user_inputs=json_data,
+        user_model=user_model,
+        user=user,
+        broadcast_func=broadcast_func,
     )
 
     return {"id": process_id}
@@ -196,31 +212,31 @@ def new_process(
     status_code=HTTPStatus.NO_CONTENT,
     dependencies=[Depends(check_global_lock, use_cache=False)],
 )
-def resume_process_endpoint(
+async def resume_process_endpoint(
     process_id: UUID,
     request: Request,
     json_data: JSON = Body(...),
     user: str = Depends(user_name),
     user_model: OIDCUserModel | None = Depends(authenticate),
 ) -> None:
-    process = _get_process(process_id)
+    process = await asyncio.to_thread(_get_process, process_id)
 
     if not can_be_resumed(process.last_status):
         raise_status(HTTPStatus.CONFLICT, f"Resuming a {process.last_status.lower()} workflow is not possible")
 
     pstat = load_process(process)
     auth_resume, auth_retry = get_auth_callbacks(get_steps_to_evaluate_for_rbac(pstat), pstat.workflow)
     if process.last_status == ProcessStatus.SUSPENDED:
-        if auth_resume is not None and not auth_resume(user_model):
+        if auth_resume is not None and not (await auth_resume(user_model)):
             raise_status(HTTPStatus.FORBIDDEN, "User is not authorized to resume step")
     elif process.last_status in (ProcessStatus.FAILED, ProcessStatus.WAITING):
-        if auth_retry is not None and not auth_retry(user_model):
+        if auth_retry is not None and not (await auth_retry(user_model)):
             raise_status(HTTPStatus.FORBIDDEN, "User is not authorized to retry step")
 
-    broadcast_invalidate_status_counts()
+    await broadcast_invalidate_status_counts_async()
     broadcast_func = api_broadcast_process_data(request)
 
-    resume_process(process, user=user, user_inputs=json_data, broadcast_func=broadcast_func)
+    await asyncio.to_thread(resume_process, process, user=user, user_inputs=json_data, broadcast_func=broadcast_func)
 
 
 @router.post(

orchestrator/graphql/schemas/workflow.py

@@ -38,4 +38,4 @@ async def is_allowed(self, info: OrchestratorInfo) -> bool:
         workflow_table = get_original_model(self, WorkflowTable)
         workflow = get_workflow(workflow_table.name)
 
-        return workflow.authorize_callback(oidc_user)  # type: ignore
+        return await workflow.authorize_callback(oidc_user)  # type: ignore

orchestrator/services/processes.py

@@ -421,10 +421,6 @@ def run() -> WFProcess:
     return process_id
 
 
-def error_message_unauthorized(workflow_key: str) -> str:
-    return f"User is not authorized to execute '{workflow_key}' workflow"
-
-
 def create_process(
     workflow_key: str,
     user_inputs: list[State] | None = None,
@@ -442,9 +438,6 @@ def create_process(
     if not workflow:
         raise_status(HTTPStatus.NOT_FOUND, "Workflow does not exist")
 
-    if not workflow.authorize_callback(user_model):
-        raise_status(HTTPStatus.FORBIDDEN, error_message_unauthorized(workflow_key))
-
     initial_state = {
         "process_id": process_id,
         "reporter": user,

orchestrator/utils/auth.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from typing import TypeAlias, TypeVar
 
 from oauth2_lib.fastapi import OIDCUserModel
@@ -7,4 +7,4 @@
 
 # Can instead use "type Authorizer = ..." in later Python versions.
 T = TypeVar("T", bound=OIDCUserModel)
-Authorizer: TypeAlias = Callable[[T | None], bool]
+Authorizer: TypeAlias = Callable[[T | None], Awaitable[bool]]

orchestrator/websocket/__init__.py

@@ -105,6 +105,19 @@ async def invalidate_subscription_cache(subscription_id: UUID | UUIDstr, invalid
     await broadcast_invalidate_cache({"type": "subscriptions", "id": str(subscription_id)})
 
 
+async def broadcast_invalidate_status_counts_async() -> None:
+    """Broadcast message to invalidate the status counts of the connected websocket clients.
+
+    This breaks the pattern of `sync_` prefixes to maintain backwards compatibility of
+    broadcast_invalidate_status_counts, a sync function.
+    """
+    if not websocket_manager.enabled:
+        logger.debug("WebSocketManager is not enabled. Skip broadcasting through websocket.")
+        return
+
+    await broadcast_invalidate_cache({"type": "processStatusCounts"})
+
+
 def broadcast_invalidate_status_counts() -> None:
     """Broadcast message to invalidate the status counts of the connected websocket clients."""
     if not websocket_manager.enabled:
@@ -148,4 +161,5 @@ async def broadcast_process_update_to_websocket_async(
     "broadcast_process_update_to_websocket_async",
     "WS_CHANNELS",
     "broadcast_invalidate_status_counts",
+    "broadcast_invalidate_status_counts_async",
 ]

orchestrator/workflow.py

@@ -193,7 +193,7 @@ def form_generator(state: State) -> FormGenerator:
     return form_generator
 
 
-def allow(_: OIDCUserModel | None) -> bool:
+async def allow(_: OIDCUserModel | None) -> bool:
     """Default function to return True in absence of user-defined authorize function."""
     return True
 

test/unit_tests/api/test_processes.py

@@ -544,12 +544,26 @@ def test_resume_all_processes_value_error(test_client, mocked_processes_resumeal
 )
 def test_create_process_reporter(test_client, fastapi_app, oidc_user, reporter, expected_user):
     # given
+    async def allow(_: object) -> bool:
+        return True
+
+    fake_workflow = make_workflow(
+        f=lambda _: None,
+        description="fake",
+        initial_input_form=None,
+        target=Target.CREATE,
+        steps=StepList([]),
+        authorize_callback=allow,
+        retry_auth_callback=allow,
+    )
     url_params = {"reporter": reporter} if reporter is not None else {}
     fastapi_depends = {authenticate: lambda: oidc_user}
     with (
+        mock.patch("orchestrator.api.api_v1.endpoints.processes.get_workflow") as mock_get_workflow,
         mock.patch("orchestrator.api.api_v1.endpoints.processes.start_process") as mock_start_process,
         mock.patch.dict(fastapi_app.dependency_overrides, fastapi_depends),
     ):
+        mock_get_workflow.return_value = fake_workflow
         mock_start_process.return_value = uuid.uuid4()
         # when
         response = test_client.post("/api/processes/fake_workflow", json=[], params=url_params)
@@ -609,7 +623,7 @@ def test_new_process_higher_version_invalid(test_client, generic_subscription_1)
 
 
 def test_unauthorized_to_run_process(test_client):
-    def disallow(_: OIDCUserModel | None = None) -> bool:
+    async def disallow(_: OIDCUserModel | None = None) -> bool:
         return False
 
     @workflow("unauthorized_workflow", target=Target.CREATE, authorize_callback=disallow)
@@ -623,10 +637,10 @@ def unauthorized_workflow():
 
 @pytest.fixture
 def authorize_resume_workflow():
-    def disallow(_: OIDCUserModel | None = None) -> bool:
+    async def disallow(_: OIDCUserModel | None = None) -> bool:
         return False
 
-    def allow(_: OIDCUserModel | None = None) -> bool:
+    async def allow(_: OIDCUserModel | None = None) -> bool:
         return True
 
     class ConfirmForm(FormPage):
@@ -719,19 +733,19 @@ def test_unauthorized_resume_input_step(test_client, process_on_unauthorized_res
     assert HTTPStatus.FORBIDDEN == response.status_code
 
 
-def _A(_: OIDCUserModel) -> bool:
+async def _A(_: OIDCUserModel) -> bool:
     return True
 
 
-def _B(_: OIDCUserModel) -> bool:
+async def _B(_: OIDCUserModel) -> bool:
     return True
 
 
-def _C(_: OIDCUserModel) -> bool:
+async def _C(_: OIDCUserModel) -> bool:
     return True
 
 
-def _D(_: OIDCUserModel) -> bool:
+async def _D(_: OIDCUserModel) -> bool:
     return True
 
 
@@ -844,10 +858,10 @@ def test_continue_awaiting_process_endpoint_wrong_process_status(test_client, pr
 
 @pytest.fixture
 def authorize_step_group_retry_workflow():
-    def disallow(_: OIDCUserModel | None = None) -> bool:
+    async def disallow(_: OIDCUserModel | None = None) -> bool:
         return False
 
-    def allow(_: OIDCUserModel | None = None) -> bool:
+    async def allow(_: OIDCUserModel | None = None) -> bool:
         return True
 
     steps = StepList([])
@@ -919,10 +933,10 @@ def test_unauthorized_step_group_retry(test_client, process_on_unretriable_step_
 
 @pytest.fixture
 def authorize_step_retry_workflow():
-    def disallow(_: OIDCUserModel | None = None) -> bool:
+    async def disallow(_: OIDCUserModel | None = None) -> bool:
         return False
 
-    def allow(_: OIDCUserModel | None = None) -> bool:
+    async def allow(_: OIDCUserModel | None = None) -> bool:
         return True
 
     @step("authorized_retry", retry_auth_callback=allow)
@@ -998,10 +1012,10 @@ def test_unauthorized_step_retry(test_client, process_on_unretriable_step):
 
 @pytest.fixture
 def authorize_retrystep_retry_workflow():
-    def disallow(_: OIDCUserModel | None = None) -> bool:
+    async def disallow(_: OIDCUserModel | None = None) -> bool:
         return False
 
-    def allow(_: OIDCUserModel | None = None) -> bool:
+    async def allow(_: OIDCUserModel | None = None) -> bool:
         return True
 
     @retrystep("authorized_retry", retry_auth_callback=allow)

test/unit_tests/graphql/test_workflows.py

@@ -246,7 +246,7 @@ def test_workflows_sort_by_resource_type_desc(test_client):
 def test_workflows_not_allowed(test_client):
     forbidden_workflow_name = "unauthorized_workflow"
 
-    def disallow(_: OIDCUserModel | None = None) -> bool:
+    async def disallow(_: OIDCUserModel | None = None) -> bool:
         return False
 
     @workflow(forbidden_workflow_name, target=Target.CREATE, authorize_callback=disallow)