diff --git a/changelog/fixed/joining-pathfinder.md b/changelog/fixed/joining-pathfinder.md new file mode 100644 index 0000000000..ba1580578e --- /dev/null +++ b/changelog/fixed/joining-pathfinder.md @@ -0,0 +1 @@ +Preserve the result of augmented assignments such as `path /= "child"` in the workflow sandbox, while retaining restrictions on replacement objects. diff --git a/temporalio/worker/workflow_sandbox/_restrictions.py b/temporalio/worker/workflow_sandbox/_restrictions.py index 6a061cac24..9ca5ad4e4f 100644 --- a/temporalio/worker/workflow_sandbox/_restrictions.py +++ b/temporalio/worker/workflow_sandbox/_restrictions.py @@ -958,9 +958,16 @@ def __init__( super().__init__(access_func, fallback_func=fallback_func) def bind_f(instance: _RestrictedProxy, obj: Any) -> Callable: - def i_op(self: Any, other: Any) -> _RestrictedProxy: - access_func(self, other) # type: ignore - return instance + def i_op(self: Any, other: Any) -> Any: + result = access_func(self, other) # type: ignore + if result is self: + return instance + if _is_restrictable(result): + state = _RestrictionState.from_proxy(instance) + return _RestrictedProxy( + state.name, result, state.context, state.matcher + ) + return result return i_op.__get__(obj, type(obj)) # type: ignore diff --git a/tests/worker/workflow_sandbox/test_importer.py b/tests/worker/workflow_sandbox/test_importer.py index df183b2079..205a8df3b1 100644 --- a/tests/worker/workflow_sandbox/test_importer.py +++ b/tests/worker/workflow_sandbox/test_importer.py @@ -25,6 +25,21 @@ from .testmodules import restrictions +def test_workflow_sandbox_importer_path_in_place_division(): + context = RestrictionContext() + context.is_runtime = True + with Importer(SandboxRestrictions.default, context).applied(): + from pathlib import Path + + path = original = Path("root") + path /= "child" + + assert str(path) == str(Path("root") / "child") + assert str(original) == "root" + with pytest.raises(RestrictedWorkflowAccessError): + path.exists() + + @pytest.mark.parametrize("first", ["host", "sandbox"]) def test_workflow_sandbox_importer_pydantic_constraints(first: str): pytest.importorskip("pydantic", minversion="2") diff --git a/tests/worker/workflow_sandbox/test_restrictions.py b/tests/worker/workflow_sandbox/test_restrictions.py index ec001ba194..c7424f4192 100644 --- a/tests/worker/workflow_sandbox/test_restrictions.py +++ b/tests/worker/workflow_sandbox/test_restrictions.py @@ -1,8 +1,11 @@ from __future__ import annotations +import operator import pathlib +from collections.abc import Callable from dataclasses import dataclass -from typing import ClassVar +from datetime import timedelta +from typing import Any, ClassVar import pytest @@ -82,14 +85,47 @@ def test_restricted_proxy_in_place_operations_preserve_proxy(): RestrictionContext(), SandboxMatcher(), ) - values = restricted_list([1]) + values = original = restricted_list([1]) values += [2] assert type(values) is _RestrictedProxy + assert values is original assert RestrictionContext.unwrap_if_proxied(values) == [1, 2] +@pytest.mark.parametrize( + "value, other, operation", + [ + ((1,), (2,), operator.iadd), + ((1,), 2, operator.imul), + (frozenset({1}), {2}, operator.ior), + ], +) +def test_restricted_proxy_in_place_operations_replace_value( + value: Any, other: Any, operation: Callable[[Any, Any], Any] +): + proxy = _RestrictedProxy("value", value, RestrictionContext(), SandboxMatcher()) + + result = operation(proxy, other) + + assert RestrictionContext.unwrap_if_proxied(result) == operation(value, other) + assert type(result) is _RestrictedProxy + assert result is not proxy + assert RestrictionContext.unwrap_if_proxied(proxy) is value + + +def test_restricted_proxy_in_place_operation_returns_primitive(): + value = _RestrictedProxy( + "timedelta", timedelta(seconds=4), RestrictionContext(), SandboxMatcher() + ) + + value /= timedelta(seconds=2) + + assert value == 2.0 + assert type(value) is float + + def test_workflow_sandbox_restricted_proxy(): obj_class = _RestrictedProxy( "RestrictableObject", diff --git a/tests/worker/workflow_sandbox/test_runner.py b/tests/worker/workflow_sandbox/test_runner.py index acf1fb7e46..d357a68724 100644 --- a/tests/worker/workflow_sandbox/test_runner.py +++ b/tests/worker/workflow_sandbox/test_runner.py @@ -232,6 +232,11 @@ async def test_workflow_sandbox_restrictions(client: Client): # General library allowed calls "import datetime\ndatetime.date(2001, 1, 1)", "import uuid\nuuid.uuid5(uuid.NAMESPACE_DNS, 'example.com')", + "from pathlib import Path\n" + "path = original = Path('root')\n" + "path /= 'child'\n" + "assert str(path) == str(Path('root') / 'child')\n" + "assert str(original) == 'root'", # Other imports we want to allow "from concurrent.futures import ThreadPoolExecutor", ]