Python: Prevent pickle deserialization of untrusted HITL HTTP input (#4566)

* fix: prevent pickle deserialization of untrusted HITL input

Add strip_pickle_markers() to sanitize HTTP input before it reaches
pickle.loads() via the checkpoint decoding path. Applied as a 3-layer
defence-in-depth:

1. _app.py: sanitize req.get_json() at the HTTP boundary
2. _workflow.py: sanitize in _deserialize_hitl_response() before decode
3. _serialization.py: sanitize in reconstruct_to_type() as final guard

Any dict containing __pickled__ or __type__ markers from untrusted
sources is replaced with None, blocking arbitrary code execution via
crafted payloads to POST /workflow/respond/{instanceId}/{requestId}.

Includes 12 new unit tests covering the sanitizer and end-to-end
attack prevention.

* refactor: address review concerns for pickle fix

1. Remove deserialize_value() fallback in _deserialize_hitl_response
   untrusted HITL data now returns as-is when no type hint is available,
   never flowing into pickle.loads().

2. Move strip_pickle_markers() out of reconstruct_to_type()  the function
   is general-purpose again; untrusted-data callers are responsible for
   sanitizing first (documented with NOTE comment).

3. Define _PICKLE_MARKER/_TYPE_MARKER as local constants with import-time
   assertions against core's values  decouples from private names while
   failing loudly if core ever changes them.

4. Update tests to reflect new responsibility boundaries.

* fix: simplify warning message and fix ruff RUF001 lint

* fix: suppress pyright reportPrivateUsage on core marker imports

* Lower marker-strip log from warning to debug to avoid log flooding

* Replace assert with RuntimeError for marker sync checks (ruff S101)

* Fix pyright and ruff CI errors in security fix

- Use cast() for dict/list comprehensions in strip_pickle_markers (pyright)
- type: ignore for narrowed dict return in _workflow.py (pyright)
- Simplify marker imports: use core constants directly, remove local copies
- Remove duplicate pyright ignore comment

* Remove duplicate end-to-end test in TestStripPickleMarkers

* Suppress mypy redundant-cast on list cast needed by pyright
This commit is contained in:
Ahmed Muhsin
2026-03-10 14:29:33 -05:00
committed by GitHub
Unverified
parent 55fc882ca8
commit 09b3e2e4f0
4 changed files with 141 additions and 10 deletions
@@ -21,6 +21,7 @@ from agent_framework_azurefunctions._serialization import (
deserialize_value,
reconstruct_to_type,
serialize_value,
strip_pickle_markers,
)
@@ -353,7 +354,11 @@ class TestReconstructToType:
assert result.comment == "Great"
def test_reconstruct_from_checkpoint_markers(self) -> None:
"""Test that data with checkpoint markers is decoded via deserialize_value."""
"""Test that data with checkpoint markers is decoded via deserialize_value.
reconstruct_to_type is general-purpose and handles trusted checkpoint
data. Untrusted HITL callers must call strip_pickle_markers() first.
"""
original = SampleData(value=99, name="marker-test")
encoded = serialize_value(original)
@@ -372,3 +377,73 @@ class TestReconstructToType:
result = reconstruct_to_type(data, Unrelated)
assert result == data
def test_reconstruct_strips_injected_pickle_markers(self) -> None:
"""End-to-end: strip_pickle_markers + reconstruct_to_type blocks attack.
This mirrors the real HITL flow where callers sanitize before reconstruction.
"""
malicious = {"__pickled__": "gASVDgAAAAAAAACMBHRlc3SULg==", "__type__": "builtins:str"}
sanitized = strip_pickle_markers(malicious)
result = reconstruct_to_type(sanitized, str)
assert result is None
class TestStripPickleMarkers:
"""Security tests for strip_pickle_markers — the defence-in-depth layer
that prevents untrusted HTTP input from reaching pickle.loads()."""
def test_strips_top_level_pickle_marker(self) -> None:
"""A dict containing __pickled__ must be replaced with None."""
data = {"__pickled__": "PAYLOAD", "__type__": "os:system"}
assert strip_pickle_markers(data) is None
def test_strips_top_level_type_marker_only(self) -> None:
"""Even __type__ alone (without __pickled__) must be neutralised."""
data = {"__type__": "os:system", "other": "value"}
assert strip_pickle_markers(data) is None
def test_strips_nested_pickle_marker(self) -> None:
"""Pickle markers nested inside a dict must be neutralised."""
data = {"safe": "value", "nested": {"__pickled__": "PAYLOAD", "__type__": "os:system"}}
result = strip_pickle_markers(data)
assert result == {"safe": "value", "nested": None}
def test_strips_pickle_marker_in_list(self) -> None:
"""Pickle markers inside a list element must be neutralised."""
data = [{"__pickled__": "PAYLOAD"}, "safe"]
result = strip_pickle_markers(data)
assert result == [None, "safe"]
def test_strips_deeply_nested_marker(self) -> None:
"""Deeply nested pickle markers must be neutralised."""
data = {"a": {"b": {"c": {"__pickled__": "deep"}}}}
result = strip_pickle_markers(data)
assert result == {"a": {"b": {"c": None}}}
def test_preserves_safe_dict(self) -> None:
"""Dicts without pickle markers must be left untouched."""
data = {"approved": True, "reason": "Looks good"}
assert strip_pickle_markers(data) == data
def test_preserves_primitives(self) -> None:
"""Primitive values must pass through unchanged."""
assert strip_pickle_markers("hello") == "hello"
assert strip_pickle_markers(42) == 42
assert strip_pickle_markers(None) is None
assert strip_pickle_markers(True) is True
def test_preserves_safe_list(self) -> None:
"""Lists without pickle markers must be left untouched."""
data = [1, "two", {"key": "value"}]
assert strip_pickle_markers(data) == data
def test_mixed_safe_and_malicious(self) -> None:
"""Only the malicious entries should be stripped; safe entries remain."""
data = {
"user_input": "hello",
"evil": {"__pickled__": "PAYLOAD", "__type__": "os:system"},
"count": 42,
}
result = strip_pickle_markers(data)
assert result == {"user_input": "hello", "evil": None, "count": 42}