From 4ec795beb9aa283ced78a41bf6239fbd0422f114 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 29 Oct 2026 12:43:00 +0100 Subject: [PATCH] Python: reject metaclasses when restoring checkpoint data _RestrictedUnpickler allowlisted builtins:type, which pickle's REDUCE opcode can call as (name, bases, namespace) to construct a brand-new class whose module/qualname are chosen by the payload. Several related routes could reach the same class factory even without the direct allowlist entry: restricted getattr(dict, "__class__"), and module attributes that re-export type under another name within an otherwise auto-allowed agent_framework.* prefix (for example "from builtins import type as builtin_type"). Add an explicit check, applied everywhere a class is resolved, that the unpickler never hands back a metaclass (a subclass of type), regardless of which allowlisted name or prefix was used to reach it. This closes all of the above routes with one invariant instead of patching each individually. builtins:type itself is resolved to a call restricted to its introspective one-argument form over a non-class value (type(x)), rather than removed from the allowlist outright: that form only returns an already-constructed value's own class and cannot construct a new one, so checkpoints that hold type(None)/etc. as data remain restorable while the three-argument construction form stays blocked. Calling the restricted form on a class argument is also rejected: every ordinary class is itself an instance of type, so type(dict) (or type(list), etc.) would otherwise return the bare type object as a plain value; pickle's REDUCE opcode calls whatever object is on top of its stack without re-resolving it through find_class, so that returned value could then be memoized and reused as the callable of a second, unchecked REDUCE to construct an arbitrary class anyway. Rejecting a class argument closes that route, because type(x) can only return a metaclass when x is itself a class. A related class-factory path exists through the Enum functional API: calling an Enum-family class with more than one positional argument (EnumMeta.__call__) constructs a new Enum subclass the same way. Python refuses to subclass an Enum that already has members on every supported version, so only member-less Enum-family classes (the base Enum/IntEnum/Flag/IntFlag/StrEnum classes themselves, reachable through the same auto-allowed module prefix) are rejected; populated application and framework enums are unaffected. Ordinary classes (dataclasses, populated enums, dict/list subclasses, named tuples, nested classes, etc.) continue to round-trip exactly as before; only the class object type(None) and member-less Enum classes stop being restorable by default, and FileCheckpointStorage already rejects anything it could not restore at save time. Adds focused regression tests for the direct global, both known in-framework type aliases, the enum-metaclass re-export, the getattr route, the member-less Enum re-export, explicitly-configured metaclass and type entries, the type(None) single-argument compatibility round trip, the two-REDUCE metaclass-recovery chain described above, and an end-to-end FileCheckpointStorage test that edits a saved checkpoint file on disk. The reusable pickle-payload builders used by these tests live in a dedicated non-test helper module. --- .../_workflows/_checkpoint_encoding.py | 105 ++++++++- .../workflow/checkpoint_pickle_builders.py | 89 ++++++++ .../test_checkpoint_unrestricted_pickle.py | 212 ++++++++++++++++++ 3 files changed, 399 insertions(+), 7 deletions(-) create mode 100644 python/packages/core/tests/workflow/checkpoint_pickle_builders.py diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 209d1ffea9..89713db8bd 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -129,7 +129,10 @@ def register_checkpoint_type(cls: type[Any]) -> None: "builtins:list", "builtins:dict", "builtins:tuple", - "builtins:type", + # Deliberately absent: "builtins:type". It is resolved separately (see + # _RESTRICTED_TYPE_GLOBAL_KEYS below) to a call restricted to its introspective one-argument + # form, because the three-argument form, type(name, bases, namespace), constructs a new class + # whose __module__/__qualname__ are chosen by the pickle data itself. # getattr is used by pickle to reconstruct nested types. find_class # substitutes a restricted implementation during unpickling. "builtins:getattr", @@ -158,6 +161,68 @@ def register_checkpoint_type(cls: type[Any]) -> None: "__builtin__:getattr", }) +_RESTRICTED_TYPE_GLOBAL_KEYS: frozenset[str] = frozenset({ + "builtins:type", + # Protocol 2 pickles use Python 2's module name for the same global. + "__builtin__:type", +}) + + +def _restricted_type(*args: Any, **kwargs: Any) -> Any: + """Restrict ``type`` to its introspective one-argument form over a non-class value. + + The three-argument form, ``type(name, bases, namespace)``, constructs a new class whose + identity is chosen entirely by its arguments -- the same hazard the metaclass and + unpopulated-enum checks below exist to prevent. Calling the one-argument form on a *value* + only returns that value's existing class and cannot construct a new one, so restricting to + that shape keeps checkpoints that hold ``type(None)``/etc. as data restorable. Calling it on a + *class* is a second route to the same hazard: every ordinary class (``dict``, ``list``, + ``int``, and any class without a custom metaclass) is itself an instance of ``type``, so + ``type(dict)`` returns the bare ``type`` object -- the pickle stream could then reuse that + returned value as the callable of a *second*, unchecked ``REDUCE`` (``find_class`` is never + consulted for values already on the unpickling stack) to construct an arbitrary class after + all. Rejecting a class argument closes that route; ``type(x)`` can only return a metaclass + when ``x`` is itself a class, because only classes are instances of a metaclass. + """ + if len(args) != 1 or kwargs: + raise pickle.UnpicklingError( + "Checkpoint deserialization blocked for 'builtins:type' called to construct a new " + "class; only the single-argument introspective form is permitted." + ) + value: object = args[0] + if isinstance(value, type): + raise pickle.UnpicklingError( + "Checkpoint deserialization blocked for 'builtins:type' called on a class; this would " + "return a metaclass that could be reused to construct an arbitrary class." + ) + return type(value) + + +def _is_metaclass(resolved: type) -> bool: + """Return True if instantiating ``resolved`` would itself construct a new class. + + A metaclass (``type`` or any subclass of it) is a class factory: calling it with + ``(qualname, bases, namespace)`` creates a brand-new class whose identity is chosen by + the caller. Checkpoint deserialization must never hand one back, regardless of which + allowlisted name or prefix was used to reach it, or a pickle could construct a class + that satisfies any downstream "is this an allowed type" check. + """ + return issubclass(resolved, type) + + +def _is_unpopulated_enum(resolved: type) -> bool: + """Return True if ``resolved`` is an Enum-family class with no members. + + ``EnumMeta.__call__`` dispatches a call with more than one positional argument to the Enum + functional API, which constructs a brand-new Enum subclass -- the same class-construction + hazard as calling ``type`` directly. Python refuses to subclass (functional-API or otherwise) + an Enum that already has members, on every supported Python version, so a populated + application or framework enum cannot be used this way; only a member-less class (the base + ``Enum``/``IntEnum``/``Flag``/``IntFlag``/``StrEnum`` classes themselves, or an unusual custom + subclass with no members) is a reachable class factory and needs to be rejected here. + """ + return isinstance(resolved, EnumMeta) and len(resolved) == 0 + class _RestrictedUnpickler(pickle.Unpickler): # ruff:ignore[suspicious-pickle-usage] """Unpickler that restricts which classes may be instantiated. @@ -173,6 +238,8 @@ def __init__(self, data: bytes, allowed_types: frozenset[str]) -> None: self._allowed_types = allowed_types def _is_allowed_type(self, resolved: type) -> bool: + if _is_metaclass(resolved) or _is_unpopulated_enum(resolved): + return False type_key = _type_to_key(resolved) return ( type_key in _BUILTIN_ALLOWED_TYPE_KEYS @@ -215,14 +282,26 @@ def find_class(self, module: str, name: str) -> Any: if type_key in _GETATTR_GLOBAL_KEYS: return self._restricted_getattr + if type_key in _RESTRICTED_TYPE_GLOBAL_KEYS: + return _restricted_type + 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 not isinstance(resolved, type): + raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for non-type global '{type_key}'.") + if _is_metaclass(resolved): + # Explicitly listing a metaclass's own key (e.g. a custom allowed_checkpoint_types entry) + # would otherwise still hand back a class factory; there is no legitimate reconstructible + # application type that is itself a metaclass. + raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for metaclass '{type_key}'.") + if _is_unpopulated_enum(resolved): + raise pickle.UnpicklingError( + f"Checkpoint deserialization blocked for unpopulated Enum class '{type_key}'." + ) + return resolved if module.startswith(_FRAMEWORK_MODULE_PREFIXES) or module.startswith(_OPENAI_MODULE_PREFIX): # Pickle dotted names traverse attributes on an allowed module; keep the prefix allowlist to concrete @@ -230,9 +309,21 @@ def find_class(self, module: str, name: str) -> Any: if "." in name: raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for type '{type_key}'.") 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 not isinstance(resolved, type): + raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for non-type global '{type_key}'.") + if _is_metaclass(resolved): + # A module can re-export a metaclass under an allowed prefix (for example `type` imported + # as `from builtins import type as builtin_type`); judge the class itself, not just where + # the name was found. + raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for metaclass '{type_key}'.") + if _is_unpopulated_enum(resolved): + # Likewise, a module can re-export a member-less Enum-family base class (for example + # `from enum import Enum`); its functional API constructs a brand-new class the same + # way calling a metaclass does. + raise pickle.UnpicklingError( + f"Checkpoint deserialization blocked for unpopulated Enum class '{type_key}'." + ) + return resolved raise pickle.UnpicklingError( f"Checkpoint deserialization blocked for type '{type_key}'. " diff --git a/python/packages/core/tests/workflow/checkpoint_pickle_builders.py b/python/packages/core/tests/workflow/checkpoint_pickle_builders.py new file mode 100644 index 0000000000..d90bec1a22 --- /dev/null +++ b/python/packages/core/tests/workflow/checkpoint_pickle_builders.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Reusable pickle payload builders for checkpoint-deserialization tests.""" + +from __future__ import annotations + + +def _short_unicode(text: str) -> bytes: + data = text.encode("utf-8") + return b"\x8c" + bytes([len(data)]) + data + + +def _stack_global(module: str, name: str) -> bytes: + return _short_unicode(module) + _short_unicode(name) + b"\x93" + + +def dynamic_dict_subclass_pickle( + class_global_module: str, class_global_name: str, dynamic_qualname: str, items: dict[str, str] +) -> bytes: + """Hand-assemble a protocol-4 pickle equivalent to: + ``(dynamic_qualname, (dict,), {"__module__": ..., "__qualname__": ...})(items)``. + + Exercises the primitive the checkpoint-restore allowlist must reject: calling an allowlisted + class factory (``type``, or a module attribute that resolves to it) via ``REDUCE`` to construct + a new ``dict`` subclass whose ``__module__``/``__qualname__`` are chosen by the payload, then + instantiating it with caller-supplied data. + """ + class_dict = ( + b"}(" + + _short_unicode("__module__") + + _short_unicode(class_global_module) + + _short_unicode("__qualname__") + + _short_unicode(dynamic_qualname) + + b"u" + ) + payload_items = b"}(" + b"".join(_short_unicode(k) + _short_unicode(v) for k, v in items.items()) + b"u" + return ( + b"\x80\x04" + + _stack_global(class_global_module, class_global_name) + + _short_unicode(dynamic_qualname) + + _stack_global("builtins", "dict") + + b"\x85" # TUPLE1 -> (dict,) + + class_dict + + b"\x87R" # TUPLE3, REDUCE -> constructed class + + payload_items + + b"\x85R." # TUPLE1, REDUCE -> constructed instance, STOP + ) + + +def metaclass_via_builtin_type_pickle( + already_allowed_class_global: tuple[str, str], dynamic_qualname: str, items: dict[str, str] +) -> bytes: + """Hand-assemble a protocol-4 pickle that recovers the bare ``type`` metaclass from a restricted + single-argument ``type(x)`` call, then reuses that recovered value -- not a ``find_class`` + lookup -- as the callable of a second ``REDUCE`` to construct an arbitrary class. + + Every ordinary class (``dict``, ``list``, any class without a custom metaclass) is itself an + instance of ``type``, so ``type()`` would legitimately return the + bare ``type`` object under a restriction that only checks argument count/kind. Pickle's + ``REDUCE`` opcode calls whatever object is on top of the stack; it does not re-resolve that + object through ``find_class``, so a value computed by one ``REDUCE`` can be memoized + (``PUT``/``GET``) and reused as the callable of a *later* ``REDUCE`` with different arguments. + This checks that a single-argument ``type(x)`` restriction also rejects ``x`` being a class, + which is the only way ``type(x)`` can itself return a metaclass. + """ + class_dict = ( + b"}(" + + _short_unicode("__module__") + + _short_unicode("builtins") + + _short_unicode("__qualname__") + + _short_unicode(dynamic_qualname) + + b"u" + ) + payload_items = b"}(" + b"".join(_short_unicode(k) + _short_unicode(v) for k, v in items.items()) + b"u" + return ( + b"\x80\x04" + + _stack_global("builtins", "type") + + _stack_global(*already_allowed_class_global) + + b"\x85R" # TUPLE1, REDUCE -> type(already_allowed_class_global) -- recovers bare `type` + + b"q\x00" # BINPUT 0 -- memoize the recovered `type` object + + b"h\x00" # BINGET 0 -- reuse it as the callable of the next REDUCE, without going through find_class + + _short_unicode(dynamic_qualname) + + _stack_global("builtins", "dict") + + b"\x85" # TUPLE1 -> (dict,) + + class_dict + + b"\x87R" # TUPLE3, REDUCE -> type(qualname, (dict,), {...}) -- constructed class + + payload_items + + b"\x85R." # TUPLE1, REDUCE -> constructed instance, STOP + ) 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 9fdb0fe2e4..66dc6e5850 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -14,6 +14,7 @@ import asyncio import base64 import enum +import json import os import pickle import tempfile @@ -34,6 +35,7 @@ decode_checkpoint_value, encode_checkpoint_value, ) +from tests.workflow.checkpoint_pickle_builders import dynamic_dict_subclass_pickle, metaclass_via_builtin_type_pickle class MaliciousPayload: @@ -81,6 +83,13 @@ def __reduce_ex__(self, proto: object) -> tuple[Any, tuple[type[enum.Enum], str] return (getattr, (type(self), self.name)) +class _PopulatedColor(enum.Enum): + """A populated enum used to prove member-less rejection does not affect ordinary round trips.""" + + RED = 1 + BLUE = 2 + + class _EnumNonMemberAttributePayload: def __reduce__(self): return (getattr, (_NamedPickleColor, "__members__")) @@ -420,6 +429,209 @@ def test_restricted_decode_rejects_non_type_global_under_prefix(): ) +def test_restricted_decode_blocks_builtins_type_class_construction(): + """``type(qualname, (dict,), {...})(items)`` constructs a new ``dict`` subclass whose + ``__module__``/``__qualname__`` are chosen by the payload; the restricted unpickler must block + this three-argument form before the new class is ever instantiated, even though ``dict`` alone + remains allowed. + """ + payload = dynamic_dict_subclass_pickle("builtins", "type", "DynamicConfig", {"token": "value"}) + checkpoint_value = { + _PICKLE_MARKER: base64.b64encode(payload).decode("ascii"), + _TYPE_MARKER: "builtins:DynamicConfig", + } + + with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"): + decode_checkpoint_value(checkpoint_value, allowed_types=frozenset()) + + +def test_restricted_decode_allows_type_single_argument_form(): + """``type(x)`` (the introspective one-argument form) remains restorable: it only returns the + type of an already-constructed value and cannot construct a new class, so a checkpoint saved by + an earlier release that holds a bare ``NoneType`` value (``type(None)``, which pickles via this + same single-argument ``REDUCE`` call) must keep loading. + """ + encoded = encode_checkpoint_value(type(None)) + decoded = decode_checkpoint_value(encoded, allowed_types=frozenset()) + + assert decoded is type(None) + + +def test_restricted_decode_blocks_explicit_three_argument_type_call(): + """``find_class`` resolves ``builtins:type`` to a restricted call that still rejects the + dangerous three-argument (or any keyword-argument) construction form directly. + """ + from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler + + unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset()) + restricted_type = unpickler.find_class("builtins", "type") + assert restricted_type(1) is int + with pytest.raises(pickle.UnpicklingError, match="deserialization blocked"): + restricted_type("DynamicConfig", (dict,), {}) + + +def test_restricted_decode_type_remains_restricted_even_when_explicitly_allowed(): + """Explicitly listing ``builtins:type`` in ``allowed_types`` does not relax it to the dangerous + construction form; it always resolves to the same restricted single-argument call. + """ + from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler + + unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset({"builtins:type"})) + restricted_type = unpickler.find_class("builtins", "type") + assert restricted_type(1) is int + with pytest.raises(pickle.UnpicklingError, match="deserialization blocked"): + restricted_type("DynamicConfig", (dict,), {}) + + +@pytest.mark.parametrize("already_allowed_class_global", [("builtins", "dict"), ("builtins", "list")]) +def test_restricted_decode_blocks_type_called_on_a_class(already_allowed_class_global: tuple[str, str]): + """``type(x)`` must reject ``x`` being a class: every ordinary class (``dict``, ``list``, and any + class without a custom metaclass) is itself an instance of ``type``, so ``type(dict)`` would + otherwise return the bare ``type`` metaclass as a plain value. + """ + from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler + + unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset()) + restricted_type = unpickler.find_class("builtins", "type") + module, name = already_allowed_class_global + resolved_class = unpickler.find_class(module, name) + with pytest.raises(pickle.UnpicklingError, match="deserialization blocked"): + restricted_type(resolved_class) + + +def test_restricted_decode_blocks_two_step_type_metaclass_chain(): + """End-to-end regression for the two-``REDUCE`` chain: pickle's ``REDUCE`` opcode calls whatever + object is on top of the stack without re-resolving it through ``find_class``, so a value + returned by one restricted call can be memoized and reused as the callable of a *second* + ``REDUCE`` with different arguments. Recovering the bare ``type`` metaclass via + ``type(dict)`` and then calling it with three arguments must still be rejected, not accepted as + a constructed class with caller-supplied data. + """ + payload = metaclass_via_builtin_type_pickle(("builtins", "dict"), "DynamicConfig", {"token": "value"}) + checkpoint_value = { + _PICKLE_MARKER: base64.b64encode(payload).decode("ascii"), + _TYPE_MARKER: "builtins:DynamicConfig", + } + + with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"): + decode_checkpoint_value(checkpoint_value, allowed_types=frozenset()) + + +@pytest.mark.parametrize( + ("alias_module", "alias_name"), + [ + ("agent_framework._workflows._edge", "builtin_type"), + ("agent_framework._workflows._request_info_mixin", "builtin_type"), + ], +) +def test_restricted_decode_blocks_type_reexported_under_framework_prefix(alias_module: str, alias_name: str): + """A framework module attribute that happens to be ``builtins.type`` under an alias (for example + ``from builtins import type as builtin_type``) must not provide a second route to class + construction once the direct ``builtins:type`` key resolves to the restricted call instead. + """ + payload = dynamic_dict_subclass_pickle(alias_module, alias_name, "DynamicConfig", {"token": "value"}) + checkpoint_value = { + _PICKLE_MARKER: base64.b64encode(payload).decode("ascii"), + # Matches what the dynamically constructed class actually reports, so this exercises the + # metaclass check rather than the (otherwise orthogonal) _verify_type integrity check. + _TYPE_MARKER: f"{alias_module}:DynamicConfig", + } + + with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"): + decode_checkpoint_value(checkpoint_value, allowed_types=frozenset()) + + +def test_restricted_decode_blocks_enum_metaclass_reexported_under_framework_prefix(): + """``enum.EnumMeta``/``EnumType``, imported into the encoding module for isinstance checks, is a + metaclass and must not be resolvable as a checkpoint global either. + """ + from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler + + unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset()) + with pytest.raises(pickle.UnpicklingError, match="metaclass"): + unpickler.find_class("agent_framework._workflows._checkpoint_encoding", "EnumMeta") + + +def test_restricted_decode_blocks_type_via_restricted_getattr(): + """``getattr(dict, "__class__")`` is a second route to the raw ``type`` factory; the restricted + ``getattr`` helper must reject it even though ``dict`` itself is allowed. + """ + from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler + + unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset()) + with pytest.raises(pickle.UnpicklingError, match="nested type"): + unpickler._restricted_getattr(dict, "__class__") + + +def test_restricted_decode_blocks_unpopulated_enum_reexported_under_framework_prefix(): + """``agent_framework._serialization`` re-exports the stdlib ``Enum`` base class (used for + ``isinstance`` checks). ``Enum`` has no members of its own, so ``EnumMeta.__call__`` would + dispatch a two-argument call to the Enum functional API and construct a brand-new class -- + the same class-construction hazard as calling ``type`` directly. A populated application or + framework enum is not affected: Python refuses to subclass (functional-API or + otherwise) an Enum that already has members, on every supported Python version, so only + member-less Enum-family classes need to be rejected here. + """ + from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler + + unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset()) + with pytest.raises(pickle.UnpicklingError, match="deserialization blocked"): + unpickler.find_class("agent_framework._serialization", "Enum") + + +def test_restricted_decode_allows_populated_enum_value_round_trip(): + """A populated enum explicitly allowed by key continues to round-trip by value, proving the + member-less rejection in the previous test does not affect ordinary enum checkpointing. + """ + type_key = f"{_PopulatedColor.__module__}:{_PopulatedColor.__qualname__}" + encoded = encode_checkpoint_value(_PopulatedColor.RED) + decoded = decode_checkpoint_value(encoded, allowed_types=frozenset({type_key})) + + assert decoded is _PopulatedColor.RED + + +def test_restricted_decode_blocks_explicitly_allowed_metaclass(): + """Explicitly listing a metaclass's own key in ``allowed_types`` must not un-block class + construction; there is no legitimate reconstructible application type that is itself a + metaclass (``builtins:type`` is handled separately -- see the restricted-call tests above). + """ + from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler + + unpickler = _RestrictedUnpickler(pickle.dumps(object), frozenset({"enum:EnumMeta"})) + with pytest.raises(pickle.UnpicklingError, match="metaclass"): + unpickler.find_class("enum", "EnumMeta") + + +async def test_file_storage_rejects_edited_checkpoint_class_construction(): + """End-to-end regression: a checkpoint file edited on disk to replace a saved value's pickle + envelope with a REDUCE-based dynamic-class-construction payload must be rejected on ``load()``, + not accepted as the constructed class. + """ + from agent_framework import WorkflowCheckpoint + + with tempfile.TemporaryDirectory() as tmpdir: + storage = FileCheckpointStorage(tmpdir) + checkpoint = WorkflowCheckpoint( + workflow_name="test", + graph_signature_hash="hash", + state={"_executor_state": {"worker": {"token": "initial"}}}, + ) + checkpoint_id = await storage.save(checkpoint) + + # Edit the checkpoint file directly on disk, independent of any in-memory state. + file_path = Path(tmpdir) / f"{checkpoint_id}.json" + payload = dynamic_dict_subclass_pickle("builtins", "type", "DynamicConfig", {"token": "modified"}) + data = json.loads(file_path.read_text(encoding="utf-8")) + data["state"]["_executor_state"]["worker"] = { + _PICKLE_MARKER: base64.b64encode(payload).decode("ascii"), + _TYPE_MARKER: "builtins:DynamicConfig", + } + file_path.write_text(json.dumps(data), encoding="utf-8") + + with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"): + await storage.load(checkpoint_id) + + @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."""