fix(picklescan): safely parse bounded PyTorch tensor batches - #1783
fix(picklescan): safely parse bounded PyTorch tensor batches#1783mldangelo-oai wants to merge 32 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR fixes a false-positive condition in the PyTorch ZIP “trusted storage reference” parsing path when state dictionaries are serialized using large SETITEMS batches (common in batched state dicts). It keeps the existing overall stack-depth bound while decoupling it from the smaller tuple-width limit, and adds regression coverage across both the standalone picklescan package and the root scanner integration.
Changes:
- Adjust the PyTorch storage trust parser to apply
SETITEMSbatching limits using the stack-depth bound rather than the tuple-width bound. - Add deterministic unit/regression tests for large batched state dicts (benign and malicious near-match) and explicit stack/tuple bound behaviors.
- Add changelog entries for both
modelauditandmodelaudit-picklescan.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/scanners/test_pytorch_zip_scanner.py | Adds an integration regression ensuring large batched torch.save state dicts remain clean (no _rebuild_tensor_v2 false positives). |
| packages/modelaudit-picklescan/tests/test_api.py | Adds deterministic ZIP fixtures + regressions for large batched SETITEMS, plus stack/tuple bound unit tests. |
| packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py | Updates marked-tuple parsing to accept a configurable width and uses stack-depth for SETITEMS batches. |
| packages/modelaudit-picklescan/CHANGELOG.md | Notes the bug fix in the picklescan package changelog. |
| CHANGELOG.md | Notes the bug fix in the root package changelog. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Performance BenchmarksCompared
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6ab180a928
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 117dc52da6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
mldangelo-oai
changed the title
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 165da9330a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3e5b2486fa
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py:2098
isinstance()cannot take a PEP 604 union (dict | _PytorchOrderedDictState) as its second argument; this will raiseTypeError: isinstance() argument 2 cannot be a unionat runtime on Python 3.10+ when the compaction path executes. Use a tuple of types instead.
if not isinstance(target, dict | _PytorchOrderedDictState):
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 46a8fb7a2d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py:2103
isinstance()cannot take a PEP 604 union (dict | _PytorchOrderedDictState) at runtime; this raisesTypeErrorand will cause the trust parser to bail out early (caught by the broadexcept Exception, returningparse_complete=False). Use a tuple of types instead.
continue
target = stack[marker_index - 1]
if not isinstance(target, dict | _PytorchOrderedDictState):
continue
if isinstance(target, _PytorchOrderedDictState) and target.used_as_hooks:
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d9f9103262
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if ( | ||
| len(setitem_pairs) * 2 > _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH | ||
| and not batch_has_canonical_tensor | ||
| and not trusted_canonical_batch_seen |
There was a problem hiding this comment.
When a valid state_dict contains one tensor and at least 33 module _metadata records, CPython serializes the tensor with SETITEM, not SETITEMS. Because only the SETITEMS branch sets trusted_canonical_batch_seen, this condition treats the later safe metadata batch as unproven and STOP returns parse_complete=False, restoring persistent-ID and _rebuild_tensor_v2 warnings. Fresh evidence beyond the earlier post-tensor report is that the current real-module regression has two tensor entries, while the one-entry OrderedDict regression strips the state dict's _metadata; have SETITEM recognize a canonical tensor and update the pending-batch state too.
AGENTS.md reference: AGENTS.md:L115-L117
Useful? React with 👍 / 👎.
| or stack[-1] is not canonical_batch_target | ||
| or (len(setitem_items) - 1) % 2 != 0 | ||
| ): | ||
| clear_stack_after_malformed_provenance() |
There was a problem hiding this comment.
When SETITEMS has odd arity or no target, pop_marked_tuple() has already removed its operands before this call clears the remaining stack, so a raw _PytorchStorageRef among those operands is never recorded as discarded. In a stream containing a newly accepted wide safe batch followed by this malformed operation and a later small canonical batch, the parser reports parse_complete=True with discarded_tracked_storage_references=False, and _trusted_pytorch_data_pkl_from_storage_member_sizes() consequently trusts the discarded storage key. A pickle hidden beyond the expanded 64 KiB member probe is then missed while its persistent-ID finding is suppressed; inspect setitem_items for tracked provenance before clearing it.
AGENTS.md reference: AGENTS.md:L115-L117
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9226edde21
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| except ValueError as error: | ||
| message = str(error).lower() | ||
| if "opcode" in message and "unknown" in message: | ||
| return False |
There was a problem hiding this comment.
When a 33+-entry canonical state batch marks a storage member for an expanded probe, a member beginning \x80\x04NNNN\x00 (a binary header, harmless opcodes, then malformed padding) is rejected here before the 64 KiB sample is read. The same member is accepted by _zip_entry_looks_like_pickle() after its valid opcode prefix, allowing the Rust follow-on scanner to inspect a dangerous pickle placed later in the member; this early return instead prevents the member from ever being added to pickle_entries, so that payload is skipped. Keep this optimization at least as conservative as the generic probe.
AGENTS.md reference: AGENTS.md:L115-L115
Useful? React with 👍 / 👎.
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| message = str(error).lower() | ||
| if "opcode" in message and "unknown" in message: | ||
| return False | ||
| except Exception: |
| 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 |
This restates ruff's current default select list, so it is a no-op today and a liability later: it pins the package to today's defaults, silently opting out of rules ruff adds to the default set. Lint passes without it, and it is unrelated to the batched-SETITEMS fix this PR makes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.
Suppressed comments (1)
packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py:2112
isinstance()cannot take a PEP 604 union (dict | _PytorchOrderedDictState) as its second argument; this raisesTypeErrorat runtime and will break the streaming compaction path for large batched state_dict pickles. Use a tuple of types instead.
continue
target = stack[marker_index - 1]
if not isinstance(target, dict | _PytorchOrderedDictState):
continue
Routing expanded-trust members through the 64 KiB probe instead of the 4 KiB trusted probe changed sample_is_prefix, which is derived from how much of the member was read. For members between the two probe sizes the flag flips True to False, and the binary-pickle predicate refuses to scan a non-STOP-terminated sample once it is no longer a prefix. A malicious pickle hidden in archive/data/N was therefore scanned before this PR and silently skipped after it. Reproduced end to end: a 5000-byte os.system payload in a 33-entry compacted state dict reports MALICIOUS with the fix and CLEAN without it. The genops guard added by this PR does not catch it either - the short probe is 16 bytes, so genops fails with a stringnl error whose message contains neither 'opcode' nor 'unknown' and the guard falls through. Re-checking the shorter view keeps a wider probe strictly additive. Applied to the proto-0/1 predicate as well, which derives sample_is_prefix the same way. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.
Suppressed comments (2)
packages/modelaudit-picklescan/tests/test_api.py:3431
- This bytes literal combines two pickle opcodes (BINPUT + POP) into a single string ("q\x000"), which is easy to misread as a single opcode/argument sequence. Splitting it into explicit opcode bytes makes the intended pickle stream clearer and reduces the chance of accidentally breaking the test payload.
+ b"q\x000"
packages/modelaudit-picklescan/tests/test_api.py:3407
- This bytes literal combines two pickle opcodes (BINPUT + POP) into a single string ("q\x000"), which is easy to misread as a single opcode/argument sequence. Splitting it into explicit opcode bytes makes the intended pickle stream clearer and reduces the chance of accidentally breaking the test payload.
This issue also appears on line 3431 of the same file.
+ b"q\x000"
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 21bc219377
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| 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 | ||
| ): |
There was a problem hiding this comment.
Fresh evidence beyond the previously reported malformed SETITEMS case is a LIST, TUPLE, or DICT without a matching marker: pop_marked_tuple() drains its operands before this helper inspects stack, so a raw _PytorchStorageRef among those operands is never marked discarded. After an accepted 33-entry tensor batch, a raw storage PID followed by LIST, NONE, and STOP produces parse_complete=True, discarded_tracked_storage_references=False, and exact trust for every storage key; a malicious pickle placed beyond that member's 64 KiB expanded probe is therefore skipped. Pass the popped operands into this cleanup or inspect them before draining the stack.
AGENTS.md reference: AGENTS.md:L115-L117
Useful? React with 👍 / 👎.
Summary
Verification