diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 3a7d361b61..f6fd376692 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -403,7 +403,7 @@ def _base64_to_unpickle(encoded: str, *, allowed_types: frozenset[str] | None = format is incompatible, or a disallowed type is encountered. """ try: - pickled = base64.b64decode(encoded.encode("ascii")) + pickled = base64.b64decode(encoded.encode("ascii"), validate=True) if allowed_types is not None: return _RestrictedUnpickler(pickled, allowed_types).load() return pickle.loads(pickled) # nosec # ruff:ignore[suspicious-pickle-usage] diff --git a/python/packages/core/tests/workflow/test_checkpoint_decode.py b/python/packages/core/tests/workflow/test_checkpoint_decode.py index a8115d17c2..9038410482 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_decode.py +++ b/python/packages/core/tests/workflow/test_checkpoint_decode.py @@ -8,6 +8,7 @@ from agent_framework import WorkflowCheckpointException from agent_framework._workflows._checkpoint_encoding import ( + _PICKLE_MARKER, # type: ignore _TYPE_MARKER, # type: ignore decode_checkpoint_value, encode_checkpoint_value, @@ -188,6 +189,16 @@ def test_decode_raises_on_type_mismatch() -> None: decode_checkpoint_value(encoded) +def test_decode_raises_on_malformed_base64() -> None: + """Test that decoding rejects non-Base64 characters in a pickle payload.""" + encoded = encode_checkpoint_value((1, 2, 3)) + assert isinstance(encoded, dict) + encoded[_PICKLE_MARKER] = cast(str, encoded[_PICKLE_MARKER]) + "!!!!" + + with pytest.raises(WorkflowCheckpointException, match="Failed to decode pickled checkpoint data"): + decode_checkpoint_value(encoded) + + class NotADataclass: # noqa: B903 """A regular class that is not a dataclass."""