fix(picklescan): safely parse bounded PyTorch tensor batches - #1783
fix(picklescan): safely parse bounded PyTorch tensor batches#1783mldangelo-oai wants to merge 30 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: 2bdd951c24
ℹ️ 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 |
Summary
Verification