diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py index cbb337f650..20f4d8e2e4 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py @@ -504,6 +504,7 @@ async def start_workflow_orchestration( # keys, so stripping them here keeps untrusted input off the orchestrator's # trusted-deserialization path (see strip_subworkflow_markers). client_input = strip_subworkflow_markers(client_input) + client_input = strip_pickle_markers(client_input) instance_id = await client.start_new(orchestrator_name, client_input=client_input) diff --git a/python/packages/azurefunctions/tests/test_azurefunctions_workflow_initial_input.py b/python/packages/azurefunctions/tests/test_azurefunctions_workflow_initial_input.py new file mode 100644 index 0000000000..7ec3a9f492 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_azurefunctions_workflow_initial_input.py @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Behavioral tests for Azure Functions workflow HTTP initial input.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, TypeVar +from unittest.mock import AsyncMock, Mock, patch + +from agent_framework import Executor, Workflow, WorkflowBuilder, WorkflowContext, handler + +from agent_framework_azurefunctions import AgentFunctionApp + +FuncT = TypeVar("FuncT", bound=Callable[..., Any]) + + +def _identity_decorator(*args: Any, **kwargs: Any) -> Callable[[FuncT], FuncT]: + def decorator(func: FuncT) -> FuncT: + return func + + return decorator + + +class _Start(Executor): + def __init__(self) -> None: + super().__init__(id="start") + + @handler(input=str | dict, workflow_output=str) + async def run(self, message: str | dict, ctx: WorkflowContext) -> None: + pass + + +def _capture_run_handler(workflow: Workflow) -> Callable[..., Any]: + captured_routes: dict[str, Callable[..., Any]] = {} + + def capture_route(*args: Any, **kwargs: Any) -> Callable[[FuncT], FuncT]: + def decorator(func: FuncT) -> FuncT: + captured_routes[kwargs["route"]] = func + return func + + return decorator + + with ( + patch.object(AgentFunctionApp, "function_name", new=_identity_decorator), + patch.object(AgentFunctionApp, "route", new=capture_route), + patch.object(AgentFunctionApp, "durable_client_input", new=_identity_decorator), + patch.object(AgentFunctionApp, "activity_trigger", new=_identity_decorator), + patch.object(AgentFunctionApp, "orchestration_trigger", new=_identity_decorator), + ): + AgentFunctionApp(workflow=workflow, enable_health_check=False) + + return captured_routes[f"workflow/{workflow.name}/run"] + + +async def test_workflow_run_route_neutralizes_reserved_marker_shaped_input() -> None: + """The workflow run route schedules neutralized framework-reserved metadata.""" + executor = _Start() + workflow = WorkflowBuilder(name="input_boundary", start_executor=executor, output_from=[executor]).build() + handler = _capture_run_handler(workflow) + request = Mock() + request.get_json.return_value = { + "__pickled__": "not-checkpoint-data", + "__type__": "builtins:int", + } + request.url = "https://example.test/api/workflow/input_boundary/run" + client = AsyncMock() + client.start_new.return_value = "instance-1" + + await handler(request, client) + + assert client.start_new.await_args.kwargs["client_input"] is None diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py index 346d6a3985..2d5d7e2d91 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py @@ -746,7 +746,7 @@ def _coerce_initial_input(workflow: Workflow, raw_value: Any) -> Any: input_type = _select_primary_input_type(start_executor) if input_type is None: - return raw_value + return strip_pickle_markers(raw_value) # The initial payload is untrusted external input (HTTP body / client input) with no # legitimate checkpoint type markers, so neutralize any pickle-marker injection before # it can reach deserialize_value() inside reconstruct_to_type() (avoids pickle RCE). diff --git a/python/packages/durabletask/tests/test_durabletask_workflow_initial_input.py b/python/packages/durabletask/tests/test_durabletask_workflow_initial_input.py new file mode 100644 index 0000000000..a948d64616 --- /dev/null +++ b/python/packages/durabletask/tests/test_durabletask_workflow_initial_input.py @@ -0,0 +1,112 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Behavioral tests for Durable Task workflow initial input.""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from typing import Any + +from agent_framework import Executor, Workflow, WorkflowBuilder, WorkflowContext, handler + +from agent_framework_durabletask import execute_workflow_activity, run_workflow_orchestrator + + +class _InlineWorkflowHost: + """Run activity tasks inline while exercising the public orchestration surface.""" + + def __init__(self, workflow: Workflow) -> None: + self.workflow = workflow + + @property + def instance_id(self) -> str: + return "test-instance" + + @property + def is_replaying(self) -> bool: + return False + + @property + def supports_event_streaming(self) -> bool: + return False + + @property + def current_utc_datetime(self) -> datetime: + return datetime.now(timezone.utc) + + def prepare_agent_task(self, executor_id: str, message: str, orchestration_instance_id: str) -> Any: + raise AssertionError("This test workflow has no agent executors") + + def prepare_activity_task(self, activity_name: str, input_json: str) -> str: + activity_input = json.loads(input_json) + executor = self.workflow.executors[activity_input["executor_id"]] + return execute_workflow_activity(executor, input_json, self.workflow) + + def call_sub_orchestrator(self, name: str, input: Any, instance_id: str | None = None) -> Any: + raise AssertionError("This test workflow has no sub-workflows") + + def task_all(self, tasks: list[Any]) -> list[Any]: + return tasks + + def task_any(self, tasks: list[Any]) -> Any: + raise AssertionError("This test workflow does not wait for competing tasks") + + def wait_for_external_event(self, name: str) -> Any: + raise AssertionError("This test workflow has no external events") + + def create_timer(self, fire_at: datetime) -> Any: + raise AssertionError("This test workflow has no timers") + + def set_custom_status(self, status: Any) -> None: + pass + + def new_uuid(self) -> str: + return "test-uuid" + + def cancel_task(self, task: Any) -> None: + pass + + def get_task_result(self, task: Any) -> Any: + return task + + +class _NullableUnionStart(Executor): + def __init__(self) -> None: + super().__init__(id="start") + + @handler(input=str | dict | None, workflow_output=str) + async def run(self, message: str | dict[str, Any] | None, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output("neutralized" if message is None else "accepted") + + +def _run_inline(workflow: Workflow, initial_input: Any) -> list[Any] | dict[str, Any]: + host = _InlineWorkflowHost(workflow) + orchestration = run_workflow_orchestrator(host, workflow, initial_input) + yielded = next(orchestration) + + while True: + try: + yielded = orchestration.send(yielded) + except StopIteration as completed: + return completed.value + + +def test_reserved_marker_shaped_initial_input_is_neutralized() -> None: + """Framework-reserved serialization metadata is not delivered as application input.""" + executor = _NullableUnionStart() + workflow = WorkflowBuilder(start_executor=executor, output_from=[executor]).build() + initial_input = { + "__pickled__": "not-checkpoint-data", + "__type__": "builtins:int", + } + + assert _run_inline(workflow, initial_input) == ["neutralized"] + + +def test_regular_union_initial_input_is_preserved() -> None: + """Ordinary JSON input keeps the union-typed workflow's existing behavior.""" + executor = _NullableUnionStart() + workflow = WorkflowBuilder(start_executor=executor, output_from=[executor]).build() + + assert _run_inline(workflow, {"message": "hello"}) == ["accepted"]