Preserve augmented-assignment results in the workflow sandbox

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.

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
 

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")

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",

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",
         ]