Sandbox importing (#1187)

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

* create modules to lazily import rather than relying on a random dependency

* Do some renaming, apply linter.

* fix some naming and typos

* Rename error type and shorten pre-defined import policy name

* Cleanup new UnintentionalPassthroughError

* Apply some naming suggestions

* run formatter

* remove static policy definitions

* Add default option for SandboxRestrictions.import_notification_policy rather than requiring it.

Author: VegetarianOrc

temporalio/worker/workflow_sandbox/__init__.py

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

temporalio/worker/workflow_sandbox/_importer.py

@@ -42,6 +42,7 @@
     RestrictedWorkflowAccessError,
     RestrictionContext,
     SandboxRestrictions,
+    UnintentionalPassthroughError,
 )
 
 logger = logging.getLogger(__name__)
@@ -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_notification_policy_applied(
+                    temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT
+                )
+            ):
+                warnings.warn(
+                    f"Module {full_name} was imported after initial workflow load."
+                )
+
             # Make sure not an entirely invalid module
             self._assert_valid_module(full_name)
 
@@ -282,13 +294,36 @@ def module_configured_passthrough(self, name: str) -> bool:
                 break
         return True
 
+    def _is_import_notification_policy_applied(
+        self, policy: temporalio.workflow.SandboxImportNotificationPolicy
+    ) -> bool:
+        override_policy = (
+            temporalio.workflow.unsafe.current_import_notification_policy_override()
+        )
+        if override_policy:
+            return policy in override_policy
+
+        return policy in self.restrictions.import_notification_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_notification_policy_applied(
+                temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH
+            ):
+                raise UnintentionalPassthroughError(name)
+
+            if self._is_import_notification_policy_applied(
+                temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_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

@@ -42,6 +42,7 @@
 except ImportError:
     HAVE_PYDANTIC = False
 
+import temporalio.exceptions
 import temporalio.workflow
 
 logger = logging.getLogger(__name__)
@@ -82,6 +83,21 @@ def default_message(qualified_name: str) -> str:
         )
 
 
+class UnintentionalPassthroughError(temporalio.exceptions.TemporalError):
+    """Error that occurs when a workflow unintentionally passes an import to the sandbox when
+    the import notification policy includes :py:attr:`temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_NON_PASSTHROUGH`.
+
+    Attributes:
+        qualified_name: Fully qualified name of what was passed through to the sandbox.
+    """
+
+    def __init__(self, qualified_name: str) -> None:
+        """Create an unintentional passthrough error."""
+        super().__init__(
+            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 +126,13 @@ class methods (including __init__, etc). The check compares the against the
     fully qualified path to the item.
     """
 
+    import_notification_policy: temporalio.workflow.SandboxImportNotificationPolicy = (
+        temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT
+    )
+    """
+    The import notification policy to use when an import is triggered during workflow loading or execution. See :py:class:`temporalio.workflow.SandboxImportNotificationPolicy` for options.
+    """
+
     passthrough_all_modules: bool = False
     """
     Pass through all modules, do not sandbox any modules. This is the equivalent
@@ -170,6 +193,12 @@ def with_passthrough_all_modules(self) -> SandboxRestrictions:
         """
         return dataclasses.replace(self, passthrough_all_modules=True)
 
+    def with_import_notification_policy(
+        self, policy: temporalio.workflow.SandboxImportNotificationPolicy
+    ) -> SandboxRestrictions:
+        """Create a new restriction set with the given import notification policy as the :py:attr:`import_policy`."""
+        return dataclasses.replace(self, import_notification_policy=policy)
+
 
 # We intentionally use specific fields instead of generic "matcher" callbacks
 # for optimization reasons.
@@ -305,10 +334,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(
@@ -819,6 +850,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 (
@@ -1410,6 +1410,20 @@ async def wait_condition(
 _sandbox_unrestricted = threading.local()
 _in_sandbox = threading.local()
 _imports_passed_through = threading.local()
+_sandbox_import_notification_policy_override = threading.local()
+
+
+class SandboxImportNotificationPolicy(Flag):
+    """Defines the behavior taken when modules are imported into the sandbox after the workflow is initially loaded or unintentionally missing from the passthrough list."""
+
+    SILENT = auto()
+    """Allow imports that do not violate sandbox restrictions and no warnings are generated."""
+    WARN_ON_DYNAMIC_IMPORT = auto()
+    """Allows dynamic imports that do not violate sandbox restrictions but issues a warning when an import is triggered in the sandbox after initial workflow load."""
+    WARN_ON_UNINTENTIONAL_PASSTHROUGH = auto()
+    """Allows imports that do not violate sandbox restrictions but issues a warning when an import is triggered in the sandbox that was unintentionally passed through."""
+    RAISE_ON_UNINTENTIONAL_PASSTHROUGH = auto()
+    """Raise an error when an import is triggered in the sandbox that was unintentionally passed through."""
 
 
 class unsafe:
@@ -1498,6 +1512,35 @@ def imports_passed_through() -> Iterator[None]:
         finally:
             _imports_passed_through.value = False
 
+    @staticmethod
+    def current_import_notification_policy_override() -> (
+        Optional[SandboxImportNotificationPolicy]
+    ):
+        """Gets the current import notification policy override if one is set."""
+        applied_policy = getattr(
+            _sandbox_import_notification_policy_override,
+            "value",
+            None,
+        )
+        return applied_policy
+
+    @staticmethod
+    @contextmanager
+    def sandbox_import_notification_policy(
+        policy: SandboxImportNotificationPolicy,
+    ) -> Iterator[None]:
+        """Context manager to apply the given import notification policy."""
+        original_policy = _sandbox_import_notification_policy_override.value = getattr(
+            _sandbox_import_notification_policy_override,
+            "value",
+            None,
+        )
+        _sandbox_import_notification_policy_override.value = policy
+        try:
+            yield None
+        finally:
+            _sandbox_import_notification_policy_override.value = original_policy
+
 
 class LoggerAdapter(logging.LoggerAdapter):
     """Adapter that adds details to the log about the running workflow.