Python: fix(core): restrict unpickler module-prefix allowlist to types only

python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py

@@ -20,9 +20,9 @@
 format uses Python's ``pickle`` module which can execute arbitrary code during
 deserialization.  The ``RestrictedUnpickler`` provides a defense-in-depth
 allowlist that limits instantiable classes, but it is **not** a security
-boundary — certain allowlisted builtins (e.g. ``getattr``) are required for
-legitimate object reconstruction (enums, named tuples) and cannot be removed
-without breaking compatibility.
+boundary — certain reconstruction helpers are required for legitimate object
+reconstruction and application-defined types can provide custom pickle
+behavior.
 
 Developers **must** ensure that:
 
@@ -50,6 +50,7 @@
 import io
 import logging
 import pickle  # nosec  # ruff:ignore[suspicious-pickle-import]
+from enum import EnumMeta
 from typing import Any, cast
 
 from ..exceptions import WorkflowCheckpointException
@@ -81,9 +82,9 @@
     "agent_framework._workflows._checkpoint_encoding:encode_checkpoint_value",
 })
 
-# Built-in types considered safe for checkpoint deserialization.
+# Built-in globals considered safe for checkpoint deserialization.
 # Each entry is a ``module:qualname`` string matching the format produced by
-# :func:`_type_to_key`.  These are the classes for which pickle's
+# :func:`_type_to_key`.  These are the globals for which pickle's
 # ``find_class`` will be called when unpickling common Python value types.
 _BUILTIN_ALLOWED_TYPE_KEYS: frozenset[str] = frozenset({
     # builtins
@@ -103,7 +104,8 @@
     "builtins:dict",
     "builtins:tuple",
     "builtins:type",
-    # getattr is used by pickle to reconstruct enum members
+    # getattr is used by pickle to reconstruct nested types. find_class
+    # substitutes a restricted implementation during unpickling.
     "builtins:getattr",
     # copyreg helpers used by pickle for object reconstruction
     "copyreg:_reconstructor",
@@ -123,6 +125,12 @@
     "collections:deque",
 })
 
+_GETATTR_GLOBAL_KEYS: frozenset[str] = frozenset({
+    "builtins:getattr",
+    # Protocol 2 pickles use Python 2's module name for the same global.
+    "__builtin__:getattr",
+})
+
 
 class _RestrictedUnpickler(pickle.Unpickler):  # ruff:ignore[suspicious-pickle-usage]
     """Unpickler that restricts which classes may be instantiated.
@@ -137,15 +145,58 @@ def __init__(self, data: bytes, allowed_types: frozenset[str]) -> None:
         super().__init__(io.BytesIO(data))
         self._allowed_types = allowed_types
 
-    def find_class(self, module: str, name: str) -> type:
+    def _is_allowed_type(self, resolved: type) -> bool:
+        type_key = _type_to_key(resolved)
+        return (
+            type_key in _BUILTIN_ALLOWED_TYPE_KEYS
+            or type_key in self._allowed_types
+            or resolved.__module__.startswith(_FRAMEWORK_MODULE_PREFIX)
+            or resolved.__module__.startswith(_OPENAI_MODULE_PREFIX)
+        )
+
+    def _restricted_getattr(self, obj: Any, name: str) -> Any:
+        """Resolve an allowlisted nested type or named enum member."""
+        if not isinstance(obj, type) or not isinstance(name, str):
+            raise pickle.UnpicklingError("Checkpoint deserialization blocked for unsafe attribute traversal.")
+
+        resolved = getattr(obj, name)
+        if isinstance(resolved, type):
+            if not self._is_allowed_type(resolved):
+                type_key = _type_to_key(resolved)
+                raise pickle.UnpicklingError(
+                    f"Checkpoint deserialization blocked for nested type '{type_key}'. "
+                    f"Include the nested type in 'allowed_types' before loading the checkpoint."
+                )
+            return resolved
+
+        # enum.pickle_by_enum_name reconstructs a member with getattr(EnumClass, member_name).
+        if isinstance(obj, EnumMeta) and self._is_allowed_type(obj):
+            members = obj.__members__
+            if name in members and members[name] is resolved:
+                return resolved
+
+        raise pickle.UnpicklingError(
+            f"Checkpoint deserialization blocked for non-type attribute '{obj.__module__}:{obj.__qualname__}.{name}'."
+        )
+
+    def find_class(self, module: str, name: str) -> Any:
         type_key = f"{module}:{name}"
 
         if type_key in _BLOCKED_FRAMEWORK_GLOBAL_KEYS:
             raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for type '{type_key}'.")
 
-        if type_key in _BUILTIN_ALLOWED_TYPE_KEYS or type_key in self._allowed_types:
+        if type_key in _GETATTR_GLOBAL_KEYS:
+            return self._restricted_getattr
+
+        if type_key in _BUILTIN_ALLOWED_TYPE_KEYS:
             return super().find_class(module, name)  # nosec
 
+        if type_key in self._allowed_types:
+            resolved = super().find_class(module, name)  # nosec
+            if isinstance(resolved, type):
+                return resolved
+            raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for non-type global '{type_key}'.")
+
         if module.startswith(_FRAMEWORK_MODULE_PREFIX) or module.startswith(_OPENAI_MODULE_PREFIX):
             # Pickle dotted names traverse attributes on an allowed module; keep the prefix allowlist to concrete
             # top-level classes rather than helper callables reachable through module attributes.
@@ -154,6 +205,7 @@ def find_class(self, module: str, name: str) -> type:
             resolved = super().find_class(module, name)  # nosec
             if isinstance(resolved, type):
                 return resolved
+            raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for non-type global '{type_key}'.")
 
         raise pickle.UnpicklingError(
             f"Checkpoint deserialization blocked for type '{type_key}'. "

python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py

@@ -12,6 +12,7 @@
 """
 
 import base64
