From 633f7260d07f9da468eafa72162af39f06d0b7cc Mon Sep 17 00:00:00 2001 From: jhLi Date: Tue, 19 May 2026 00:37:57 +0800 Subject: [PATCH] fix(core): restrict unpickler module-prefix allowlist to types only The `_RestrictedUnpickler.find_class` method allows any attribute from modules matching the `agent_framework.` and `openai.types.` prefixes. This permits non-type globals such as functions and sub-modules, which can be chained through `builtins.getattr` to reach `pickle.loads` and execute unrestricted nested pickle payloads. Limit module-prefix allowlisting to actual types by checking that the resolved global is an instance of `type` before returning it. Co-Authored-By: Claude Opus 4.7 --- .../_workflows/_checkpoint_encoding.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index dd1fb3d704..979389f49e 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -99,14 +99,24 @@ def __init__(self, data: bytes, allowed_types: frozenset[str]) -> None: def find_class(self, module: str, name: str) -> type: type_key = f"{module}:{name}" - if ( - type_key in _BUILTIN_ALLOWED_TYPE_KEYS - or type_key in self._allowed_types - or module.startswith(_FRAMEWORK_MODULE_PREFIX) - or module.startswith(_OPENAI_MODULE_PREFIX) - ): + if type_key in _BUILTIN_ALLOWED_TYPE_KEYS or type_key in self._allowed_types: return super().find_class(module, name) # type: ignore[no-any-return] # nosec + if module.startswith(_FRAMEWORK_MODULE_PREFIX) or module.startswith( + _OPENAI_MODULE_PREFIX + ): + resolved = super().find_class(module, name) # nosec + # Reject non-type globals from allowed package prefixes. A broad + # module-prefix allowlist returns arbitrary module attributes such + # as functions and sub-modules, which can be combined with + # builtins.getattr to call unrestricted pickle.loads. + if isinstance(resolved, type): + return resolved # type: ignore[return-value] + raise pickle.UnpicklingError( + f"Checkpoint deserialization blocked for non-type global " + f"'{type_key}'." + ) + raise pickle.UnpicklingError( f"Checkpoint deserialization blocked for type '{type_key}'. " f"To allow this type, either include its 'module:qualname' key in the "