From 1dab04ae79cd62d632f03100be6694165d637b5d Mon Sep 17 00:00:00 2001 From: White-Mouse <15983334+White-Mouse@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:17:36 +0800 Subject: [PATCH 1/2] fix(core): harden restricted pickle attribute resolution --- .../_workflows/_checkpoint_encoding.py | 40 ++++++++++--- .../test_checkpoint_unrestricted_pickle.py | 59 +++++++++++++++++++ 2 files changed, 91 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index d5c1703a46a..9a684df3754 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/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: @@ -81,9 +81,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 +103,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", @@ -124,6 +125,19 @@ }) +def _restricted_getattr(obj: Any, name: str) -> type: + """Resolve type-valued attributes needed for legitimate pickle reconstruction.""" + 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 not isinstance(resolved, type): + raise pickle.UnpicklingError( + f"Checkpoint deserialization blocked for non-type attribute '{obj.__module__}:{obj.__qualname__}.{name}'." + ) + return resolved + + class _RestrictedUnpickler(pickle.Unpickler): # noqa: S301 """Unpickler that restricts which classes may be instantiated. @@ -137,15 +151,24 @@ 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 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 == "builtins:getattr": + return _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 +177,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}'. " diff --git a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py index 7ebd6f13a05..8cc7b8c5f95 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -49,6 +49,11 @@ def __reduce__(self) -> tuple[Any, tuple[str]]: return (_base64_to_unpickle, (self.nested_payload,)) +class _NestedTypeContainer: + class NestedType: + pass + + 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 +305,60 @@ 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", + "encode_checkpoint_value", + ) + + +def test_restricted_getattr_allows_nested_type_resolution(): + """The restricted getattr replacement preserves nested-type reconstruction.""" + from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler + + unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset()) + restricted_getattr = unpickler.find_class("builtins", "getattr") + + assert restricted_getattr(_NestedTypeContainer, "NestedType") is _NestedTypeContainer.NestedType + + +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="non-type attribute"): + 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 From 9dd038e1444fb74939665a289f04ff29ac5f13cd Mon Sep 17 00:00:00 2001 From: White-Mouse <15983334+White-Mouse@users.noreply.github.com> Date: Tue, 14 Jul 2026 20:42:31 +0800 Subject: [PATCH 2/2] fix(core): validate nested pickle types against allowlist --- .../_workflows/_checkpoint_encoding.py | 56 ++++++-- .../test_checkpoint_unrestricted_pickle.py | 121 ++++++++++++++++-- 2 files changed, 154 insertions(+), 23 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 9a684df3754..4a741ec4ae0 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -50,6 +50,7 @@ import io import logging import pickle # nosec # noqa: S403 +from enum import EnumMeta from typing import Any, cast from ..exceptions import WorkflowCheckpointException @@ -124,18 +125,11 @@ "collections:deque", }) - -def _restricted_getattr(obj: Any, name: str) -> type: - """Resolve type-valued attributes needed for legitimate pickle reconstruction.""" - 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 not isinstance(resolved, type): - raise pickle.UnpicklingError( - f"Checkpoint deserialization blocked for non-type attribute '{obj.__module__}:{obj.__qualname__}.{name}'." - ) - return resolved +_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): # noqa: S301 @@ -151,14 +145,48 @@ def __init__(self, data: bytes, allowed_types: frozenset[str]) -> None: super().__init__(io.BytesIO(data)) self._allowed_types = allowed_types + 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 == "builtins:getattr": - return _restricted_getattr + 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 diff --git a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py index 8cc7b8c5f95..05e42317fb4 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -12,6 +12,7 @@ """ import base64 +import enum import os import pickle import tempfile @@ -50,8 +51,34 @@ def __reduce__(self) -> tuple[Any, tuple[str]]: class _NestedTypeContainer: + @dataclass class NestedType: - pass + 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 + + __reduce_ex__ = enum.pickle_by_enum_name # type: ignore[attr-defined] + + +class _EnumNonMemberAttributePayload: + def __reduce__(self): + return (getattr, (_NamedPickleColor, "__members__")) def test_restricted_decode_blocks_arbitrary_callable(): @@ -322,18 +349,94 @@ def test_restricted_decode_rejects_non_type_global_under_prefix(): with pytest.raises(pickle.UnpicklingError, match="non-type global"): unpickler.find_class( "agent_framework._workflows._checkpoint_encoding", - "encode_checkpoint_value", + "_value_type_to_key", ) -def test_restricted_getattr_allows_nested_type_resolution(): - """The restricted getattr replacement preserves nested-type reconstruction.""" - from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler +@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__}", + }) - unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset()) - restricted_getattr = unpickler.find_class("builtins", "getattr") + 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 restricted_getattr(_NestedTypeContainer, "NestedType") is _NestedTypeContainer.NestedType + 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(): @@ -355,7 +458,7 @@ def test_restricted_decode_blocks_getattr_globals_pickle_loads_chain(): _TYPE_MARKER: "builtins:int", } - with pytest.raises(WorkflowCheckpointException, match="non-type attribute"): + with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"): decode_checkpoint_value(checkpoint_value, allowed_types=frozenset())