Sandbox fixes (#197)

Author: cretz

temporalio/worker/workflow_sandbox/_importer.py

@@ -67,6 +67,9 @@ def __init__(
         self.new_modules: Dict[str, types.ModuleType] = {
             "sys": sys,
             "builtins": builtins,
+            # Even though we don't want to, we have to have __main__ because
+            # stdlib packages like inspect and others expect it to be present
+            "__main__": types.ModuleType("__main__"),
         }
         self.modules_checked_for_restrictions: Set[str] = set()
         self.import_func = self._import if not LOG_TRACE else self._traced_import
@@ -334,7 +337,7 @@ def unapplied(self) -> Iterator[None]:
 
 
 class _ThreadLocalSysModules(
-    _ThreadLocalOverride[MutableMapping[str, types.ModuleType]],
+    _ThreadLocalOverride[Dict[str, types.ModuleType]],
     MutableMapping[str, types.ModuleType],
 ):
     def __contains__(self, key: object) -> bool:
@@ -355,6 +358,35 @@ def __iter__(self) -> Iterator[str]:
     def __setitem__(self, key: str, value: types.ModuleType) -> None:
         self.current[key] = value
 
+    # Below methods are not in mutable mapping. Python chose not to put
+    # everything in MutableMapping they do in dict (see
+    # https://bugs.python.org/issue22101). So when someone calls
+    # sys.modules.copy() it breaks (which is exactly what the inspect module
+    # does sometimes).
+
+    def __or__(
+        self, other: Mapping[str, types.ModuleType]
+    ) -> Dict[str, types.ModuleType]:
+        if sys.version_info < (3, 9):
+            raise NotImplementedError
+        return self.current.__or__(other)
+
+    def __ior__(
+        self, other: Mapping[str, types.ModuleType]
+    ) -> Dict[str, types.ModuleType]:
+        if sys.version_info < (3, 9):
+            raise NotImplementedError
+        return self.current.__ior__(other)
+
+    __ror__ = __or__
+
+    def copy(self) -> Dict[str, types.ModuleType]:
+        return self.current.copy()
+
+    @classmethod
+    def fromkeys(cls, *args, **kwargs) -> Any:
+        return dict.fromkeys(*args, **kwargs)
+
 
 _thread_local_sys_modules = _ThreadLocalSysModules(sys.modules)
 

temporalio/worker/workflow_sandbox/_runner.py

@@ -156,7 +156,9 @@ def _run_code(self, code: str, **extra_globals: Any) -> None:
         for k, v in extra_globals.items():
             self.globals_and_locals[k] = v
         try:
+            temporalio.workflow.unsafe._set_in_sandbox(True)
             exec(code, self.globals_and_locals, self.globals_and_locals)
         finally:
+            temporalio.workflow.unsafe._set_in_sandbox(False)
             for k, v in extra_globals.items():
                 del self.globals_and_locals[k]

temporalio/workflow.py

@@ -713,6 +713,7 @@ async def wait_condition(
 
 
 _sandbox_unrestricted = threading.local()
+_in_sandbox = threading.local()
 
 
 class unsafe:
@@ -723,6 +724,19 @@ class unsafe:
     def __init__(self) -> None:  # noqa: D107
         raise NotImplementedError
 
+    @staticmethod
+    def in_sandbox() -> bool:
+        """Whether the code is executing on a sandboxed thread.
+
+        Returns:
+            True if the code is executing in the sandbox thread.
+        """
+        return getattr(_in_sandbox, "value", False)
+
+    @staticmethod
+    def _set_in_sandbox(v: bool) -> None:
+        _in_sandbox.value = v
+
     @staticmethod
     def is_replaying() -> bool:
         """Whether the workflow is currently replaying.

tests/worker/test_workflow.py

@@ -51,7 +51,6 @@
 from temporalio.converter import (
     DataConverter,
     DefaultFailureConverterWithEncodedAttributes,
-    FailureConverter,
     PayloadCodec,
     PayloadConverter,
 )

tests/worker/workflow_sandbox/test_importer.py

@@ -1,6 +1,12 @@
+import sys
+
 import pytest
 
-from temporalio.worker.workflow_sandbox._importer import Importer
+from temporalio.worker.workflow_sandbox._importer import (
+    Importer,
+    _thread_local_sys_modules,
+    _ThreadLocalSysModules,
+)
 from temporalio.worker.workflow_sandbox._restrictions import (
     RestrictedWorkflowAccessError,
     RestrictionContext,
@@ -76,3 +82,22 @@ def test_workflow_sandbox_importer_invalid_module_members():
         err.value.qualified_name
         == "tests.worker.workflow_sandbox.testmodules.invalid_module_members.invalid_function.__call__"
     )
+
+
+def test_thread_local_sys_module_attrs():
+    if sys.version_info < (3, 9):
+        pytest.skip("Dict or methods only in >= 3.9")
+    # Python chose not to put everything in MutableMapping they do in dict, see
+    # https://bugs.python.org/issue22101. Therefore we manually confirm that
+    # every attribute of sys modules is also in thread local sys modules to
+    # ensure compatibility.
+    for attr in dir(sys.modules):
+        getattr(_thread_local_sys_modules, attr)
+
+    # Let's also test "or" and "copy"
+    norm = {"foo": 123}
+    thread_local = _ThreadLocalSysModules({"foo": 123})
+    assert (norm | {"bar": 456}) == (thread_local | {"bar": 456})
+    norm |= {"baz": 789}
+    thread_local |= {"baz": 789}
+    assert norm.copy() == thread_local.copy()

tests/worker/workflow_sandbox/test_runner.py

@@ -2,6 +2,7 @@
 
 import dataclasses
 import functools
+import inspect
 import os
 import time
 import uuid
@@ -267,6 +268,48 @@ async def test_workflow_sandbox_operator(client: Client):
         )
 
 
+global_in_sandbox = workflow.unsafe.in_sandbox()
+
+
+@workflow.defn
+class InSandboxWorkflow:
+    def __init__(self) -> None:
+        assert global_in_sandbox
+        assert workflow.unsafe.in_sandbox()
+
+    @workflow.run
+    async def run(self) -> None:
+        assert global_in_sandbox
+        assert workflow.unsafe.in_sandbox()
+
+
+async def test_workflow_sandbox_assert(client: Client):
+    async with new_worker(client, InSandboxWorkflow) as worker:
+        assert not global_in_sandbox
+        assert not workflow.unsafe.in_sandbox()
+        await client.execute_workflow(
+            InSandboxWorkflow.run,
+            id=f"workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+
+
+@workflow.defn
+class AccessStackWorkflow:
+    @workflow.run
+    async def run(self) -> str:
+        return inspect.stack()[0].function
+
+
+async def test_workflow_sandbox_access_stack(client: Client):
+    async with new_worker(client, AccessStackWorkflow) as worker:
+        assert "run" == await client.execute_workflow(
+            AccessStackWorkflow.run,
+            id=f"workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+
+
 def new_worker(
     client: Client,
     *workflows: Type,