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."""