+import enum
 import os
 import pickle
 import tempfile
@@ -49,6 +50,39 @@ def __reduce__(self) -> tuple[Any, tuple[str]]:
         return (_base64_to_unpickle, (self.nested_payload,))
 
 
+class _NestedTypeContainer:
+    @dataclass
+    class NestedType:
+        value: int
+
+
+class _NestedSetStateContainer:
+    class Evil:
+        setstate_called = False
+
+        def __init__(self) -> None:
+            self.value = 42
+
+        def __setstate__(self, state: dict[str, object]) -> None:
+            type(self).setstate_called = True
+            self.__dict__.update(state)
+
+
+class _NamedPickleColor(enum.Enum):
+    RED = 1
+    BLUE = 2
+    CRIMSON = 1
+
+    def __reduce_ex__(self, proto: object) -> tuple[Any, tuple[type[enum.Enum], str]]:
+        del proto
+        return (getattr, (type(self), self.name))
+
+
+class _EnumNonMemberAttributePayload:
+    def __reduce__(self):
+        return (getattr, (_NamedPickleColor, "__members__"))
+
+
 def test_restricted_decode_blocks_arbitrary_callable():
     """Restricted decoding blocks arbitrary module-level callables."""
     pickled = pickle.dumps(os.getpid, protocol=pickle.HIGHEST_PROTOCOL)
@@ -300,6 +334,136 @@ def test_restricted_unpickler_raises_pickle_error():
         unpickler.load()
 
 
