fix(pickle): bound legacy PyTorch control streams - #1619
Conversation
Performance BenchmarksCompared Top improvements:
|
![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: 9affa2894d
ℹ️ 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 result.metadata.get("legacy_pytorch_container") is True: | ||
| return None |
There was a problem hiding this comment.
When legacy_pytorch_container is true but storage_key_count is 0, a valid legacy PyTorch file has no storage payload after pickle_end; appended ELF/PE/shell bytes should still hit the binary-tail detector. This early return skips that coverage, weakening detections. guidance
Useful? React with 👍 / 👎.
| if legacy_layout is not None: | ||
| result = self._scan_standalone_bytes(payload[: legacy_layout.pickle_end], source=source) | ||
| self._annotate_legacy_pytorch_layout(result, legacy_layout) |
There was a problem hiding this comment.
For non-seekable streams whose legacy control pickles fit in the buffer but tensor storage extends past max_known_stream_read_bytes, this branch still falls through to _add_stream_truncation_check, so the scan is marked inconclusive with non_seekable_stream_truncated even though only raw tensor storage was omitted. That keeps large legacy PyTorch streams failing the coverage check that this path is meant to avoid.
Useful? React with 👍 / 👎.
| result = self._scan_standalone_bytes(raw_data[: legacy_layout.pickle_end], source=source) | ||
| self._annotate_legacy_pytorch_layout(result, legacy_layout, position_offset=start_position) |
There was a problem hiding this comment.
When scan_stream is called on a seekable stream that is already positioned after a wrapper prefix, this fallback replaces the native scan_stream result with scan_bytes, which has no position_offset. Any malicious globals in the legacy control pickles are still detected, but their reported positions/import-reference metadata are shifted back to zero, unlike the normal stream path that reports offsets relative to the actual stream position.
Useful? React with 👍 / 👎.
…gacy-pytorch-stream-boundaries
|
@codex review Please review current head |
|
Codex Review: Didn't find any major issues. Swish! ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
|
@codex review Re-requesting review for current head |
|
Codex Review: Didn't find any major issues. Nice work! ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
1 participant
Summary
This fixes the campaign's
known_stream_truncated, binary-tail, 64-stream CVE coverage, and storage-byteEXT1/EXT2false positives as one container-boundary root cause.Tests
pytest tests/scanners/test_pickle_scanner.py -q(154 passed)ruff check modelaudit/scanners/pickle_scanner.py tests/scanners/test_pickle_scanner.pyruff format --check modelaudit/scanners/pickle_scanner.py tests/scanners/test_pickle_scanner.pymypy modelaudit/scanners/pickle_scanner.py tests/scanners/test_pickle_scanner.pygit diff --check