From 91843c6ec9d379c1be9a9d846f338a937e97641a Mon Sep 17 00:00:00 2001 From: Herdiyan Adam Putra Date: Sat, 22 Aug 2026 22:12:59 +0700 Subject: [PATCH] Harden FoundryCheckpointStore: apply RestrictedUnpickler allowlist on load FoundryCheckpointStore was the only checkpoint backend calling decode_checkpoint_value(item.value) without allowed_types, falling back to unrestricted pickle.loads. FileCheckpointStorage and CosmosCheckpointStorage both thread the RestrictedUnpickler allowlist through. Thread allowed_checkpoint_types through FoundryCheckpointStore and CheckpointStoreProvider, and pass allowed_types on both load() and list_checkpoints(). Default (frozenset()) restricts deserialization to the built-in safe set plus framework and OpenAI SDK types, matching the other backends. Legitimate checkpoints round-trip unchanged. --- .../_state_store.py | 51 +++++++++++++++++-- .../foundry_hosting/tests/test_state_store.py | 45 ++++++++++++++++ 2 files changed, 91 insertions(+), 5 deletions(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py index 9df38a4332..aa2c225258 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py @@ -68,23 +68,43 @@ def get_store( class FoundryCheckpointStore: - """Checkpoint store backed by the `FoundryStateStore`.""" + """Checkpoint store backed by the `FoundryStateStore`. + + By default, checkpoint deserialization is restricted to a built-in set of safe Python types + (primitives, datetime, uuid, ...), all ``agent_framework`` internal types, and OpenAI SDK types + (``openai.types``). To allow additional application-specific types, pass them via the + ``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format. This mirrors the + behavior of ``FileCheckpointStorage`` and ``CosmosCheckpointStorage``, ensuring the + ``RestrictedUnpickler`` allowlist is applied on every hosted checkpoint load rather than + falling back to unrestricted ``pickle.loads``. + """ DEFAULT_ROOT_SCOPE = "checkpoints" - def __init__(self, context_id: str, platform_context: FoundryAgentRequestContext) -> None: + def __init__( + self, + context_id: str, + platform_context: FoundryAgentRequestContext, + *, + allowed_checkpoint_types: list[str] | None = None, + ) -> None: """Initialize a Foundry-scoped checkpoint store for the given context ID. Args: context_id: A string that uniquely identifies the context for which the checkpoint store is scoped. This can be used to isolate checkpoints for different workflow runs. platform_context: The request-scoped platform context for the current request. + allowed_checkpoint_types: Additional types (beyond the built-in safe set + and framework types) that are permitted during checkpoint + deserialization. Each entry should be a ``"module:qualname"`` + string (e.g., ``"my_app.models:MyState"``). """ if not context_id: raise ValueError("context_id must be provided to initialize a FoundryCheckpointStore.") self.context_id = context_id self.platform_context = platform_context + self._allowed_types: frozenset[str] = frozenset(allowed_checkpoint_types or []) async def _get_store(self) -> FoundryStateStore: return await FoundryStateStore.get_or_create( @@ -132,7 +152,7 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: item = await store.get_item(checkpoint_id, call_id=self.platform_context.call_id) if item is None: raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}") - return WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value)) + return WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value, allowed_types=self._allowed_types)) async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: """List all workflow checkpoints for a given workflow name.""" @@ -148,7 +168,9 @@ async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoi item = await store.get_item(item_key.key, call_id=self.platform_context.call_id) if item is None: continue - checkpoint = WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value)) + checkpoint = WorkflowCheckpoint.from_dict( + decode_checkpoint_value(item.value, allowed_types=self._allowed_types) + ) if checkpoint.workflow_name == workflow_name: checkpoints.append(checkpoint) if not page.has_more or page.last_id is None: @@ -181,8 +203,23 @@ class CheckpointStoreProvider(ContextScopedStoreProvider[CheckpointStorage]): and clean up older checkpoints without affecting another workflow context. This defaults to using the `FoundryCheckpointStore` in all environments. + + Checkpoint deserialization is restricted to the built-in safe type set plus framework and + OpenAI SDK types by default. Applications that persist custom state types can widen the + allowlist by passing ``allowed_checkpoint_types`` (``"module:qualname"`` strings), which is + threaded to every ``FoundryCheckpointStore`` this provider creates. """ + def __init__(self, *, allowed_checkpoint_types: list[str] | None = None) -> None: + """Initialize the provider. + + Args: + allowed_checkpoint_types: Additional types (beyond the built-in safe set + and framework types) permitted during checkpoint deserialization, + as ``"module:qualname"`` strings. + """ + self._allowed_checkpoint_types = allowed_checkpoint_types + def get_store( self, *, @@ -194,7 +231,11 @@ def get_store( if not context_id: raise ValueError("context_id must be provided to get a checkpoint store.") - return FoundryCheckpointStore(context_id, platform_context) + return FoundryCheckpointStore( + context_id, + platform_context, + allowed_checkpoint_types=self._allowed_checkpoint_types, + ) # endregion Checkpoint persistence diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index e08d4dfabd..04c962b3e3 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -114,6 +114,51 @@ async def test_load_raises_for_missing_checkpoint() -> None: await FoundryCheckpointStore("context-1", _platform_context()).load("missing") +async def test_load_defaults_to_restricted_unpickler_and_blocks_unlisted_types() -> None: + # A raw store item that was NOT produced by encode_checkpoint_value (e.g. planted + # directly into the backing store) carries an arbitrary pickle payload. Under the + # default (restricted) unpickler this must be rejected rather than executed. + import base64 + import os + import pickle + + class _Gadget: + def __reduce__(self) -> Any: + return (os.system, ("echo should-not-run",)) + + malicious_value = { + "__pickled__": base64.b64encode(pickle.dumps(_Gadget())).decode("ascii"), + "__type__": "builtins:int", + } + + store = _store() + store.get_item = AsyncMock(return_value=SimpleNamespace(value=malicious_value)) + + with ( + patch( + "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", + new=AsyncMock(return_value=store), + ), + pytest.raises(WorkflowCheckpointException), + ): + await FoundryCheckpointStore("context-1", _platform_context()).load("checkpoint-1") + + +def test_checkpoint_store_provider_threads_allowed_types() -> None: + provider = CheckpointStoreProvider(allowed_checkpoint_types=["my_app.models:MyState"]) + store = provider.get_store( + config=_config(is_hosted=True), context_id="context-1", platform_context=_platform_context() + ) + + assert isinstance(store, FoundryCheckpointStore) + assert store._allowed_types == frozenset({"my_app.models:MyState"}) + + +def test_checkpoint_store_defaults_to_empty_allowlist() -> None: + store = FoundryCheckpointStore("context-1", _platform_context()) + assert store._allowed_types == frozenset() + + async def test_list_checkpoints_paginates_and_filters_by_workflow() -> None: store = _store() matching = _checkpoint("checkpoint-1")