diff --git a/packages/modelaudit-picklescan/CHANGELOG.md b/packages/modelaudit-picklescan/CHANGELOG.md index 0c0a1fad7..fbcba5f31 100644 --- a/packages/modelaudit-picklescan/CHANGELOG.md +++ b/packages/modelaudit-picklescan/CHANGELOG.md @@ -7,6 +7,11 @@ 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. +- Keep scanning storage members between the trusted and expanded pickle probe sizes. + ## [0.1.10](https://github.com/promptfoo/modelaudit/compare/modelaudit-picklescan-v0.1.9...modelaudit-picklescan-v0.1.10) (2026-07-22) ### Bug Fixes diff --git a/packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py b/packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py index 997b31d8f..27b4017ab 100644 --- a/packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py +++ b/packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py @@ -10,9 +10,9 @@ import time import zipfile from _collections import OrderedDict as _CANONICAL_COLLECTIONS_ORDERED_DICT -from collections.abc import Mapping +from collections.abc import Callable, Mapping from contextlib import suppress -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from importlib import import_module from importlib import metadata as importlib_metadata from pathlib import Path @@ -107,6 +107,7 @@ _MAX_PYTORCH_ZIP_STORAGE_REFERENCE_DATA_PICKLE_BYTES = 10 * 1024 * 1024 _MAX_PYTORCH_ZIP_STORAGE_REFERENCE_TOTAL_DATA_PICKLE_BYTES = 64 * 1024 * 1024 _PYTORCH_STORAGE_TRUST_MAX_OPCODES = 100_000 +_PYTORCH_STORAGE_TRUST_MAX_PROVENANCE_NODES = 100_000 _PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH = 1024 _PYTORCH_STORAGE_TRUST_MAX_MEMO_ENTRIES = 100_000 _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH = 64 @@ -336,6 +337,16 @@ class _PytorchStorageRef: class _PytorchOrderedDictState: mutated: bool = False used_as_hooks: bool = False + contains_untrusted_storage: bool = False + keys_with_tracked_provenance: set[object] = field(default_factory=set) + + +class _PytorchDictionaryState(dict[Any, Any]): + __slots__ = ("contains_untrusted_storage",) + + def __init__(self) -> None: + super().__init__() + self.contains_untrusted_storage = False @dataclass(frozen=True) @@ -378,6 +389,9 @@ class _PytorchStorageReferenceParse: canonical_tensor_rebuild_invocations: set[tuple[int, int]] parse_complete: bool all_persistent_ids_are_pytorch_storage: bool + discarded_tracked_storage_references: bool = False + used_streaming_batch_compaction: bool = False + accepted_oversized_state_batch: bool = False @dataclass(frozen=True) @@ -389,6 +403,7 @@ class _PytorchZipDataPickleTrust: @dataclass(frozen=True) class _PytorchZipStorageEntries: trusted_entry_ids: set[int] + expanded_trust_entry_ids: set[int] storage_probe_entry_ids: set[int] trusted_data_pkl_by_name: dict[str, _PytorchZipDataPickleTrust] @@ -918,7 +933,11 @@ def add_entry(entry: zipfile.ZipInfo) -> None: entry_id = id(entry) storage_probe_bytes = None if entry_id in storage_entries.trusted_entry_ids: - storage_probe_bytes = _TRUSTED_STORAGE_PICKLE_PROBE_BYTES + storage_probe_bytes = ( + _PICKLE_DISCOVERY_LONG_PROBE_BYTES + if entry_id in storage_entries.expanded_trust_entry_ids + else _TRUSTED_STORAGE_PICKLE_PROBE_BYTES + ) elif entry_id in storage_entries.storage_probe_entry_ids: storage_probe_bytes = _PICKLE_DISCOVERY_LONG_PROBE_BYTES if storage_probe_bytes is not None: @@ -968,6 +987,7 @@ def _validated_pytorch_storage_entry_ids( entries_by_name.setdefault(name, []).append(entry) trusted_entry_ids: set[int] = set() + expanded_trust_entry_ids: set[int] = set() storage_probe_entry_ids: set[int] = set() trusted_data_pkl_by_name: dict[str, _PytorchZipDataPickleTrust] = {} notices: list[Notice] = [] @@ -1108,7 +1128,10 @@ def _validated_pytorch_storage_entry_ids( ) validated_storage_keys = trusted_storage_keys - storage_size_mismatch_keys exact_trusted_storage_keys = ( - validated_storage_keys if reference_parse.all_persistent_ids_are_pytorch_storage else set() + validated_storage_keys + if reference_parse.all_persistent_ids_are_pytorch_storage + and not reference_parse.discarded_tracked_storage_references + else set() ) storage_probe_keys = trusted_storage_keys - exact_trusted_storage_keys if exact_trusted_storage_keys and not missing_storage_keys: @@ -1119,13 +1142,17 @@ def _validated_pytorch_storage_entry_ids( ), ) for storage_key in exact_trusted_storage_keys: - trusted_entry_ids.add(id(storage_entries_by_key[storage_key])) + entry_id = id(storage_entries_by_key[storage_key]) + trusted_entry_ids.add(entry_id) + if reference_parse.used_streaming_batch_compaction or reference_parse.accepted_oversized_state_batch: + expanded_trust_entry_ids.add(entry_id) for storage_key in storage_probe_keys: storage_probe_entry_ids.add(id(storage_entries_by_key[storage_key])) return ( _PytorchZipStorageEntries( trusted_entry_ids=trusted_entry_ids, + expanded_trust_entry_ids=expanded_trust_entry_ids, storage_probe_entry_ids=storage_probe_entry_ids, trusted_data_pkl_by_name=trusted_data_pkl_by_name, ), @@ -1151,6 +1178,7 @@ def _trusted_pytorch_data_pkl_from_storage_member_sizes( not reference_parse.parse_complete or not reference_parse.referenced_keys or not reference_parse.all_persistent_ids_are_pytorch_storage + or reference_parse.discarded_tracked_storage_references ): return None if not reference_parse.referenced_keys <= storage_member_sizes.keys(): @@ -1254,6 +1282,16 @@ def _trusted_storage_zip_entry_looks_like_pickle( is_frame_first_candidate = prefix.startswith(_PICKLE_FRAME_OPCODE) if not is_binary_pickle_candidate and not is_frame_first_candidate and prefix[0] not in _PROTO0_1_START_BYTES: return False + if max_probe_bytes > _TRUSTED_STORAGE_PICKLE_PROBE_BYTES: + try: + for _opcode, _arg, _position in pickletools.genops(prefix): + pass + except ValueError as error: + message = str(error).lower() + if "opcode" in message and "unknown" in message: + return False + except Exception: + pass sample = prefix if entry.file_size > len(prefix): @@ -1265,10 +1303,35 @@ def _trusted_storage_zip_entry_looks_like_pickle( deadline, ) if is_binary_pickle_candidate: - return _binary_pickle_probe_should_scan(sample, sample_is_prefix=entry.file_size > len(sample)) + return _binary_pickle_probe_should_scan( + sample, sample_is_prefix=entry.file_size > len(sample) + ) or _expanded_probe_preserves_trusted_scan(entry, sample, max_probe_bytes, _binary_pickle_probe_should_scan) if is_frame_first_candidate: return _frame_first_trusted_storage_probe_should_scan(sample) - return _proto0_or_1_trusted_storage_probe_should_scan(sample, sample_is_prefix=entry.file_size > len(sample)) + return _proto0_or_1_trusted_storage_probe_should_scan( + sample, sample_is_prefix=entry.file_size > len(sample) + ) or _expanded_probe_preserves_trusted_scan( + entry, sample, max_probe_bytes, _proto0_or_1_trusted_storage_probe_should_scan + ) + + +def _expanded_probe_preserves_trusted_scan( + entry: zipfile.ZipInfo, + sample: bytes, + max_probe_bytes: int, + predicate: Callable[..., bool], +) -> bool: + """Return whether the shorter trusted probe would have scanned this member. + + ``sample_is_prefix`` is derived from how much of the member the probe read, so widening the + probe can flip it from True to False and silently drop a member that the 4 KiB trusted probe + scanned. Members between the two probe sizes are exactly the window affected. Re-checking the + shorter view keeps a wider probe strictly additive for coverage. + """ + if max_probe_bytes <= _TRUSTED_STORAGE_PICKLE_PROBE_BYTES: + return False + short_sample = sample[:_TRUSTED_STORAGE_PICKLE_PROBE_BYTES] + return predicate(short_sample, sample_is_prefix=entry.file_size > len(short_sample)) def _binary_pickle_probe_should_scan(sample: bytes, *, sample_is_prefix: bool) -> bool: @@ -1847,6 +1910,9 @@ def _merge_pytorch_storage_reference_parses( canonical_tensor_rebuild_invocations=set(), parse_complete=True, all_persistent_ids_are_pytorch_storage=all(parsed.all_persistent_ids_are_pytorch_storage for parsed in parses), + discarded_tracked_storage_references=any(parsed.discarded_tracked_storage_references for parsed in parses), + used_streaming_batch_compaction=any(parsed.used_streaming_batch_compaction for parsed in parses), + accepted_oversized_state_batch=any(parsed.accepted_oversized_state_batch for parsed in parses), ) @@ -1874,6 +1940,12 @@ 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 + trusted_canonical_batch_seen = False + pending_uncanonical_metadata_batch = False memo: dict[int, Any] = {} stack: list[Any] = [] referenced_keys: set[str] = set() @@ -1883,16 +1955,33 @@ def _pytorch_storage_keys_from_pickle_bytes( tensor_rebuild_uses: set[tuple[int, int]] = set() tensor_rebuild_proof_valid = True all_persistent_ids_are_pytorch_storage = True + discarded_tracked_storage_references = False + used_streaming_batch_compaction = False + accepted_oversized_state_batch = False + provenance_nodes_inspected = 0 def invalidate_tensor_rebuild_proof() -> None: nonlocal tensor_rebuild_proof_valid tensor_rebuild_proof_valid = False def clear_stack_after_malformed_provenance() -> None: + nonlocal canonical_batch_target, trusted_canonical_batch_seen, discarded_tracked_storage_references + + if any(value_contains_tracked_provenance(value) for value in stack if value is not marker) or any( + value_contains_tracked_provenance(value) for _key, value in canonical_batch_entries + ): + discarded_tracked_storage_references = True invalidate_tensor_rebuild_proof() stack.clear() + canonical_batch_entries.clear() + canonical_batch_target = None + trusted_canonical_batch_seen = False def poison_stack_top() -> None: + nonlocal discarded_tracked_storage_references + + if value_contains_tracked_provenance(stack[-1]): + discarded_tracked_storage_references = True invalidate_tensor_rebuild_proof() stack[-1] = None @@ -1907,9 +1996,25 @@ def ordered_dict_is_empty_hooks(value: object) -> bool: value.used_as_hooks = True return True - def value_contains_tracked_provenance(value: object, seen: set[int] | None = None) -> bool: - if isinstance(value, _PytorchStorageRef | _PytorchOrderedDictState): + def value_contains_tracked_provenance( + value: object, + seen: set[int] | None = None, + *, + storage_only: bool = False, + ) -> bool: + nonlocal provenance_nodes_inspected + + provenance_nodes_inspected += 1 + if provenance_nodes_inspected > _PYTORCH_STORAGE_TRUST_MAX_PROVENANCE_NODES: + raise ValueError("PyTorch storage trust parser exceeded its bounded provenance workload") + if deadline is not None and provenance_nodes_inspected % 1024 == 0: + _check_pytorch_zip_deadline(deadline) + if value is canonical_tensor: + return not storage_only + if isinstance(value, _PytorchStorageRef): return True + if isinstance(value, _PytorchOrderedDictState): + return value.contains_untrusted_storage if storage_only else True if not isinstance(value, (tuple, list, dict)): return False if seen is None: @@ -1919,34 +2024,157 @@ def value_contains_tracked_provenance(value: object, seen: set[int] | None = Non return False seen.add(value_id) if isinstance(value, dict): - return any(value_contains_tracked_provenance(item, seen) for item in value.items()) - return any(value_contains_tracked_provenance(item, seen) for item in value) + if storage_only and isinstance(value, _PytorchDictionaryState): + return value.contains_untrusted_storage + return any( + value_contains_tracked_provenance(key, seen, storage_only=storage_only) + or value_contains_tracked_provenance(item, seen, storage_only=storage_only) + for key, item in value.items() + ) + return any(value_contains_tracked_provenance(item, seen, storage_only=storage_only) for item in value) + + def setitems_entry_is_safe(key: object, value: object) -> bool: + return isinstance(key, str) and ( + value is canonical_tensor + or ( + isinstance(value, (str, int, float, bytes, type(None), tuple, list, dict)) + and not value_contains_tracked_provenance(value, storage_only=True) + and ( + not trusted_canonical_batch_seen + or not value_contains_tracked_provenance(value) + or setitems_entry_contains_canonical_tensor(value) + ) + ) + ) + + def setitems_entry_contains_canonical_tensor(value: object) -> bool: + return value is canonical_tensor or ( + isinstance(value, (tuple, list, dict)) + and not value_contains_tracked_provenance(value, storage_only=True) + and value_contains_tracked_provenance(value) + ) def apply_setitems_to_target(items: tuple[tuple[Any, Any], ...]) -> None: + nonlocal discarded_tracked_storage_references + + canonical_storage_context = ( + trusted_canonical_batch_seen + or bool(canonical_batch_entries) + or any(value is canonical_tensor for value in stack) + or any(setitems_entry_contains_canonical_tensor(value) for _key, value in items) + ) + items_contain_untrusted_storage = any( + value_contains_tracked_provenance(key, storage_only=True) + or (value is not canonical_tensor and value_contains_tracked_provenance(value, storage_only=True)) + for key, value in items + ) + first_marker_index = next((index for index, value in enumerate(stack) if value is marker), len(stack)) + prior_state_contains_untrusted_storage = any( + value_contains_tracked_provenance(value, storage_only=True) + for index, value in enumerate(stack) + if isinstance(value, (tuple, list, dict, _PytorchOrderedDictState)) + and ( + index < first_marker_index + or index == len(stack) - 1 + or (index + 1 < len(stack) and stack[index + 1] is marker) + ) + ) + if canonical_storage_context and (items_contain_untrusted_storage or prior_state_contains_untrusted_storage): + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() if isinstance(stack[-1], _PytorchStorageRef): poison_stack_top() elif isinstance(stack[-1], _PytorchOrderedDictState): - mutate_tracked_ordered_dict(stack[-1]) + target = stack[-1] + target.contains_untrusted_storage = target.contains_untrusted_storage or items_contain_untrusted_storage + for key, value in items: + if key in target.keys_with_tracked_provenance: + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() + if value_contains_tracked_provenance(value): + target.keys_with_tracked_provenance.add(key) + else: + target.keys_with_tracked_provenance.discard(key) + mutate_tracked_ordered_dict(target) elif isinstance(stack[-1], dict): - stack[-1].update(items) + if isinstance(stack[-1], _PytorchDictionaryState): + stack[-1].contains_untrusted_storage |= items_contain_untrusted_storage + for key, value in items: + if key in stack[-1] and value_contains_tracked_provenance(stack[-1][key]): + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() + stack[-1][key] = value else: + if any( + value_contains_tracked_provenance(key) or value_contains_tracked_provenance(value) + for key, value in items + ): + discarded_tracked_storage_references = True 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, used_streaming_batch_compaction + + if len(canonical_batch_entries) >= _PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH: + return False + for marker_index in range(len(stack)): + item = stack[marker_index] + 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 setitems_entry_is_safe(key, value): + continue + if pair_index + 2 < len(stack) and stack[pair_index + 2] is marker: + 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] + used_streaming_batch_compaction = True + 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 ) @@ -2027,11 +2255,14 @@ def rebuild_tensor_v2_args_are_canonical(args: Any) -> bool: ) def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: + nonlocal discarded_tracked_storage_references + if isinstance(function, _PickleGlobalRef) and (function.module, function.name) == ( "collections", "OrderedDict", ): if args != () and value_contains_tracked_provenance(args): + discarded_tracked_storage_references = True invalidate_tensor_rebuild_proof() return _PytorchOrderedDictState() if args == () else None if isinstance(function, _PickleGlobalRef) and (function.module, function.name) == ( @@ -2041,10 +2272,14 @@ 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: + if value_contains_tracked_provenance(args): + discarded_tracked_storage_references = True invalidate_tensor_rebuild_proof() return None if value_contains_tracked_provenance(function) or value_contains_tracked_provenance(args): + discarded_tracked_storage_references = True invalidate_tensor_rebuild_proof() return None @@ -2063,6 +2298,19 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: if opcode_name in {"PROTO", "FRAME"}: continue if opcode_name == "STOP": + if canonical_batch_entries or (pending_uncanonical_metadata_batch and not trusted_canonical_batch_seen): + return _PytorchStorageReferenceParse(set(), {}, set(), set(), False, False) + if ( + stack + and isinstance(stack[-1], (tuple, list, dict, _PytorchOrderedDictState)) + and canonical_tensor_rebuild_invocations + and (value_contains_tracked_provenance(stack[-1], storage_only=True)) + ): + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() + if any(value_contains_tracked_provenance(value) for value in stack[:-1] if value is not marker): + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() if not isinstance(_pos, int) or _pos + 1 != len(pickle_data): invalidate_tensor_rebuild_proof() continue @@ -2089,8 +2337,13 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: if len(stack) < 2: clear_stack_after_malformed_provenance() continue - name = _coerce_pickle_string_arg(stack.pop()) - module = _coerce_pickle_string_arg(stack.pop()) + name_value = stack.pop() + module_value = stack.pop() + if value_contains_tracked_provenance(name_value) or value_contains_tracked_provenance(module_value): + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() + name = _coerce_pickle_string_arg(name_value) + module = _coerce_pickle_string_arg(module_value) stack.append( _PickleGlobalRef(module, name, position) if module is not None and name is not None else None ) @@ -2099,7 +2352,7 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: elif opcode_name == "EMPTY_LIST": stack.append([]) elif opcode_name == "EMPTY_DICT": - stack.append({}) + stack.append(_PytorchDictionaryState()) elif opcode_name == "TUPLE": tuple_value = pop_marked_tuple() if tuple_value is None: @@ -2125,7 +2378,20 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: if dict_items is None or len(dict_items) % 2 != 0: clear_stack_after_malformed_provenance() else: - stack.append({dict_items[index]: dict_items[index + 1] for index in range(0, len(dict_items), 2)}) + built_dict = _PytorchDictionaryState() + for index in range(0, len(dict_items), 2): + key = dict_items[index] + value = dict_items[index + 1] + if key in built_dict and value_contains_tracked_provenance(built_dict[key]): + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() + if value_contains_tracked_provenance(key, storage_only=True) or ( + value is not canonical_tensor + and value_contains_tracked_provenance(value, storage_only=True) + ): + built_dict.contains_untrusted_storage = True + built_dict[key] = value + stack.append(built_dict) elif opcode_name in {"BININT", "BININT1", "BININT2", "LONG", "LONG1", "LONG4", "INT", "FLOAT", "BINFLOAT"}: stack.append(arg) elif opcode_name == "NONE": @@ -2155,11 +2421,18 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: stack.append(memo[key]) elif opcode_name == "POP": if stack: - stack.pop() + popped_value = stack.pop() + if value_contains_tracked_provenance(popped_value): + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() else: invalidate_tensor_rebuild_proof() elif opcode_name == "POP_MARK": - if pop_marked_tuple() is None: + popped_values = pop_marked_tuple() + if popped_values is None: + invalidate_tensor_rebuild_proof() + elif any(value_contains_tracked_provenance(value) for value in popped_values): + discarded_tracked_storage_references = True invalidate_tensor_rebuild_proof() elif opcode_name == "DUP": if stack: @@ -2174,6 +2447,8 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: if isinstance(stack[-1], list): stack[-1].append(value) else: + if value_contains_tracked_provenance(value): + discarded_tracked_storage_references = True poison_stack_top() elif opcode_name == "APPENDS": appended_items = pop_marked_tuple() @@ -2183,6 +2458,8 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: if isinstance(stack[-1], list): stack[-1].extend(appended_items) else: + if any(value_contains_tracked_provenance(value) for value in appended_items): + discarded_tracked_storage_references = True poison_stack_top() elif opcode_name == "SETITEM": if len(stack) < 3: @@ -2192,13 +2469,54 @@ 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) + ) + batch_has_canonical_tensor = any( + setitems_entry_contains_canonical_tensor(value) for _key, value in setitem_pairs ) + if len(setitem_pairs) * 2 > _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH: + accepted_oversized_state_batch = True + if len(setitem_pairs) * 2 > _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH and ( + not tensor_rebuild_proof_valid + or not all(setitems_entry_is_safe(key, value) for key, value in setitem_pairs) + ): + return _PytorchStorageReferenceParse(set(), {}, set(), set(), False, False) + if ( + len(setitem_pairs) * 2 > _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH + and not batch_has_canonical_tensor + and not trusted_canonical_batch_seen + ): + pending_uncanonical_metadata_batch = True + apply_setitems_to_target(setitem_pairs) + if batch_has_canonical_tensor: + trusted_canonical_batch_seen = True + pending_uncanonical_metadata_batch = False elif opcode_name == "BINPERSID": pid = stack.pop() if stack else None storage_ref = storage_ref_from_pid(pid) @@ -2245,20 +2563,27 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: state = stack.pop() obj = stack.pop() if isinstance(obj, _PytorchStorageRef): + discarded_tracked_storage_references = True invalidate_tensor_rebuild_proof() stack.append(None) elif isinstance(obj, _PytorchOrderedDictState): if state is not None: + if value_contains_tracked_provenance(state, storage_only=True): + discarded_tracked_storage_references = True + invalidate_tensor_rebuild_proof() mutate_tracked_ordered_dict(obj) stack.append(obj) else: if value_contains_tracked_provenance(obj) or value_contains_tracked_provenance(state): + discarded_tracked_storage_references = True invalidate_tensor_rebuild_proof() stack.append(obj if state is None else None) else: clear_stack_after_malformed_provenance() if not within_limits(): return _PytorchStorageReferenceParse(set(), {}, set(), set(), False, False) + except _PytorchZipDeadlineExceeded: + raise except Exception: return _PytorchStorageReferenceParse(set(), {}, set(), set(), False, False) return _PytorchStorageReferenceParse( @@ -2272,6 +2597,9 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any: ), parse_complete=True, all_persistent_ids_are_pytorch_storage=all_persistent_ids_are_pytorch_storage, + discarded_tracked_storage_references=discarded_tracked_storage_references, + used_streaming_batch_compaction=used_streaming_batch_compaction, + accepted_oversized_state_batch=accepted_oversized_state_batch, ) diff --git a/packages/modelaudit-picklescan/tests/test_api.py b/packages/modelaudit-picklescan/tests/test_api.py index 4d24cfb22..ba6855a3f 100644 --- a/packages/modelaudit-picklescan/tests/test_api.py +++ b/packages/modelaudit-picklescan/tests/test_api.py @@ -3404,10 +3404,40 @@ def test_scan_file_suppresses_rebuild_tensor_v2_for_multiple_canonical_memo_uses data_pkl = ( b"\x80\x04" + _global(b"torch._utils", b"_rebuild_tensor_v2") - + b"q\x00" + + b"q\x000" + + _pytorch_empty_ordered_dict_reduce_expr() + + b"(" + + _short_binunicode(b"first") + _pytorch_rebuild_tensor_v2_reduce_expr() + + _short_binunicode(b"second") + _pytorch_rebuild_tensor_v2_reduce_expr() - + b"." + + b"u." + ) + _write_pytorch_zip_data_pickle(archive_path, data_pkl) + + report = _scan_file_report_dict_subprocess(archive_path) + + assert report["status"] == "complete" + assert report["verdict"] == "clean" + assert _torch_rebuild_tensor_v2_warning_dicts(report) == [] + + +def test_scan_file_suppresses_rebuild_tensor_v2_for_nested_canonical_tensor(tmp_path: Path) -> None: + _require_torch_distribution() + archive_path = tmp_path / "nested-canonical-state.pt" + data_pkl = ( + b"\x80\x04" + + _global(b"torch._utils", b"_rebuild_tensor_v2") + + b"q\x000" + + _pytorch_empty_ordered_dict_reduce_expr() + + b"(" + + _short_binunicode(b"weight") + + _pytorch_rebuild_tensor_v2_reduce_expr() + + _short_binunicode(b"_extra_state") + + b"}" + + _short_binunicode(b"calibration") + + _pytorch_rebuild_tensor_v2_reduce_expr() + + b"su." ) _write_pytorch_zip_data_pickle(archive_path, data_pkl) @@ -3418,6 +3448,136 @@ def test_scan_file_suppresses_rebuild_tensor_v2_for_multiple_canonical_memo_uses assert _torch_rebuild_tensor_v2_warning_dicts(report) == [] +@pytest.mark.parametrize("memoized_final_tensor", [False, True]) +def test_scan_file_preserves_hidden_malicious_storage_after_stacked_canonical_tensors( + tmp_path: Path, + memoized_final_tensor: bool, +) -> None: + archive_path = tmp_path / f"malicious-stacked-canonical-{memoized_final_tensor}.pt" + hidden_payload = b"S'" + b"A" * 5000 + b"'\n0cos\nsystem\n(S'echo discarded-same-storage'\ntR." + canonical_tensor = _pytorch_rebuild_tensor_v2_reduce_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + final_tensor = b"q\x01h\x01" if memoized_final_tensor else canonical_tensor + data_pkl = ( + b"\x80\x04" + + _global(b"torch._utils", b"_rebuild_tensor_v2") + + b"q\x00" + + canonical_tensor + + final_tensor + + b"." + ) + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", data_pkl) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + + report = scan_file(archive_path) + + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == Severity.CRITICAL + and finding.details.get("module") in {"os", "posix", "nt"} + and finding.details.get("name") == "system" + for finding in report.findings + ) + + +@pytest.mark.parametrize("raw_storage_first", [False, True]) +def test_scan_file_preserves_hidden_malicious_storage_inside_enclosing_list( + tmp_path: Path, + raw_storage_first: bool, +) -> None: + archive_path = tmp_path / f"malicious-enclosing-list-{raw_storage_first}.pt" + hidden_payload = b"S'" + b"A" * 5000 + b"'\n0cos\nsystem\n(S'echo enclosing-list-storage'\ntR." + entries: list[bytes] = [] + for index in range(600): + key = _short_binunicode(f"weight_{index}".encode("ascii")) + if index == 0: + value = _pytorch_rebuild_tensor_v2_payload(key=str(index + 1)).removeprefix(b"\x80\x04").removesuffix(b".") + else: + value = _pytorch_rebuild_tensor_v2_reduce_expr(key=str(index + 1)) + entries.append(key + value) + raw_storage = _pytorch_storage_binpersid_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + canonical_dictionary = _pytorch_empty_ordered_dict_reduce_expr() + b"(" + b"".join(entries) + b"u" + ordered_entries = (raw_storage, canonical_dictionary) if raw_storage_first else (canonical_dictionary, raw_storage) + data_pkl = b"\x80\x04]" + b"".join(item + b"a" for item in ordered_entries) + b"." + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", data_pkl) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + for index in range(600): + archive.writestr(f"archive/data/{index + 1}", b"\x00" * 24) + + report = scan_file(archive_path) + + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == Severity.CRITICAL + and finding.details.get("module") in {"os", "posix", "nt"} + and finding.details.get("name") == "system" + for finding in report.findings + ) + + +@pytest.mark.parametrize("padding_entries", [2, 5]) +def test_scan_file_preserves_hidden_malicious_storage_after_padded_nested_dictionary_overwrite( + tmp_path: Path, + padding_entries: int, +) -> None: + archive_path = tmp_path / f"malicious-padded-nested-state-{padding_entries}.pt" + hidden_payload = b"S'" + b"A" * 5000 + b"'\n0cos\nsystem\n(S'echo nested-dictionary-storage'\ntR." + hidden_tensor = ( + _pytorch_rebuild_tensor_v2_payload( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + .removeprefix(b"\x80\x04") + .removesuffix(b".") + ) + padding = b"".join(_short_binunicode(f"pad_{index}".encode("ascii")) + b"K\x00" for index in range(padding_entries)) + nested_dictionary = b"}(" + padding + _short_binunicode(b"hidden") + hidden_tensor + b"u" + entries = [ + _short_binunicode(b"outer") + nested_dictionary, + _short_binunicode(b"outer") + b"K\x00", + _short_binunicode(b"safe") + _pytorch_rebuild_tensor_v2_reduce_expr(key="1"), + ] + entries.extend(_short_binunicode(f"metadata_{index}".encode("ascii")) + b"K\x01" for index in range(40)) + data_pkl = b"\x80\x04" + _pytorch_empty_ordered_dict_reduce_expr() + b"(" + b"".join(entries) + b"u." + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", data_pkl) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + archive.writestr("archive/data/1", b"\x00" * 24) + + report = scan_file(archive_path) + + assert ( + package_api._trusted_pytorch_data_pkl_from_storage_member_sizes( + data_pkl, + {"0": len(hidden_payload), "1": 24}, + ) + is None + ) + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == Severity.CRITICAL + and finding.details.get("module") in {"os", "posix", "nt"} + and finding.details.get("name") == "system" + for finding in report.findings + ) + + def test_scan_file_suppresses_rebuild_tensor_v2_for_real_ordered_state_dict(tmp_path: Path) -> None: _require_torch_distribution() import torch @@ -3432,6 +3592,909 @@ 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, + leading_metadata_count: int = 0, + leading_canonical_container: Literal["list", "tuple", "dictionary"] | None = None, + split_leading_metadata_batch: bool = False, + metadata_value: bytes = b"K\x01", + trailing_metadata_count: int = 0, + trailing_metadata_fields: int = 1, +) -> 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 + metadata_value) + if leading_canonical_container is not None: + nested_tensor = _pytorch_rebuild_tensor_v2_payload(key="0").removeprefix(b"\x80\x04").removesuffix(b".") + nested_value = { + "list": b"]" + nested_tensor + b"a", + "tuple": nested_tensor + b"\x85", + "dictionary": b"}" + _short_binunicode(b"calibration") + nested_tensor + b"s", + }[leading_canonical_container] + entries.append(_short_binunicode(b"_extra_state") + nested_value) + for index in range(leading_metadata_count): + metadata_key = f"_extra_state_{index}".encode("ascii") + entries.append(b"X" + len(metadata_key).to_bytes(4, "little") + metadata_key + metadata_value) + 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 and leading_canonical_container is None: + 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(" + if split_leading_metadata_batch: + prefix_count = leading_metadata_count + int(leading_metadata) + payload += b"".join(entries[:prefix_count]) + b"u(" + payload += b"".join(entries[prefix_count:]) + b"u" + else: + payload += b"".join(entries) + b"u" + if trailing_metadata_count: + metadata_key = b"_metadata" + metadata_entries = [] + for index in range(trailing_metadata_count): + module_name = f"layer_{index}".encode("ascii") + fields = [] + for field_index in range(trailing_metadata_fields): + version_key = b"version" if field_index == 0 else f"extra_{field_index}".encode("ascii") + fields.append(b"X" + len(version_key).to_bytes(4, "little") + version_key + b"K\x01") + version_dict = b"}(" + b"".join(fields) + b"u" + metadata_entries.append(b"X" + len(module_name).to_bytes(4, "little") + module_name + version_dict) + payload += b"(" + b"X" + len(metadata_key).to_bytes(4, "little") + metadata_key + b"}(" + payload += b"".join(metadata_entries) + b"uu" + payload += b"." + 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", "metadata_value"), + [(600, False, b"K\x01"), (1000, False, b"K\x01"), (600, True, b"K\x01"), (600, True, b"}")], +) +def test_pytorch_storage_trust_parses_large_batched_state_dict_without_framework( + tmp_path: Path, + entry_count: int, + leading_metadata: bool, + metadata_value: bytes, +) -> 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, + metadata_value=metadata_value, + ) + 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_pytorch_storage_trust_parses_large_metadata_after_canonical_batch(tmp_path: Path) -> None: + archive_path = tmp_path / "large-batched-metadata-state.pt" + _write_large_batched_pytorch_state_dict(archive_path, trailing_metadata_count=600) + 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) == 600 + assert len(parsed.canonical_tensor_rebuild_invocations) == 600 + + +@pytest.mark.parametrize("container_kind", ["list", "tuple", "dictionary"]) +def test_pytorch_storage_trust_parses_large_batch_with_leading_nested_canonical_tensor( + tmp_path: Path, + container_kind: Literal["list", "tuple", "dictionary"], +) -> None: + archive_path = tmp_path / f"large-batched-nested-canonical-{container_kind}.pt" + _write_large_batched_pytorch_state_dict( + archive_path, + leading_canonical_container=container_kind, + ) + 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) == 600 + assert len(parsed.canonical_tensor_rebuild_invocations) == 601 + + +@pytest.mark.parametrize("container_kind", ["list", "tuple", "dictionary"]) +def test_pytorch_storage_trust_parses_nested_canonical_tensor_in_later_state_batch( + tmp_path: Path, + container_kind: Literal["list", "tuple", "dictionary"], +) -> None: + archive_path = tmp_path / f"later-batch-nested-canonical-{container_kind}.pt" + _write_large_batched_pytorch_state_dict(archive_path, entry_count=1000) + with zipfile.ZipFile(archive_path) as archive: + first_batch = archive.read("archive/data.pkl") + nested_tensor = _pytorch_rebuild_tensor_v2_reduce_expr(key="1000") + nested_value = { + "list": b"]" + nested_tensor + b"a", + "tuple": nested_tensor + b"\x85", + "dictionary": b"}" + _short_binunicode(b"calibration") + nested_tensor + b"s", + }[container_kind] + later_entries = [_short_binunicode(b"_extra_state") + nested_value] + later_entries.extend( + _short_binunicode(f"weight_{index}".encode("ascii")) + _pytorch_rebuild_tensor_v2_reduce_expr(key=str(index)) + for index in range(1001, 1033) + ) + payload = first_batch.removesuffix(b".") + b"(" + b"".join(later_entries) + b"u." + + parsed = package_api._pytorch_storage_keys_from_pickle_bytes(payload) + + assert parsed.parse_complete is True + assert len(parsed.referenced_keys) == 1033 + assert len(parsed.canonical_tensor_rebuild_invocations) == 1033 + + +@pytest.mark.parametrize(("leading_metadata_count", "split_batch"), [(510, False), (1000, False), (1000, True)]) +def test_pytorch_storage_trust_parses_large_metadata_prefix_before_canonical_tensor( + tmp_path: Path, + leading_metadata_count: int, + split_batch: bool, +) -> None: + archive_path = tmp_path / f"large-metadata-prefix-{leading_metadata_count}.pt" + _write_large_batched_pytorch_state_dict( + archive_path, + entry_count=1, + leading_metadata_count=leading_metadata_count, + split_leading_metadata_batch=split_batch, + ) + 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 parsed.referenced_keys == {"0"} + assert len(parsed.canonical_tensor_rebuild_invocations) == 1 + + +def test_scan_file_preserves_malicious_call_after_large_metadata_prefix(tmp_path: Path) -> None: + archive_path = tmp_path / "large-metadata-prefix-malicious.pt" + _write_large_batched_pytorch_state_dict( + archive_path, + entry_count=2, + leading_metadata_count=510, + malicious=True, + ) + + 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_parses_nested_multifield_metadata_batch(tmp_path: Path) -> None: + archive_path = tmp_path / "large-multifield-metadata-state.pt" + _write_large_batched_pytorch_state_dict( + archive_path, + trailing_metadata_count=510, + trailing_metadata_fields=2, + ) + 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.canonical_tensor_rebuild_invocations) == 600 + + +def test_scan_file_suppresses_rebuild_tensor_v2_for_real_module_state_dict(tmp_path: Path) -> None: + _require_torch_distribution() + import torch + + archive_path = tmp_path / "module-state-with-metadata.pt" + torch.save(torch.nn.Linear(3, 2).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) == [] + + +def test_pytorch_storage_trust_checks_deadline_during_metadata_provenance( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + archive_path = tmp_path / "large-batched-metadata-deadline.pt" + _write_large_batched_pytorch_state_dict(archive_path, trailing_metadata_count=600) + with zipfile.ZipFile(archive_path) as archive: + payload = archive.read("archive/data.pkl") + + def stop_during_metadata_inspection(_deadline: float) -> None: + if sys._getframe(1).f_code.co_name == "value_contains_tracked_provenance": + raise package_api._PytorchZipDeadlineExceeded + + monkeypatch.setattr(package_api, "_check_pytorch_zip_deadline", stop_during_metadata_inspection) + + with pytest.raises(package_api._PytorchZipDeadlineExceeded): + package_api._pytorch_storage_keys_from_pickle_bytes(payload, deadline=1.0) + + +def test_pytorch_storage_trust_bounds_metadata_provenance_work( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + archive_path = tmp_path / "large-batched-metadata-budget.pt" + _write_large_batched_pytorch_state_dict(archive_path, trailing_metadata_count=600) + with zipfile.ZipFile(archive_path) as archive: + payload = archive.read("archive/data.pkl") + + monkeypatch.setattr(package_api, "_PYTORCH_STORAGE_TRUST_MAX_PROVENANCE_NODES", 32) + + 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_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(("entry_count", "padding_bytes"), [(33, 5000), (100, 5000), (600, 5000), (600, 32_000)]) +def test_scan_file_preserves_hidden_malicious_storage_inside_compacted_canonical_tensor( + tmp_path: Path, + entry_count: int, + padding_bytes: int, +) -> None: + archive_path = tmp_path / f"malicious-compacted-canonical-{entry_count}-{padding_bytes}.pt" + hidden_payload = b"S'" + b"A" * padding_bytes + b"'\n0cos\nsystem\n(S'echo compacted-storage'\ntR." + entries: list[bytes] = [] + for index in range(entry_count): + key = _short_binunicode(f"weight_{index}".encode("ascii")) + if index == 0: + value = ( + _pytorch_rebuild_tensor_v2_payload( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + .removeprefix(b"\x80\x04") + .removesuffix(b".") + ) + else: + value = _pytorch_rebuild_tensor_v2_reduce_expr(key=str(index)) + entries.append(key + value) + data_pkl = b"\x80\x04" + _pytorch_empty_ordered_dict_reduce_expr() + b"(" + b"".join(entries) + b"u." + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", data_pkl) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + for index in range(1, entry_count): + archive.writestr(f"archive/data/{index}", b"\x00" * 24) + + parsed = package_api._pytorch_storage_keys_from_pickle_bytes(data_pkl) + report = scan_file(archive_path) + + assert parsed.parse_complete is True + assert parsed.accepted_oversized_state_batch is True + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == Severity.CRITICAL + and finding.details.get("module") in {"os", "posix", "nt"} + and finding.details.get("name") == "system" + and finding.location is not None + and "archive/data/0" in finding.location + for finding in report.findings + ) + + +def test_scan_file_bounds_expanded_probes_for_benign_pickle_like_tensor_prefixes(tmp_path: Path) -> None: + archive_path = tmp_path / "large-pickle-like-storage-prefixes.pt" + entry_count = 600 + storage_size = 64 * 1024 + entries: list[bytes] = [] + for index in range(entry_count): + key = _short_binunicode(f"weight_{index}".encode("ascii")) + if index == 0: + value = ( + _pytorch_rebuild_tensor_v2_payload( + key="0", + storage_name="ByteStorage", + element_count=storage_size, + ) + .removeprefix(b"\x80\x04") + .removesuffix(b".") + ) + else: + value = _pytorch_rebuild_tensor_v2_reduce_expr( + key=str(index), + storage_name="ByteStorage", + element_count=storage_size, + ) + entries.append(key + value) + data_pkl = b"\x80\x04" + _pytorch_empty_ordered_dict_reduce_expr() + b"(" + b"".join(entries) + b"u." + storage = b"N" + b"\x00" * (storage_size - 1) + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", data_pkl) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + for index in range(entry_count): + archive.writestr(f"archive/data/{index}", storage) + + report = scan_file(archive_path) + + assert not any(notice.code == "pytorch_zip_pickle_discovery_probe_budget" for notice in report.notices) + try: + importlib_metadata.distribution("torch") + except importlib_metadata.PackageNotFoundError: + assert report.status == ScanStatus.INCONCLUSIVE + else: + assert report.status == ScanStatus.COMPLETE + + +@pytest.mark.parametrize("container_kind", ["dict-setitem", "dict-setitems", "list", "tuple", "ordered-dict"]) +def test_scan_file_preserves_retained_benign_storage_container( + tmp_path: Path, + container_kind: str, +) -> None: + archive_path = tmp_path / f"retained-storage-{container_kind}.pt" + storage = _pytorch_storage_binpersid_expr() + key = _short_binunicode(b"storage") + container = { + "dict-setitem": b"}" + key + storage + b"s", + "dict-setitems": b"}(" + key + storage + b"u", + "list": b"]" + storage + b"a", + "tuple": storage + b"\x85", + "ordered-dict": _pytorch_empty_ordered_dict_reduce_expr() + key + storage + b"s", + }[container_kind] + data_pkl = b"\x80\x04" + container + b"." + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", data_pkl) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", b"\x00" * 24) + + report = scan_file(archive_path) + + assert report.status == ScanStatus.COMPLETE + assert report.verdict == SafetyVerdict.CLEAN + assert not any(finding.rule_code == "PERSISTENT_ID" for finding in report.findings) + + +@pytest.mark.parametrize(("leading_metadata", "metadata_value"), [(False, b"K\x01"), (True, b"}")]) +def test_scan_file_preserves_malicious_call_in_large_batched_state_dict( + tmp_path: Path, + leading_metadata: bool, + metadata_value: bytes, +) -> 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, + metadata_value=metadata_value, + ) + + 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_preserves_noncanonical_setitems_width_limit() -> None: + entries = package_api._PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH // 2 + 1 + payload = ( + b"\x80\x04}(" + + b"".join( + b"X" + len(key := f"key_{index}".encode("ascii")).to_bytes(4, "little") + key + b"K\x00" + for index in range(entries) + ) + + b"u." + ) + + 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() + + +@pytest.mark.parametrize( + ("with_prior_tensor_batch", "storage_wrapper"), + [ + (False, None), + (True, None), + (True, (b"builtins", b"list")), + (True, (b"builtins", b"tuple")), + (True, (b"collections", b"OrderedDict")), + (True, "popped"), + (True, "pop_mark"), + (True, "overwrite_dict"), + (True, "overwrite_setitems"), + (True, "popped_tensor"), + (True, "pop_mark_tensor"), + (True, "overwrite_tensor_dict"), + (True, "overwrite_tensor_setitems"), + (True, "nested_tensor_tuple"), + ], +) +def test_scan_file_preserves_hidden_malicious_storage_after_noncanonical_setitems( + tmp_path: Path, + with_prior_tensor_batch: bool, + storage_wrapper: tuple[bytes, bytes] + | Literal[ + "popped", + "pop_mark", + "overwrite_dict", + "overwrite_setitems", + "popped_tensor", + "pop_mark_tensor", + "overwrite_tensor_dict", + "overwrite_tensor_setitems", + "nested_tensor_tuple", + ] + | None, +) -> None: + archive_path = tmp_path / f"malicious-noncanonical-state-{with_prior_tensor_batch}.pt" + hidden_payload = b"S'" + b"A" * 5000 + b"'\n0cos\nsystem\n(S'echo malicious-near-match'\ntR." + entries: list[bytes] = [] + for index in range(package_api._PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH // 2 + 1): + key = f"key_{index}".encode("ascii") + entries.append(b"X" + len(key).to_bytes(4, "little") + key + b"K\x00") + storage_key = b"storage" + storage_reference = _pytorch_storage_binpersid_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + if storage_wrapper in { + "popped_tensor", + "pop_mark_tensor", + "overwrite_tensor_dict", + "overwrite_tensor_setitems", + "nested_tensor_tuple", + }: + storage_reference = _pytorch_rebuild_tensor_v2_reduce_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + if storage_wrapper in {"popped", "popped_tensor"}: + storage_reference += b"0K\x00" + elif storage_wrapper in {"pop_mark", "pop_mark_tensor"}: + storage_reference = b"(" + storage_reference + b"1K\x00" + elif storage_wrapper in {"overwrite_dict", "overwrite_tensor_dict"}: + slot = _short_binunicode(b"slot") + storage_reference = b"}" + slot + storage_reference + b"s" + slot + b"K\x00s" + elif storage_wrapper in {"overwrite_setitems", "overwrite_tensor_setitems"}: + slot = _short_binunicode(b"slot") + storage_reference = b"}(" + slot + storage_reference + b"u(" + slot + b"K\x00u" + elif storage_wrapper == "nested_tensor_tuple": + storage_reference += b"\x85" + elif storage_wrapper is not None: + storage_reference = _global(*storage_wrapper) + b"(" + storage_reference + b"tR" + entries.append(b"X" + len(storage_key).to_bytes(4, "little") + storage_key + storage_reference) + if with_prior_tensor_batch: + tensor_key = b"weight" + tensor_value = _pytorch_rebuild_tensor_v2_payload(key="1").removeprefix(b"\x80\x04").removesuffix(b".") + prefix = ( + b"\x80\x04" + + _global(b"collections", b"OrderedDict") + + b")R(" + + b"X" + + len(tensor_key).to_bytes(4, "little") + + tensor_key + + tensor_value + + b"u(" + + b"X\x09\x00\x00\x00_metadata}(" + ) + suffix = b"uu." + else: + prefix = b"\x80\x04}(" + suffix = b"u." + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", prefix + b"".join(entries) + suffix) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + if with_prior_tensor_batch: + archive.writestr("archive/data/1", b"\x00" * 24) + + report = scan_file(archive_path) + + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == Severity.CRITICAL + and finding.details.get("module") in {"os", "posix", "nt"} + and finding.details.get("name") == "system" + for finding in report.findings + ) + + +@pytest.mark.parametrize( + "duplicate_variant", + [ + "ordered_batch", + "ordered_setitem", + "nested_dict", + "ordered_build", + "list_reduce", + "tuple_reduce", + "ordered_reduce", + "ordered_storage_key", + "ordered_storage_value", + "nested_storage_key", + "leading_nested_ordered_storage", + "leading_plain_dict_storage_batch", + "leading_plain_dict_storage_setitem", + ], +) +def test_scan_file_preserves_hidden_malicious_storage_after_duplicate_state_key( + tmp_path: Path, + duplicate_variant: str, +) -> None: + archive_path = tmp_path / f"malicious-duplicate-state-{duplicate_variant}.pt" + hidden_payload = b"S'" + b"A" * 5000 + b"'\n0cos\nsystem\n(S'echo duplicate-state-key'\ntR." + separate_storage = duplicate_variant in { + "ordered_build", + "list_reduce", + "tuple_reduce", + "ordered_reduce", + "ordered_storage_key", + "ordered_storage_value", + "nested_storage_key", + "leading_nested_ordered_storage", + "leading_plain_dict_storage_batch", + "leading_plain_dict_storage_setitem", + } + + def encoded_key(value: str) -> bytes: + raw = value.encode("ascii") + return b"X" + len(raw).to_bytes(4, "little") + raw + + tensor_value = ( + _pytorch_rebuild_tensor_v2_payload( + key="1" if separate_storage else "0", + storage_name="LongStorage" if separate_storage else "ByteStorage", + element_count=3 if separate_storage else len(hidden_payload), + ) + .removeprefix(b"\x80\x04") + .removesuffix(b".") + ) + entries = [encoded_key("same") + tensor_value] + entries.extend(encoded_key(f"metadata_{index}") + b"K\x01" for index in range(40)) + suffix = b"u." + if duplicate_variant == "ordered_batch": + entries.insert(1, encoded_key("same") + b"K\x00") + elif duplicate_variant == "ordered_setitem": + suffix = b"u" + encoded_key("same") + b"K\x00s." + elif duplicate_variant == "ordered_build": + storage_reference = _pytorch_storage_binpersid_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + suffix = b"u}(" + encoded_key("hidden") + storage_reference + b"ub." + elif duplicate_variant in {"list_reduce", "tuple_reduce", "ordered_reduce"}: + storage_reference = _pytorch_storage_binpersid_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + module_name = b"collections" if duplicate_variant == "ordered_reduce" else b"builtins" + callable_name = { + "list_reduce": b"list", + "tuple_reduce": b"tuple", + "ordered_reduce": b"OrderedDict", + }[duplicate_variant] + suffix = ( + b"u" + encoded_key("converted") + _global(module_name, callable_name) + b"(" + storage_reference + b"tRs." + ) + elif duplicate_variant in { + "ordered_storage_key", + "ordered_storage_value", + "nested_storage_key", + "leading_nested_ordered_storage", + "leading_plain_dict_storage_batch", + "leading_plain_dict_storage_setitem", + }: + storage_reference = _pytorch_storage_binpersid_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + if duplicate_variant == "ordered_storage_key": + suffix = b"u" + storage_reference + b"K\x00s." + elif duplicate_variant == "ordered_storage_value": + suffix = b"u" + encoded_key("hidden") + storage_reference + b"s." + elif duplicate_variant == "nested_storage_key": + suffix = b"u" + encoded_key("hidden") + b"}" + storage_reference + b"K\x00ss." + else: + storage_reference = _pytorch_storage_binpersid_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + nested_dict = b"(" + encoded_key("slot") + storage_reference + encoded_key("slot") + b"K\x00d" + entries.append(encoded_key("nested") + nested_dict) + + payload = b"\x80\x04" + _global(b"collections", b"OrderedDict") + b")R(" + b"".join(entries) + suffix + if duplicate_variant == "leading_nested_ordered_storage": + ordered_dict = _global(b"collections", b"OrderedDict") + b")R" + payload = ( + b"\x80\x04" + + ordered_dict + + encoded_key("pre") + + ordered_dict + + encoded_key("inside") + + storage_reference + + b"ss(" + + b"".join(entries) + + b"u." + ) + elif duplicate_variant in {"leading_plain_dict_storage_batch", "leading_plain_dict_storage_setitem"}: + hidden_entry = encoded_key("hidden") + storage_reference + initial_storage = b"(" + hidden_entry + b"u" if duplicate_variant.endswith("batch") else hidden_entry + b"s" + payload = b"\x80\x04}" + initial_storage + b"(" + b"".join(entries) + b"u." + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", payload) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + if separate_storage: + archive.writestr("archive/data/1", b"\x00" * 24) + + report = scan_file(archive_path) + + storage_sizes = {"0": len(hidden_payload)} + if separate_storage: + storage_sizes["1"] = 24 + assert package_api._trusted_pytorch_data_pkl_from_storage_member_sizes(payload, storage_sizes) is None + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == Severity.CRITICAL + and finding.details.get("module") in {"os", "posix", "nt"} + and finding.details.get("name") == "system" + for finding in report.findings + ) + + +@pytest.mark.parametrize( + "discard_variant", + [ + "stop_extra_storage", + "stop_extra_storage_none", + "stop_nested_storage", + "stop_extra_canonical_storage", + "unsupported_frozenset", + "unsupported_additems", + "append_wrong_target", + "setitem_storage_target", + "append_value", + "appends_value", + "setitem_value", + "setitem_key", + "stackglobal_name", + "stackglobal_module", + ], +) +def test_scan_file_preserves_hidden_malicious_storage_after_stack_discard( + tmp_path: Path, + discard_variant: str, +) -> None: + archive_path = tmp_path / f"malicious-stack-discard-{discard_variant}.pt" + hidden_payload = b"S'" + b"A" * 5000 + b"'\n0cos\nsystem\n(S'echo discarded-storage'\ntR." + + def encoded_key(value: str) -> bytes: + raw = value.encode("ascii") + return b"X" + len(raw).to_bytes(4, "little") + raw + + tensor = _pytorch_rebuild_tensor_v2_payload(key="1").removeprefix(b"\x80\x04").removesuffix(b".") + entries = [encoded_key("weight") + tensor] + entries.extend(encoded_key(f"metadata_{index}") + b"K\x01" for index in range(40)) + storage_reference = _pytorch_storage_binpersid_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + trailers = { + "stop_extra_storage": storage_reference + b"h\x01.", + "stop_extra_storage_none": storage_reference + b"N.", + "stop_nested_storage": b"]" + storage_reference + b"ah\x01.", + "stop_extra_canonical_storage": b".", + "unsupported_frozenset": b"(" + storage_reference + b"\x910h\x01.", + "unsupported_additems": b"\x8f(" + storage_reference + b"\x900h\x01.", + "append_wrong_target": storage_reference + b"K\x00ah\x01.", + "setitem_storage_target": storage_reference + b"K\x00K\x01sh\x01.", + "append_value": encoded_key("converted") + b"K\x00" + storage_reference + b"as.", + "appends_value": encoded_key("converted") + b"K\x00(" + storage_reference + b"es.", + "setitem_value": encoded_key("converted") + b"K\x00" + encoded_key("slot") + storage_reference + b"ss.", + "setitem_key": encoded_key("converted") + b"K\x00" + storage_reference + b"K\x00ss.", + "stackglobal_name": encoded_key("converted") + _short_binunicode(b"fake") + storage_reference + b"\x93s.", + "stackglobal_module": encoded_key("converted") + storage_reference + _short_binunicode(b"fake") + b"\x93s.", + } + payload = ( + b"\x80\x04" + + _global(b"collections", b"OrderedDict") + + b")Rq\x01(" + + b"".join(entries) + + b"u" + + trailers[discard_variant] + ) + if discard_variant == "stop_extra_canonical_storage": + payload = ( + b"\x80\x04" + + _global(b"torch._utils", b"_rebuild_tensor_v2") + + b"q\x00" + + _pytorch_rebuild_tensor_v2_reduce_expr( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + + _pytorch_rebuild_tensor_v2_reduce_expr(key="1") + + b"." + ) + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", payload) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + archive.writestr("archive/data/1", b"\x00" * 24) + + report = scan_file(archive_path) + + assert ( + package_api._trusted_pytorch_data_pkl_from_storage_member_sizes(payload, {"0": len(hidden_payload), "1": 24}) + is None + ) + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == Severity.CRITICAL + and finding.details.get("module") in {"os", "posix", "nt"} + and finding.details.get("name") == "system" + for finding in report.findings + ) + + +@pytest.mark.parametrize("discard_opcode", [b"\x91", b"\x82\x01", b"\x90"]) +def test_scan_file_preserves_hidden_malicious_storage_after_compacted_batch_discard( + tmp_path: Path, + discard_opcode: bytes, +) -> None: + archive_path = tmp_path / f"malicious-compacted-discard-{discard_opcode.hex()}.pt" + hidden_payload = b"S'" + b"A" * 5000 + b"'\n0cos\nsystem\n(S'echo compacted-storage'\ntR." + + def encoded_key(value: str) -> bytes: + raw = value.encode("ascii") + return b"X" + len(raw).to_bytes(4, "little") + raw + + tensor = ( + _pytorch_rebuild_tensor_v2_payload(key="0", storage_name="ByteStorage", element_count=len(hidden_payload)) + .removeprefix(b"\x80\x04") + .removesuffix(b".") + ) + entries = [encoded_key("weight") + tensor] + entries.extend(encoded_key(f"metadata_{index}") + b"K\x01" for index in range(600)) + payload = b"\x80\x04}q\x01(" + b"".join(entries) + discard_opcode + b"0h\x01." + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", payload) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + + report = scan_file(archive_path) + + assert package_api._trusted_pytorch_data_pkl_from_storage_member_sizes(payload, {"0": len(hidden_payload)}) is None + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == 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_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_invalidates_unfinished_compacted_batch(tmp_path: Path) -> None: + archive_path = tmp_path / "unfinished-canonical-batch.pt" + _write_large_batched_pytorch_state_dict(archive_path) + with zipfile.ZipFile(archive_path) as archive: + payload = archive.read("archive/data.pkl") + + parsed = package_api._pytorch_storage_keys_from_pickle_bytes(payload[:-2] + b".") + + assert parsed.parse_complete is False + assert parsed.referenced_keys == set() + assert parsed.canonical_tensor_rebuild_invocations == set() + + +def test_pytorch_storage_trust_clears_compacted_batch_after_malformed_provenance(tmp_path: Path) -> None: + archive_path = tmp_path / "malformed-canonical-batch.pt" + _write_large_batched_pytorch_state_dict(archive_path) + with zipfile.ZipFile(archive_path) as archive: + payload = archive.read("archive/data.pkl") + + parsed = package_api._pytorch_storage_keys_from_pickle_bytes(payload[:-2] + b"\x82\x01.") + + assert parsed.parse_complete is True + 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: @@ -14212,3 +15275,55 @@ def test_scan_stream_preserves_absolute_offsets_from_current_stream_position() - ] assert finding_positions assert all(position >= len(prefix) for position in finding_positions) + + +def test_scan_file_scans_binary_storage_member_inside_expanded_probe_window(tmp_path: Path) -> None: + """Widening the storage probe must not drop members the 4 KiB trusted probe scanned. + + ``sample_is_prefix`` is derived from how much of the member the probe read, so a member between + the two probe sizes flips it from True to False, and the binary-pickle predicate refuses to scan + a non-STOP-terminated sample once it is no longer a prefix. + """ + hidden_payload = b"\x80\x04cos\nsystem\n(S'echo expanded-window'\ntR" + hidden_payload += b"\x00" * (5000 - len(hidden_payload)) + assert len(hidden_payload) > package_api._TRUSTED_STORAGE_PICKLE_PROBE_BYTES + assert len(hidden_payload) < package_api._PICKLE_DISCOVERY_LONG_PROBE_BYTES + + entry_count = 33 + archive_path = tmp_path / "malicious-expanded-probe-window.pt" + entries: list[bytes] = [] + for index in range(entry_count): + key = _short_binunicode(f"weight_{index}".encode("ascii")) + if index == 0: + value = ( + _pytorch_rebuild_tensor_v2_payload( + key="0", + storage_name="ByteStorage", + element_count=len(hidden_payload), + ) + .removeprefix(b"\x80\x04") + .removesuffix(b".") + ) + else: + value = _pytorch_rebuild_tensor_v2_reduce_expr(key=str(index)) + entries.append(key + value) + data_pkl = b"\x80\x04" + _pytorch_empty_ordered_dict_reduce_expr() + b"(" + b"".join(entries) + b"u." + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("archive/data.pkl", data_pkl) + archive.writestr("archive/version", "3\n") + archive.writestr("archive/byteorder", "little") + archive.writestr("archive/data/0", hidden_payload) + for index in range(1, entry_count): + archive.writestr(f"archive/data/{index}", b"\x00" * 24) + + report = scan_file(archive_path) + + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.severity == Severity.CRITICAL + and finding.details.get("module") in {"os", "posix", "nt"} + and finding.details.get("name") == "system" + and finding.location is not None + and "archive/data/0" in finding.location + for finding in report.findings + )