Implement auth for inputstep

Author: Ben Elam

orchestrator/services/processes.py

@@ -523,7 +523,11 @@ def thread_resume_process(
 
     form = pstat.log[0].form
 
-    user_input = post_form(form, pstat.state.unwrap(), user_inputs)
+    # Add OIDC user to state to be processed by form for authorization
+    state = pstat.state.unwrap()
+    state["__process_user"] = user_model
+
+    user_input = post_form(form, state, user_inputs)
 
     if user:
         pstat.update(current_user=user)

orchestrator/workflow.py

@@ -13,13 +13,14 @@
 
 
 from __future__ import annotations
-
 import contextvars
 import functools
 import inspect
 import secrets
 from collections.abc import Callable
 from dataclasses import asdict, dataclass
+from functools import update_wrapper
+from http import HTTPStatus
 from itertools import dropwhile
 from typing import (
     Any,
@@ -40,6 +41,7 @@
 
 from nwastdlib import const, identity
 from oauth2_lib.fastapi import OIDCUserModel
+from orchestrator.api.error_handling import raise_status
 from orchestrator.config.assignee import Assignee
 from orchestrator.db import db, transactional
 from orchestrator.services.settings import get_engine_settings
@@ -99,7 +101,7 @@ def __call__(self) -> NoReturn: ...
 
 
 def make_step_function(
-    f: Callable, name: str, form: InputFormGenerator | None = None, assignee: Assignee | None = Assignee.SYSTEM
+    f: Callable, name: str, form: InputFormGenerator | None = None, assignee: Assignee | None = Assignee.SYSTEM, authorize_callback: Callable[[OIDCUserModel | None], bool] | None = None
 ) -> Step:
     step_func = cast(Step, f)
 
@@ -166,14 +168,39 @@ def __repr__(self) -> str:
         return f"StepList [{', '.join(repr(x) for x in self)}]"
 
 
-def _handle_simple_input_form_generator(f: StateInputStepFunc) -> StateInputFormGenerator:
+def _handle_simple_input_form_generator(f: StateInputStepFunc, authorize_callback: Callable[[OIDCUserModel | None], bool] | None = None) -> StateInputFormGenerator:
+    """Processes f into a form generator and injects a pre-hook for user authorization"""
+    def authorize_user_from_state(state: State) -> None:
+        logger.error("authorize_user_from_state: called")
+        user_model = state.pop("__process_user", None)
+        if user_model is not None:
+            user_model = cast(OIDCUserModel, user_model)
+        else:
+            logger.error("authorize_user_from_state: no user model")
+
+        if authorize_callback is not None:
+            logger.error("authorize_user_from_state: callback found")
+            authorize_callback(user_model)
+            if not authorize_callback(user_model):
+                logger.error("authorize_user_from_state: FORBIDDEN")
+                #TODO not sure that step name is available here, but could put it on state?
+                raise_status(HTTPStatus.FORBIDDEN, "User is not authorized to execute step")
+        else:
+            logger.error("authorize_user_from_state: no callback!")
+
     if inspect.isgeneratorfunction(f):
-        return cast(StateInputFormGenerator, f)
+        def generator_wrapper(state: State):
+            authorize_user_from_state(state)
+            return f(state)
+
+        update_wrapper(generator_wrapper, f)
+        return cast(StateInputFormGenerator, generator_wrapper)
     if inspect.isgenerator(f):
         raise ValueError("Got a generator object instead of function, this is not correct")
 
     # If f is a SimpleInputFormGenerator convert to new style generator function
     def form_generator(state: State) -> FormGenerator:
+        authorize_user_from_state(state)
         user_input: FormPage = yield cast(StateSimpleInputFormGenerator, f)(state)
         return user_input.model_dump()
 
@@ -270,7 +297,7 @@ def wrapper(state: State) -> Process:
     return decorator
 
 
-def inputstep(name: str, assignee: Assignee) -> Callable[[InputStepFunc], Step]:
+def inputstep(name: str, assignee: Assignee, authorize_callback: Callable[[OIDCUserModel | None], bool] | None = None) -> Callable[[InputStepFunc], Step]:
     """Add user input step to workflow.
 
     IMPORTANT: In contrast to other workflow steps, the `@inputstep` wrapped function will not run in the
@@ -291,7 +318,7 @@ def decorator(func: InputStepFunc) -> Step:
         def wrapper(state: State) -> FormGenerator:
             form_generator_in_form_inject_args = form_inject_args(func)
 
-            form_generator = _handle_simple_input_form_generator(form_generator_in_form_inject_args)
+            form_generator = _handle_simple_input_form_generator(form_generator_in_form_inject_args, authorize_callback=authorize_callback)
 
             return form_generator(state)
 

test/unit_tests/api/test_processes.py

@@ -6,6 +6,8 @@
 from uuid import uuid4
 
 import pytest
+from pydantic_forms.core import FormPage
+from pydantic_forms.types import State
 from sqlalchemy import select
 
 from oauth2_lib.fastapi import OIDCUserModel
@@ -22,7 +24,7 @@
 from orchestrator.services.settings import get_engine_settings
 from orchestrator.settings import app_settings
 from orchestrator.targets import Target
-from orchestrator.workflow import ProcessStatus, done, init, step, workflow
+from orchestrator.workflow import ProcessStatus, done, init, inputstep, step, workflow
 from test.unit_tests.helpers import URL_STR_TYPE
 from test.unit_tests.workflows import WorkflowInstanceForTests
 
@@ -593,3 +595,50 @@ def unauthorized_workflow():
     with WorkflowInstanceForTests(unauthorized_workflow, "unauthorized_workflow"):
         response = test_client.post("/api/processes/unauthorized_workflow", json=[{}])
         assert HTTPStatus.FORBIDDEN == response.status_code
+
+
+def test_inputstep_authorization(test_client):
+    def disallow(_: OIDCUserModel | None = None) -> bool:
+        return False
+    
+    def allow(_: OIDCUserModel | None = None) -> bool:
+        return True
+
+    class ConfirmForm(FormPage):
+        confirm: bool
+
+    @inputstep("unauthorized_resume", assignee=Assignee.SYSTEM, authorize_callback=disallow)
+    def unauthorized_resume(state: State) -> State:
+        user_input = yield ConfirmForm
+        return user_input.model_dump()
+
+    @inputstep("authorized_resume", assignee=Assignee.SYSTEM, authorize_callback=allow)
+    def authorized_resume(state: State) -> State:
+        user_input = yield ConfirmForm
+        return user_input.model_dump()
+
+    @inputstep("noauth_resume", assignee=Assignee.SYSTEM)
+    def noauth_resume(state: State) -> State:
+        user_input = yield ConfirmForm
+        return user_input.model_dump()
+
+    @workflow("test_auth_workflow", target=Target.CREATE)
+    def test_auth_workflow():
+        return init >> noauth_resume >> authorized_resume >> unauthorized_resume >> done
+
+    with WorkflowInstanceForTests(test_auth_workflow, "test_auth_workflow"):
+        response = test_client.post("/api/processes/test_auth_workflow", json=[{}])
+        assert HTTPStatus.CREATED == response.status_code
+        process_id = response.json()["id"]
+        # No auth succeeds
+        response = test_client.put(f"/api/processes/{process_id}/resume", json=[{"confirm": True}])
+        assert HTTPStatus.NO_CONTENT == response.status_code
+        # Authorized succeeds
+        response = test_client.put(f"/api/processes/{process_id}/resume", json=[{"confirm": True}])
+        assert HTTPStatus.NO_CONTENT == response.status_code
+        # Unauthorized fails
+        response = test_client.put(f"/api/processes/{process_id}/resume", json=[{"confirm": True}])
+        assert HTTPStatus.FORBIDDEN == response.status_code
+
+    #TODO test how this interacts with passing a different callback to @authorize_workflow
+    # These should be as functionally independent as possible.