initial explorations into warnings and errors for dynamic imports and non-passthroughs

Author: VegetarianOrc

temporalio/worker/workflow_sandbox/__init__.py

@@ -58,11 +58,13 @@
     RestrictedWorkflowAccessError,
     SandboxMatcher,
     SandboxRestrictions,
+    DisallowedUnintentionalPassthroughError,
 )
 from ._runner import SandboxedWorkflowRunner
 
 __all__ = [
     "RestrictedWorkflowAccessError",
+    "DisallowedUnintentionalPassthroughError",
     "SandboxedWorkflowRunner",
     "SandboxMatcher",
     "SandboxRestrictions",

temporalio/worker/workflow_sandbox/_importer.py

@@ -40,6 +40,7 @@
 from ._restrictions import (
     RestrictedModule,
     RestrictedWorkflowAccessError,
+    DisallowedUnintentionalPassthroughError,
     RestrictionContext,
     SandboxRestrictions,
 )
@@ -200,6 +201,17 @@ def _import(
 
         # Check module restrictions and passthrough modules
         if full_name not in sys.modules:
+            # Issue a warning if appropriate
+            if (
+                self.restriction_context.in_activation
+                and self._is_import_policy_applied(
+                    temporalio.workflow.SandboxImportPolicy.WARN_ON_DYNAMIC_IMPORT
+                )
+            ):
+                warnings.warn(
+                    f"Module {name} was imported after initial workflow load."
+                )
+
             # Make sure not an entirely invalid module
             self._assert_valid_module(full_name)
 
@@ -282,13 +294,35 @@ def module_configured_passthrough(self, name: str) -> bool:
                 break
         return True
 
+    def _is_import_policy_applied(
+        self, policy: temporalio.workflow.SandboxImportPolicy
+    ) -> bool:
+        applied_policy = temporalio.workflow.unsafe.current_sandbox_import_policy()
+        if applied_policy != temporalio.workflow.SandboxImportPolicy.UNSET:
+            return policy in applied_policy
+
+        return policy in self.restrictions.import_policy
+
     def _maybe_passthrough_module(self, name: str) -> Optional[types.ModuleType]:
         # If imports not passed through and all modules are not passed through
         # and name not in passthrough modules, check parents
         if (
             not temporalio.workflow.unsafe.is_imports_passed_through()
             and not self.module_configured_passthrough(name)
         ):
+            if self._is_import_policy_applied(
+                temporalio.workflow.SandboxImportPolicy.RAISE_ON_NON_PASSTHROUGH
+            ):
+                # TODO(amazzeo): this is not an appropriate error type
+                raise DisallowedUnintentionalPassthroughError(name)
+
+            if self._is_import_policy_applied(
+                temporalio.workflow.SandboxImportPolicy.WARN_ON_NON_PASSTHROUGH
+            ):
+                warnings.warn(
+                    f"Module {name} was not intentionally passed through to the sandbox."
+                )
+
             return None
         # Do the pass through
         with self._unapplied():

temporalio/worker/workflow_sandbox/_restrictions.py

@@ -82,6 +82,23 @@ def default_message(qualified_name: str) -> str:
         )
 
 
+# TODO(amazzeo): is NondeterminisimError appropriate as a subclass?
+class DisallowedUnintentionalPassthroughError(temporalio.workflow.NondeterminismError):
+    def __init__(
+        self, qualified_name: str, *, override_message: Optional[str] = None
+    ) -> None:
+        super().__init__(
+            override_message
+            or DisallowedUnintentionalPassthroughError.default_message(qualified_name)
+        )
+        self.qualified_name = qualified_name
+
+    @staticmethod
+    def default_message(qualified_name: str) -> str:
+        """Get default message for restricted access."""
+        return f"Module {qualified_name} was not intentionally passed through to the sandbox."
+
+
 @dataclass(frozen=True)
 class SandboxRestrictions:
     """Set of restrictions that can be applied to a sandbox."""
@@ -110,6 +127,15 @@ class methods (including __init__, etc). The check compares the against the
     fully qualified path to the item.
     """
 
+    import_policy: temporalio.workflow.SandboxImportPolicy = (
+        temporalio.workflow.SandboxImportPolicy.WARN_ON_DYNAMIC_IMPORT
+    )
+
+    import_policy_all_warnings: ClassVar[temporalio.workflow.SandboxImportPolicy]
+    import_policy_disallow_unintentional_passthrough: ClassVar[
+        temporalio.workflow.SandboxImportPolicy
+    ]
+
     passthrough_all_modules: bool = False
     """
     Pass through all modules, do not sandbox any modules. This is the equivalent
@@ -285,6 +311,7 @@ def access_matcher(
         Returns:
             The matcher if matched.
         """
+
         # We prefer to avoid recursion
         matcher = self
         for v in child_path:
@@ -305,10 +332,12 @@ def access_matcher(
             if not child_matcher:
                 return None
             matcher = child_matcher
+
         if not context.is_runtime and matcher.only_runtime:
             return None
         if not matcher.match_self:
             return None
+
         return matcher
 
     def match_access(
@@ -496,6 +525,15 @@ def with_child_unrestricted(self, *child_path: str) -> SandboxMatcher:
     }
 )
 
+SandboxRestrictions.import_policy_all_warnings = (
+    temporalio.workflow.SandboxImportPolicy.WARN_ON_DYNAMIC_IMPORT
+    | temporalio.workflow.SandboxImportPolicy.WARN_ON_NON_PASSTHROUGH
+)
+SandboxRestrictions.import_policy_disallow_unintentional_passthrough = (
+    temporalio.workflow.SandboxImportPolicy.WARN_ON_DYNAMIC_IMPORT
+    | temporalio.workflow.SandboxImportPolicy.RAISE_ON_NON_PASSTHROUGH
+)
+
 # sys.stdlib_module_names is only available on 3.10+, so we hardcode here. A
 # test will fail if this list doesn't match the latest Python version it was
 # generated against, spitting out the expected list. This is a string instead
@@ -819,6 +857,7 @@ def unwrap_if_proxied(v: Any) -> Any:
     def __init__(self) -> None:
         """Create a restriction context."""
         self.is_runtime = False
+        self.in_activation = False
 
 
 @dataclass

temporalio/worker/workflow_sandbox/_runner.py

@@ -159,6 +159,7 @@ def activate(
         self, act: temporalio.bridge.proto.workflow_activation.WorkflowActivation
     ) -> temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion:
         self.importer.restriction_context.is_runtime = True
+        self.importer.restriction_context.in_activation = True
         try:
             self._run_code(
                 "with __temporal_importer.applied():\n"
@@ -169,6 +170,7 @@ def activate(
             return self.globals_and_locals.pop("__temporal_completion")  # type: ignore
         finally:
             self.importer.restriction_context.is_runtime = False
+            self.importer.restriction_context.in_activation = False
 
     def _run_code(self, code: str, **extra_globals: Any) -> None:
         for k, v in extra_globals.items():

temporalio/workflow.py

@@ -13,7 +13,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from enum import Enum, IntEnum
+from enum import Enum, Flag, IntEnum, auto
 from functools import partial
 from random import Random
 from typing import (
@@ -1412,6 +1412,24 @@ async def wait_condition(
 _imports_passed_through = threading.local()
 
 
+# TODO(amazzeo): this name is confusing with existing import passthrough phrasing.
+class SandboxImportPolicy(Flag):
+    # TODO(amazzeo): Need to establish consistent phrasing about the time this is applied
+    """Defines the behavior taken when modules are imported into the sandbox after the workflow is initially loaded."""
+
+    UNSET = auto()
+    """All imports that do not violate sandbox restrictions are allowed and no warning is generated."""
+    WARN_ON_DYNAMIC_IMPORT = auto()
+    """Issue a warning when an import is triggered in the sandbox after initial workflow load."""
+    WARN_ON_NON_PASSTHROUGH = auto()
+    """Issue a warning when an import is triggered in the sandbox that was not passed through."""
+    RAISE_ON_NON_PASSTHROUGH = auto()
+    """Raise an error when an import is triggered in the sandbox that was not passed through."""
+
+
+_sandbox_import_policy = threading.local()
+
+
 class unsafe:
     """Contains static methods that should not normally be called during
     workflow execution except in advanced cases.
@@ -1498,6 +1516,26 @@ def imports_passed_through() -> Iterator[None]:
         finally:
             _imports_passed_through.value = False
 
+    @staticmethod
+    def current_sandbox_import_policy() -> SandboxImportPolicy:
+        applied_policy = getattr(
+            _sandbox_import_policy, "value", SandboxImportPolicy.UNSET
+        )
+        return applied_policy
+
+    @staticmethod
+    @contextmanager
+    def sandbox_import_policy(policy: SandboxImportPolicy) -> Iterator[None]:
+        # TODO(amazzeo): the default behavior here seems inappropriate
+        original_policy = _sandbox_import_policy.value = getattr(
+            _sandbox_import_policy, "value", SandboxImportPolicy.UNSET
+        )
+        _sandbox_import_policy.value = policy
+        try:
+            yield None
+        finally:
+            _sandbox_import_policy.value = original_policy
+
 
 class LoggerAdapter(logging.LoggerAdapter):
     """Adapter that adds details to the log about the running workflow.

tests/worker/workflow_sandbox/test_runner.py

@@ -3,6 +3,7 @@
 import asyncio
 import dataclasses
 import functools
+import importlib
 import inspect
 import os
 import sys
@@ -11,19 +12,25 @@
 from dataclasses import dataclass
 from datetime import date, datetime, timedelta
 from enum import IntEnum
-from typing import Callable, Dict, List, Optional, Sequence, Set, Type
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Type
 
 import pytest
 
 from temporalio import activity, workflow
 from temporalio.client import Client, WorkflowFailureError, WorkflowHandle
 from temporalio.exceptions import ApplicationError
-from temporalio.worker import Worker
+from temporalio.worker import Worker, WorkflowInboundInterceptor
+from temporalio.worker._interceptor import (
+    ExecuteWorkflowInput,
+    Interceptor,
+    WorkflowInterceptorClassInput,
+)
 from temporalio.worker.workflow_sandbox import (
     RestrictedWorkflowAccessError,
     SandboxedWorkflowRunner,
     SandboxMatcher,
     SandboxRestrictions,
+    DisallowedUnintentionalPassthroughError,
 )
 from tests.helpers import assert_eq_eventually
 from tests.worker.workflow_sandbox.testmodules import stateful_module
@@ -483,3 +490,97 @@ def new_worker(
         activities=activities,
         workflow_runner=SandboxedWorkflowRunner(restrictions=restrictions),
     )
+
+
+class _TestWorkflowInboundInterceptor(WorkflowInboundInterceptor):
+    async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:
+        # import in the interceptor to show it will be captured
+        # applying this policy should squelch the "after initial workload" warning
+        with workflow.unsafe.sandbox_import_policy(
+            workflow.SandboxImportPolicy.WARN_ON_NON_PASSTHROUGH
+        ):
+            import opentelemetry
+
+        return await super().execute_workflow(input)
+
+
+class _TestInterceptor(Interceptor):
+    def workflow_interceptor_class(
+        self, input: WorkflowInterceptorClassInput
+    ) -> Type[_TestWorkflowInboundInterceptor]:
+        return _TestWorkflowInboundInterceptor
+
+
+@workflow.defn
+class LazyImportWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        try:
+            import opentelemetry.version
+        except DisallowedUnintentionalPassthroughError as err:
+            raise ApplicationError(
+                str(err), type="DisallowedUnintentionalPassthroughError"
+            ) from err
+
+
+async def test_workflow_sandbox_import_warnings(client: Client):
+    restrictions = dataclasses.replace(
+        SandboxRestrictions.default,
+        import_policy=SandboxRestrictions.import_policy_all_warnings,
+        # passthrough this test module to avoid a ton of noisy warnings
+        passthrough_modules=SandboxRestrictions.passthrough_modules_default
+        | {"tests.worker.workflow_sandbox.test_runner"},
+    )
+
+    async with Worker(
+        client,
+        task_queue=str(uuid.uuid4()),
+        workflows=[LazyImportWorkflow],
+        interceptors=[_TestInterceptor()],
+        workflow_runner=SandboxedWorkflowRunner(restrictions),
+    ) as worker:
+        with pytest.warns() as records:
+            await client.execute_workflow(
+                LazyImportWorkflow.run,
+                id=f"workflow-{uuid.uuid4()}",
+                task_queue=worker.task_queue,
+            )
+
+            expected_warnings = {
+                "Module opentelemetry was not intentionally passed through to the sandbox.",
+                "Module opentelemetry.version was imported after initial workflow load.",
+                "Module opentelemetry.version was not intentionally passed through to the sandbox.",
+            }
+            actual_warnings = {str(w.message) for w in records}
+
+            assert expected_warnings <= actual_warnings
+
+
+async def test_workflow_sandbox_import_errors(client: Client):
+    restrictions = dataclasses.replace(
+        SandboxRestrictions.default,
+        import_policy=SandboxRestrictions.import_policy_disallow_unintentional_passthrough,
+        # passthrough this test module to avoid a ton of noisy warnings
+        passthrough_modules=SandboxRestrictions.passthrough_modules_default
+        | {"tests.worker.workflow_sandbox.test_runner"},
+    )
+
+    async with Worker(
+        client,
+        task_queue=str(uuid.uuid4()),
+        workflows=[LazyImportWorkflow],
+        workflow_runner=SandboxedWorkflowRunner(restrictions),
+    ) as worker:
+        with pytest.raises(WorkflowFailureError) as err:
+            await client.execute_workflow(
+                LazyImportWorkflow.run,
+                id=f"workflow-{uuid.uuid4()}",
+                task_queue=worker.task_queue,
+            )
+
+        assert isinstance(err.value.cause, ApplicationError)
+        assert err.value.cause.type == "DisallowedUnintentionalPassthroughError"
+        assert (
+            "Module opentelemetry.version was not intentionally passed through to the sandbox."
+            == err.value.cause.message
+        )