+def test_restricted_decode_rejects_non_type_global_in_allowed_types():
+    """Explicit allowed_types entries must resolve to types."""
+    from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler
+
+    unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset({"os:getpid"}))
+    with pytest.raises(pickle.UnpicklingError, match="non-type global"):
+        unpickler.find_class("os", "getpid")
+
+
+def test_restricted_decode_rejects_non_type_global_under_prefix():
+    """Allowed package prefixes must not expose arbitrary module globals."""
+    from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler
+
+    unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset())
+    with pytest.raises(pickle.UnpicklingError, match="non-type global"):
+        unpickler.find_class(
+            "agent_framework._workflows._checkpoint_encoding",
+            "_value_type_to_key",
+        )
+
+
+@pytest.mark.parametrize("protocol", [2, 3])
+def test_restricted_decode_allows_explicitly_listed_nested_type(protocol: int):
+    """Protocol 2/3 nested types load when both the container and nested type are allowed."""
+    original = _NestedTypeContainer.NestedType(42)
+    checkpoint_value = {
+        _PICKLE_MARKER: base64.b64encode(pickle.dumps(original, protocol=protocol)).decode("ascii"),
+        _TYPE_MARKER: f"{type(original).__module__}:{type(original).__qualname__}",
+    }
+    allowed_types = frozenset({
+        f"{_NestedTypeContainer.__module__}:{_NestedTypeContainer.__qualname__}",
+        f"{_NestedTypeContainer.NestedType.__module__}:{_NestedTypeContainer.NestedType.__qualname__}",
+    })
+
+    decoded = decode_checkpoint_value(checkpoint_value, allowed_types=allowed_types)
+
+    assert isinstance(decoded, _NestedTypeContainer.NestedType)
+    assert decoded.value == 42
+
+
+@pytest.mark.parametrize("protocol", [2, 3])
+def test_restricted_decode_blocks_unlisted_nested_type_before_setstate(protocol: int):
+    """Allowing only a container must not allow its nested types or run their state hooks."""
+    original = _NestedSetStateContainer.Evil()
+    checkpoint_value = {
+        _PICKLE_MARKER: base64.b64encode(pickle.dumps(original, protocol=protocol)).decode("ascii"),
+        _TYPE_MARKER: f"{type(original).__module__}:{type(original).__qualname__}",
+    }
+    outer_type_key = f"{_NestedSetStateContainer.__module__}:{_NestedSetStateContainer.__qualname__}"
+    _NestedSetStateContainer.Evil.setstate_called = False
+
+    with pytest.raises(WorkflowCheckpointException, match="nested type"):
+        decode_checkpoint_value(checkpoint_value, allowed_types=frozenset({outer_type_key}))
+
+    assert not _NestedSetStateContainer.Evil.setstate_called
+
+
+@pytest.mark.parametrize("protocol", [2, 3, 4, 5])
+def test_restricted_decode_allows_allowlisted_enum_member_by_name(protocol: int):
+    """Named enum members load only through their allowlisted enum class."""
+    original = _NamedPickleColor.RED
+    checkpoint_value = {
+        _PICKLE_MARKER: base64.b64encode(pickle.dumps(original, protocol=protocol)).decode("ascii"),
+        _TYPE_MARKER: f"{type(original).__module__}:{type(original).__qualname__}",
+    }
+    enum_type_key = f"{_NamedPickleColor.__module__}:{_NamedPickleColor.__qualname__}"
+
+    decoded = decode_checkpoint_value(checkpoint_value, allowed_types=frozenset({enum_type_key}))
+
+    assert decoded is _NamedPickleColor.RED
+
+
+def test_restricted_encoder_round_trips_allowlisted_enum_member_by_name():
+    """The current checkpoint encoder's protocol preserves named enum members."""
+    enum_type_key = f"{_NamedPickleColor.__module__}:{_NamedPickleColor.__qualname__}"
+
+    encoded = encode_checkpoint_value(_NamedPickleColor.RED)
+    decoded = decode_checkpoint_value(encoded, allowed_types=frozenset({enum_type_key}))
+
+    assert decoded is _NamedPickleColor.RED
+
+
+def test_restricted_decode_blocks_non_member_attribute_on_allowlisted_enum():
+    """Allowing an enum class does not expose its other attributes through getattr."""
+    payload = pickle.dumps(_EnumNonMemberAttributePayload(), protocol=pickle.HIGHEST_PROTOCOL)
+    checkpoint_value = {
+        _PICKLE_MARKER: base64.b64encode(payload).decode("ascii"),
+        _TYPE_MARKER: "builtins:dict",
+    }
+    enum_type_key = f"{_NamedPickleColor.__module__}:{_NamedPickleColor.__qualname__}"
+
+    with pytest.raises(WorkflowCheckpointException, match="non-type attribute"):
+        decode_checkpoint_value(checkpoint_value, allowed_types=frozenset({enum_type_key}))
+
+
+def test_restricted_decode_blocks_member_of_unlisted_enum():
+    """Enum member reconstruction still requires the enum class in allowed_types."""
+    original = _NamedPickleColor.RED
+    checkpoint_value = {
+        _PICKLE_MARKER: base64.b64encode(pickle.dumps(original, protocol=pickle.HIGHEST_PROTOCOL)).decode("ascii"),
+        _TYPE_MARKER: f"{type(original).__module__}:{type(original).__qualname__}",
+    }
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        decode_checkpoint_value(checkpoint_value, allowed_types=frozenset())
+
+
+def test_restricted_decode_blocks_getattr_globals_pickle_loads_chain():
+    """Restricted decoding blocks attribute traversal to an inner unrestricted pickle load."""
+    inner_pickle = b"cbuiltins\neval\n(V40 + 2\ntR."
+    escaped_inner_pickle = inner_pickle.decode("ascii").replace("\\", "\\u005c").replace("\n", "\\u000a")
+    # The outer pickle memoizes getattr and walks __init__.__globals__["pickle"].loads.
+    payload = (
+        b"cbuiltins\ngetattr\np0\n0"
+        b"g0\n(cagent_framework._workflows._checkpoint_encoding\n_RestrictedUnpickler\nV__init__\ntRp1\n0"
+        b"g0\n(g1\nV__globals__\ntRp2\n0"
+        b"g0\n(cbuiltins\ndict\nV__getitem__\ntRp3\n0"
+        b"g3\n(g2\nVpickle\ntRp4\n0"
+        b"g0\n(g4\nVloads\ntRp5\n0"
+        b"g5\n(cbuiltins\nbytearray\n(V" + escaped_inner_pickle.encode("ascii") + b"\nVascii\ntRtR."
+    )
+    checkpoint_value = {
+        _PICKLE_MARKER: base64.b64encode(payload).decode("ascii"),
+        _TYPE_MARKER: "builtins:int",
+    }
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        decode_checkpoint_value(checkpoint_value, allowed_types=frozenset())
+
+
 def test_restricted_decode_allows_openai_types():
     """OpenAI SDK types are always allowed during restricted deserialization."""
     from openai.types.chat.chat_completion import ChatCompletion, Choice