Reject TorchScript custom state hooks before load (__setstate__ gap)#3078
Draft
edsavage wants to merge 9 commits into
Draft
Reject TorchScript custom state hooks before load (__setstate__ gap)#3078edsavage wants to merge 9 commits into
edsavage wants to merge 9 commits into
Conversation
TorchScript runs a module's __setstate__ during torch::jit::load(), before the loaded module reaches CModelGraphValidator (which only walks the inlined forward graph). A forbidden op hidden in __setstate__ is therefore invisible to the validator and has already executed by the time validation would run. Add fixtures and a self-contained repro that demonstrate the gap: - generate_malicious_models.py: add SetStateFileReaderModel and a submodule variant (aten::from_file in __setstate__, benign forward); surface __setstate__ ops in the generator output. - test/test_setstate_load_timing.py: torch-only, build-free repro proving load-time execution (benign print probe) and op-hiding (static inspection). - test_pytorch_inference_evil_models.py: register the fixtures as regression guards for the forthcoming pre-load scan. No fix yet; this only reproduces and documents the vulnerability. Co-authored-by: Cursor <cursoragent@cursor.com>
Close the load-time execution gap: torch::jit::load runs a module's __setstate__ during deserialization, before CModelGraphValidator (which walks an already-loaded module's graph) can run. A forbidden op hidden in __setstate__ therefore executes at load time and never appears in the forward graph the validator inspects. Add CModelGraphValidator::scanSerialisedCodeForForbiddenOps, which uses a PyTorchStreamReader over the buffered archive bytes to textually scan the serialised TorchScript code (code/**/*.py, including every submodule's __setstate__) for the forbidden aten ops, WITHOUT loading the model. Main.cc runs this scan before torch::jit::load and rejects via HANDLE_FATAL, so no model code executes. Expose the buffered bytes via CBufferedIStreamAdapter::buffer(). This is a defense-in-depth textual check targeted at the curated forbidden set (from_file, as_strided, save); the post-load allowlist validation is unchanged, and seccomp remains the syscall backstop. Unit tests cover both setstate fixtures, forward-graph attacks, a benign model (no false positive), and malformed input; the binary integration test now expects the setstate models to be rejected with the forbidden-operations message. Co-authored-by: Cursor <cursoragent@cursor.com>
edsavage
changed the title
|
Hi @edsavage, I've created a changelog YAML for you. |
There was a problem hiding this comment.
Pull request overview
This PR closes a security-relevant validation gap in bin/pytorch_inference: TorchScript can execute a module’s __setstate__ during torch::jit::load(), so forbidden ops hidden there could run before the existing post-load graph validation. The fix adds a pre-load static scan of the serialized TorchScript source within the .pt archive (including __setstate__) and expands test coverage to reproduce and prevent regressions.
Changes:
- Add
CModelGraphValidator::scanSerialisedCodeForForbiddenOps()to scancode/**/*.pyinside a serialized.ptarchive for forbidden aten ops without loading the model. - Invoke the pre-load scan in
pytorch_inferencebeforetorch::jit::load()(when model validation is enabled). - Add/extend Python + C++ tests and model-generation tooling to reproduce the gap and assert rejection of
__setstate__-hidden forbidden ops.
Reviewed changes
Copilot reviewed 4 out of 11 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| test/test_setstate_load_timing.py | New standalone repro script demonstrating __setstate__ executes during load and is invisible to forward-graph-only validation. |
| test/test_pytorch_inference_evil_models.py | Adds __setstate__-based “evil models” and updates expectations to require clean rejection. |
| docs/changelog/3078.yaml | Changelog entry for the security/validation bug fix. |
| dev-tools/generate_malicious_models.py | Adds __setstate__ attack fixtures and surfaces __setstate__ ops during fixture generation. |
| bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc | Adds unit tests for the new pre-load scan (setstate fixtures, benign model, garbage input). |
| bin/pytorch_inference/Main.cc | Runs the pre-load serialized-code scan before calling torch::jit::load(). |
| bin/pytorch_inference/CModelGraphValidator.h | Declares the new serialized-code scanning API and documents rationale/behavior. |
| bin/pytorch_inference/CModelGraphValidator.cc | Implements in-memory archive parsing via PyTorchStreamReader and textual scanning of code/**/*.py. |
| bin/pytorch_inference/CBufferedIStreamAdapter.h | Exposes a read-only buffer() view so the already-buffered archive can be scanned pre-load. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+137
to
+146
| proc = subprocess.run( | ||
| [sys.executable, "-c", loader], | ||
| capture_output=True, text=True, timeout=120, | ||
| ) | ||
| stdout = proc.stdout | ||
| marker_before_return = ( | ||
| _LOAD_MARKER in stdout | ||
| and stdout.index(_LOAD_MARKER) < stdout.index("LOAD_RETURNED") | ||
| ) if "LOAD_RETURNED" in stdout else (_LOAD_MARKER in stdout) | ||
|
|
Contributor
Author
There was a problem hiding this comment.
Good catch — tightened in 5957acd. The benign probe now only counts as proof when the subprocess exits 0 and LOAD_RETURNED was printed, with the marker strictly before it; on failure it now prints the returncode and a stderr tail. So a crash that happened to print the marker first is no longer a false positive.
…l-cpp into security/setstate-load-timing
Address Copilot review: the benign probe must load cleanly for the marker to prove load-time execution. Require returncode == 0 and that LOAD_RETURNED was printed, and print returncode + stderr tail on failure, so a crash that happens to print the marker first is no longer a false positive. Co-authored-by: Cursor <cursoragent@cursor.com>
This was referenced Jul 20, 2026
Match the #12621 reporter remediation by refusing archives that embed __setstate__/__getstate__, and keep the forbidden-op pre-scan as defense in depth. Allowlisted-op RCE gadgets in setstate are otherwise missed. Co-authored-by: Cursor <cursoragent@cursor.com>
edsavage
changed the title
Contributor
Author
|
Updated after comparing with the reporter's |
Custom state-hook rejection is enough for the load-time execution gap; forward-graph ops remain covered by post-load validate(). Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
torch::jit::load/__setstate__timing gap from elastic/security#12621 (HackerOne).__setstate__or__getstate__before load.Test plan
e5_with_norm.ptaccepted; garbage input does not crash