fix(picklescan): safely parse bounded PyTorch tensor batches

packages/modelaudit-picklescan/CHANGELOG.md

@@ -7,6 +7,10 @@ and this package adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Bug Fixes
+
+- Validate bounded batched PyTorch state-dictionary entries without falsely flagging canonical tensor reconstruction.
+
 ## [0.1.10](https://github.com/promptfoo/modelaudit/compare/modelaudit-picklescan-v0.1.9...modelaudit-picklescan-v0.1.10) (2026-07-22)
 
 ### Bug Fixes

packages/modelaudit-picklescan/pyproject.toml

@@ -47,6 +47,7 @@ line-length = 120
 target-version = "py310"
 
 [tool.ruff.lint]
+select = ["E4", "E7", "E9", "F"]
 extend-select = [
     "I",
     "W",

packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py

@@ -1874,6 +1874,10 @@ def _pytorch_storage_keys_from_pickle_bytes(
         )
 
     marker = object()
+    canonical_tensor = object()
+    canonical_batch_placeholder = object()
+    canonical_batch_entries: list[tuple[str, object]] = []
+    canonical_batch_target: object | None = None
     memo: dict[int, Any] = {}
     stack: list[Any] = []
     referenced_keys: set[str] = set()
@@ -1932,21 +1936,72 @@ def apply_setitems_to_target(items: tuple[tuple[Any, Any], ...]) -> None:
         else:
             poison_stack_top()
 
-    def pop_marked_tuple() -> tuple[Any, ...] | None:
+    def pop_marked_tuple(*, max_width: int = _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH) -> tuple[Any, ...] | None:
         items: list[Any] = []
         while stack:
             item = stack.pop()
             if item is marker:
                 return tuple(reversed(items))
             items.append(item)
-            if len(items) > _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH:
-                raise ValueError("PyTorch storage persistent ID tuple exceeded trust parser width")
+            if len(items) > max_width:
+                raise ValueError("PyTorch storage trust parser marked collection exceeded its width limit")
         return None
 
+    def compact_canonical_setitems_stack() -> bool:
+        nonlocal canonical_batch_target
+
+        if len(canonical_batch_entries) >= _PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH:
+            return False
+        for marker_index, item in enumerate(stack):
+            if item is not marker or marker_index == 0:
+                continue
+            target = stack[marker_index - 1]
+            if not isinstance(target, dict | _PytorchOrderedDictState):
+                continue
+            if isinstance(target, _PytorchOrderedDictState) and target.used_as_hooks:
+                continue
+            if canonical_batch_target is not None and target is not canonical_batch_target:
+                continue
+            pair_index = marker_index + 1
+            if pair_index < len(stack) and stack[pair_index] is canonical_batch_placeholder:
+                pair_index += 1
+            if pair_index + 1 >= len(stack):
+                continue
+            key = stack[pair_index]
+            value = stack[pair_index + 1]
+            if not isinstance(key, str):
+                continue
+            if value is not canonical_tensor and not isinstance(value, (str, int, float, bytes, type(None))):
+                continue
+            if value is not canonical_tensor and not (
+                any(item is canonical_tensor for item in stack[pair_index + 2 :])
+                or any(entry_value is canonical_tensor for _entry_key, entry_value in canonical_batch_entries)
+            ):
+                continue
+            if canonical_batch_target is None:
+                canonical_batch_target = target
+                stack.insert(marker_index + 1, canonical_batch_placeholder)
+                pair_index += 1
+            canonical_batch_entries.append((key, value))
+            del stack[pair_index : pair_index + 2]
+            return True
+        return False
+
     def within_limits() -> bool:
+        while len(stack) > _PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH:
+            if not compact_canonical_setitems_stack():
+                return False
+        if canonical_batch_entries and not any(
+            item is marker
+            and marker_index > 0
+            and stack[marker_index - 1] is canonical_batch_target
+            and marker_index + 1 < len(stack)
+            and stack[marker_index + 1] is canonical_batch_placeholder
+            for marker_index, item in enumerate(stack)
+        ):
+            return False
         return (
-            len(stack) <= _PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH
-            and len(memo) <= _PYTORCH_STORAGE_TRUST_MAX_MEMO_ENTRIES
+            len(memo) <= _PYTORCH_STORAGE_TRUST_MAX_MEMO_ENTRIES
             and len(referenced_keys) <= _PYTORCH_STORAGE_TRUST_MAX_REFERENCED_KEYS
         )
 
@@ -2041,6 +2096,7 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any:
             tensor_rebuild_uses.add((function.position, reduce_position))
             if rebuild_tensor_v2_args_are_canonical(args):
                 canonical_tensor_rebuild_invocations.add((function.position, reduce_position))
+                return canonical_tensor
             else:
                 invalidate_tensor_rebuild_proof()
             return None
@@ -2192,13 +2248,35 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any:
                 key = stack.pop()
                 apply_setitems_to_target(((key, value),))
             elif opcode_name == "SETITEMS":
-                setitem_items = pop_marked_tuple()
+                setitem_items = pop_marked_tuple(max_width=_PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH)
                 if setitem_items is None or len(setitem_items) % 2 != 0 or not stack:
-                    clear_stack_after_malformed_provenance()
-                    continue
-                apply_setitems_to_target(
-                    tuple((setitem_items[index], setitem_items[index + 1]) for index in range(0, len(setitem_items), 2))
-                )
+                    if (
+                        setitem_items is None
+                        or not setitem_items
+                        or setitem_items[0] is not canonical_batch_placeholder
+                        or not stack
+                        or stack[-1] is not canonical_batch_target
+                        or (len(setitem_items) - 1) % 2 != 0
+                    ):
+                        clear_stack_after_malformed_provenance()
+                        continue
+                    remaining_items = setitem_items[1:]
+                    if (
+                        len(canonical_batch_entries) + len(remaining_items) // 2
+                        > _PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH
+                    ):
+                        return _PytorchStorageReferenceParse(set(), {}, set(), set(), False, False)
+                    setitem_pairs = tuple(canonical_batch_entries) + tuple(
+                        (remaining_items[index], remaining_items[index + 1])
+                        for index in range(0, len(remaining_items), 2)
+                    )
+                    canonical_batch_entries.clear()
+                    canonical_batch_target = None
+                else:
+                    setitem_pairs = tuple(
+                        (setitem_items[index], setitem_items[index + 1]) for index in range(0, len(setitem_items), 2)
+                    )
+                apply_setitems_to_target(setitem_pairs)
             elif opcode_name == "BINPERSID":
                 pid = stack.pop() if stack else None
                 storage_ref = storage_ref_from_pid(pid)

packages/modelaudit-picklescan/tests/test_api.py

@@ -3432,6 +3432,129 @@ def test_scan_file_suppresses_rebuild_tensor_v2_for_real_ordered_state_dict(tmp_
     assert _torch_rebuild_tensor_v2_warning_dicts(report) == []
 
 
+def _write_large_batched_pytorch_state_dict(
+    path: Path,
+    *,
+    malicious: bool = False,
+    entry_count: int = 600,
+    leading_metadata: bool = False,
+) -> None:
+    entries: list[bytes] = []
+    if leading_metadata:
+        metadata_key = b"_extra_state"
+        entries.append(b"X" + len(metadata_key).to_bytes(4, "little") + metadata_key + b"K\x01")
+    for index in range(entry_count):
+        key_bytes = f"weight_{index}".encode("ascii")
+        key = b"X" + len(key_bytes).to_bytes(4, "little") + key_bytes
+        if malicious and index == entry_count - 1:
+            value = b"cos\nsystem\n(S'echo malicious-near-match'\ntR"
+        elif index == 0:
+            value = _pytorch_rebuild_tensor_v2_payload(key=str(index)).removeprefix(b"\x80\x04").removesuffix(b".")
+        else:
+            value = _pytorch_rebuild_tensor_v2_reduce_expr(key=str(index))
+        entries.append(key + value)
+
+    payload = b"\x80\x04" + _global(b"collections", b"OrderedDict") + b")R(" + b"".join(entries) + b"u."
+    with zipfile.ZipFile(path, "w") as archive:
+        archive.writestr("archive/data.pkl", payload)
+        archive.writestr("archive/version", "3\n")
+        archive.writestr("archive/byteorder", "little")
+        for index in range(entry_count - int(malicious)):
+            archive.writestr(f"archive/data/{index}", b"\x00" * 24)
+
+
+@pytest.mark.parametrize(("entry_count", "leading_metadata"), [(600, False), (1000, False), (600, True)])
+def test_pytorch_storage_trust_parses_large_batched_state_dict_without_framework(
+    tmp_path: Path,
+    entry_count: int,
+    leading_metadata: bool,
+) -> None:
+    archive_path = tmp_path / f"large-batched-state-{entry_count}-{leading_metadata}.pt"
+    _write_large_batched_pytorch_state_dict(
+        archive_path,
+        entry_count=entry_count,
+        leading_metadata=leading_metadata,
+    )
+    with zipfile.ZipFile(archive_path) as archive:
+        payload = archive.read("archive/data.pkl")
+
+    parsed = package_api._pytorch_storage_keys_from_pickle_bytes(payload)
+
+    assert parsed.parse_complete is True
+    assert len(parsed.referenced_keys) == entry_count
+    assert len(parsed.canonical_tensor_rebuild_invocations) == len(parsed.referenced_keys)
+
+
+def test_scan_file_suppresses_rebuild_tensor_v2_for_large_batched_state_dict(tmp_path: Path) -> None:
+    _require_torch_distribution()
+    archive_path = tmp_path / "large-batched-state.pt"
+    _write_large_batched_pytorch_state_dict(archive_path)
+
+    report = _scan_file_report_dict_subprocess(archive_path)
+
+    assert report["status"] == "complete"
+    assert report["verdict"] == "clean"
+    assert _torch_rebuild_tensor_v2_warning_dicts(report) == []
+    assert not any(finding["rule_code"] == "PERSISTENT_ID" for finding in report["findings"])
+
+
+@pytest.mark.parametrize("leading_metadata", [False, True])
+def test_scan_file_preserves_malicious_call_in_large_batched_state_dict(
+    tmp_path: Path,
+    leading_metadata: bool,
+) -> None:
+    archive_path = tmp_path / f"large-batched-malicious-state-{leading_metadata}.pt"
+    _write_large_batched_pytorch_state_dict(archive_path, malicious=True, leading_metadata=leading_metadata)
+
+    report = _scan_file_report_dict_subprocess(archive_path)
+
+    assert report["verdict"] == "malicious"
+    assert any(
+        finding["severity"] == "critical"
+        and finding["details"].get("module") in {"os", "posix", "nt"}
+        and finding["details"].get("name") == "system"
+        for finding in report["findings"]
+    )
+
+
+def test_pytorch_storage_trust_rejects_setitems_beyond_stack_limit() -> None:
+    entries = package_api._PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH // 2 + 1
+    oversized_payload = b"\x80\x04}(" + (b"K\x00K\x00" * entries) + b"u."
+
+    parsed = package_api._pytorch_storage_keys_from_pickle_bytes(oversized_payload)
+
+    assert parsed.parse_complete is False
+    assert parsed.referenced_keys == set()
+    assert parsed.canonical_tensor_rebuild_invocations == set()
+
+
+def test_pytorch_storage_trust_rejects_canonical_setitems_beyond_bounded_batch(tmp_path: Path) -> None:
+    archive_path = tmp_path / "oversized-canonical-batch.pt"
+    _write_large_batched_pytorch_state_dict(
+        archive_path,
+        entry_count=package_api._PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH + 1,
+    )
+    with zipfile.ZipFile(archive_path) as archive:
+        payload = archive.read("archive/data.pkl")
+
+    parsed = package_api._pytorch_storage_keys_from_pickle_bytes(payload)
+
+    assert parsed.parse_complete is False
+    assert parsed.referenced_keys == set()
+    assert parsed.canonical_tensor_rebuild_invocations == set()
+
+
+def test_pytorch_storage_trust_preserves_tuple_width_limit() -> None:
+    tuple_items = package_api._PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH + 1
+    oversized_payload = b"\x80\x04(" + (b"K\x00" * tuple_items) + b"t."
+
+    parsed = package_api._pytorch_storage_keys_from_pickle_bytes(oversized_payload)
+
+    assert parsed.parse_complete is False
+    assert parsed.referenced_keys == set()
+    assert parsed.canonical_tensor_rebuild_invocations == set()
+
+
 def test_scan_bytes_keeps_rebuild_tensor_v2_warning_for_raw_data_pickle(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: