diff --git a/bin/pytorch_inference/CBufferedIStreamAdapter.h b/bin/pytorch_inference/CBufferedIStreamAdapter.h index d1f845e75..2f9bf9b3b 100644 --- a/bin/pytorch_inference/CBufferedIStreamAdapter.h +++ b/bin/pytorch_inference/CBufferedIStreamAdapter.h @@ -51,6 +51,11 @@ class CBufferedIStreamAdapter : public caffe2::serialize::ReadAdapterInterface { std::size_t size() const override; std::size_t read(std::uint64_t pos, void* buf, std::size_t n, const char* what = "") const override; + //! Read-only view of the buffered model bytes. Valid until this object is + //! destroyed. Used to statically scan the archive before it is handed to + //! torch::jit::load (which would otherwise execute __setstate__ code). + const char* buffer() const { return m_Buffer.get(); } + CBufferedIStreamAdapter(const CBufferedIStreamAdapter&) = delete; CBufferedIStreamAdapter& operator=(const CBufferedIStreamAdapter&) = delete; diff --git a/bin/pytorch_inference/CModelGraphValidator.cc b/bin/pytorch_inference/CModelGraphValidator.cc index 01658b440..c9b2bd21e 100644 --- a/bin/pytorch_inference/CModelGraphValidator.cc +++ b/bin/pytorch_inference/CModelGraphValidator.cc @@ -15,12 +15,44 @@ #include +#include +#include #include #include +#include +#include +#include +#include namespace ml { namespace torch { +namespace { + +//! Minimal in-memory ReadAdapterInterface so a PyTorchStreamReader can parse an +//! archive that is already fully buffered in memory, without taking ownership. +class CMemoryReadAdapter final : public caffe2::serialize::ReadAdapterInterface { +public: + CMemoryReadAdapter(const char* data, std::size_t size) + : m_Data{data}, m_Size{size} {} + + std::size_t size() const override { return m_Size; } + + std::size_t + read(std::uint64_t pos, void* buf, std::size_t n, const char* /*what*/ = "") const override { + if (pos >= m_Size) { + return 0; + } + n = std::min(n, m_Size - static_cast(pos)); + std::memcpy(buf, m_Data + pos, n); + return n; + } + +private: + const char* m_Data; + std::size_t m_Size; +}; +} CModelGraphValidator::SResult CModelGraphValidator::validate(const ::torch::jit::Module& module) { @@ -95,6 +127,48 @@ void CModelGraphValidator::collectBlockOps(const ::torch::jit::Block& block, } } +CModelGraphValidator::TStringVec +CModelGraphValidator::scanArchiveForCustomStateHooks(const char* data, std::size_t size) { + TStringSet hooks; + + constexpr std::string_view SETSTATE{"__setstate__"}; + constexpr std::string_view GETSTATE{"__getstate__"}; + + try { + auto adapter = std::make_shared(data, size); + caffe2::serialize::PyTorchStreamReader reader{adapter}; + + for (const auto& name : reader.getAllRecords()) { + auto[recordData, recordSize] = reader.getRecord(name); + std::string_view bytes{static_cast(recordData.get()), recordSize}; + + // Reporter remediation (elastic/security#12621): reject any archive + // that embeds custom state hooks. Scan every record — not just + // code/*.py — so hooks cannot hide in debug_pkl / adjacent entries. + if (bytes.find(SETSTATE) != std::string_view::npos) { + hooks.emplace("__setstate__"); + } + if (bytes.find(GETSTATE) != std::string_view::npos) { + hooks.emplace("__getstate__"); + } + if (hooks.size() == 2) { + break; + } + } + } catch (const std::exception& e) { + // If the archive cannot be parsed as a stream we do not treat this as a + // rejection here: torch::jit::load will attempt the same parse and + // surface a clear error. The post-load graph validator still applies. + LOG_WARN(<< "Pre-load state-hook scan skipped (could not parse archive): " + << e.what()); + return {}; + } + + TStringVec result{hooks.begin(), hooks.end()}; + std::sort(result.begin(), result.end()); + return result; +} + void CModelGraphValidator::collectModuleOps(const ::torch::jit::Module& module, TStringSet& ops, std::size_t& nodeCount) { diff --git a/bin/pytorch_inference/CModelGraphValidator.h b/bin/pytorch_inference/CModelGraphValidator.h index 2c589dab5..198800af5 100644 --- a/bin/pytorch_inference/CModelGraphValidator.h +++ b/bin/pytorch_inference/CModelGraphValidator.h @@ -14,6 +14,7 @@ #include +#include #include #include #include @@ -71,6 +72,23 @@ class CModelGraphValidator { const std::unordered_set& allowedOps, const std::unordered_set& forbiddenOps); + //! Statically scan a model archive BEFORE torch::jit::load for custom + //! TorchScript state hooks ("__setstate__", "__getstate__"), without + //! executing any of the model's code. + //! + //! Closes the load-time execution gap exploited in elastic/security#12621: + //! torch::jit::load runs a module's __setstate__ during deserialization, + //! before post-load graph validation. Matching the reporter's remediation, + //! any archive containing those literal substrings is rejected. Ops that + //! only run when methods are invoked (e.g. forward) remain the job of the + //! post-load allowlist / forbid checks in validate(). + //! + //! \p data / \p size are the raw bytes of the .pt (ZIP) archive. Returns + //! the sorted names of any hooks found, or empty if none are found / the + //! archive cannot be parsed (in which case torch::jit::load will surface + //! the error). + static TStringVec scanArchiveForCustomStateHooks(const char* data, std::size_t size); + private: //! Collect all operation names from a block, recursing into sub-blocks. static void collectBlockOps(const ::torch::jit::Block& block, diff --git a/bin/pytorch_inference/Main.cc b/bin/pytorch_inference/Main.cc index 4aaca0049..76c57d610 100644 --- a/bin/pytorch_inference/Main.cc +++ b/bin/pytorch_inference/Main.cc @@ -73,6 +73,22 @@ void verifySafeModel(const torch::jit::script::Module& module_) { HANDLE_FATAL(<< "Model graph validation failed: " << e.what()); } } + +//! Reject models with custom TorchScript state hooks BEFORE torch::jit::load. +//! Load executes __setstate__ during deserialization, so post-load graph +//! validation runs too late. Matching elastic/security#12621 / the reporter's +//! remediation, any __setstate__/__getstate__ hooks are refused outright. +//! Forbidden / unrecognised ops in methods that only run when invoked remain +//! the job of verifySafeModel() after a successful load. +//! \p modelData / \p modelSize are the raw bytes of the buffered .pt archive. +void verifySafeModelBeforeLoad(const char* modelData, std::size_t modelSize) { + auto hooks = ml::torch::CModelGraphValidator::scanArchiveForCustomStateHooks( + modelData, modelSize); + if (hooks.empty() == false) { + std::string names = ml::core::CStringUtils::join(hooks, ", "); + HANDLE_FATAL(<< "Model archive contains custom state hooks: " << names); + } +} } torch::Tensor infer(torch::jit::script::Module& module_, @@ -315,6 +331,12 @@ int main(int argc, char** argv) { if (readAdapter->init() == false) { return EXIT_FAILURE; } + if (skipModelValidation == false) { + // Reject custom state hooks before loading so that __setstate__ + // (which torch::jit::load would execute during deserialization) + // never runs. Op allowlisting remains a post-load check. + verifySafeModelBeforeLoad(readAdapter->buffer(), readAdapter->size()); + } module_ = torch::jit::load(std::move(readAdapter)); if (skipModelValidation) { LOG_WARN(<< "Model graph validation SKIPPED — --skipModelValidation flag is set. " diff --git a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc index 9c4f7d6d5..0204123d3 100644 --- a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc +++ b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc @@ -338,6 +338,17 @@ bool hasUnrecognisedOp(const CModelGraphValidator::SResult& result, const std::s return std::find(result.s_UnrecognisedOps.begin(), result.s_UnrecognisedOps.end(), op) != result.s_UnrecognisedOps.end(); } + +std::string readFileBytes(const std::string& path) { + std::ifstream file(path, std::ios::binary); + std::ostringstream buffer; + buffer << file.rdbuf(); + return buffer.str(); +} + +bool scanContains(const CModelGraphValidator::TStringVec& ops, const std::string& op) { + return std::find(ops.begin(), ops.end(), op) != ops.end(); +} } BOOST_AUTO_TEST_CASE(testMaliciousFileReader) { @@ -439,6 +450,71 @@ BOOST_AUTO_TEST_CASE(testMaliciousRopExploit) { BOOST_REQUIRE(hasForbiddenOp(result, "aten::as_strided")); } +// --- Pre-load custom state-hook scan tests --- +// +// torch::jit::load executes a module's __setstate__ during deserialization, +// before CModelGraphValidator::validate (which inspects an already-loaded +// module) could run. Matching elastic/security#12621, custom state hooks are +// rejected outright before load. Ops that only run when methods are invoked +// remain covered by the post-load allowlist / forbid checks. These tests +// never call torch::jit::load on attack fixtures. + +BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooks) { + // The HackerOne #12621 RCE uses allowlisted ops inside __setstate__; + // rejecting the hooks themselves is the reporter's remediation. + std::string bytes = readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(hooks, "__setstate__")); + BOOST_REQUIRE(scanContains(hooks, "__getstate__")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooksInSubmodule) { + std::string bytes = readFileBytes( + "testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(hooks, "__setstate__")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsBenignModel) { + // Typical scripted transformers do not embed custom __setstate__/__getstate__. + std::string bytes = readFileBytes("testfiles/e5_with_norm.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(hooks.empty()); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsForwardOnlyMaliciousModel) { + // Forbidden ops in forward (no custom state hooks) are not a pre-load + // concern — they are caught by post-load validate() after jit::load. + std::string bytes = readFileBytes("testfiles/malicious_models/malicious_file_reader.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(hooks.empty()); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanHandlesGarbageInput) { + // Non-archive input must not throw or crash; torch::jit::load will surface + // the real error later. The scan returns empty results. + std::string garbage = "this is not a valid .pt archive"; + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + garbage.data(), garbage.size()); + BOOST_REQUIRE(hooks.empty()); +} + // --- Prepacked model compatibility tests --- // // These load TorchScript models that mirror the ops used by Elasticsearch's diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt new file mode 100644 index 000000000..7a60d4f79 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt new file mode 100644 index 000000000..18552abc1 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt differ diff --git a/dev-tools/generate_malicious_models.py b/dev-tools/generate_malicious_models.py index fdf6ddbf1..1beaf4449 100644 --- a/dev-tools/generate_malicious_models.py +++ b/dev-tools/generate_malicious_models.py @@ -26,7 +26,7 @@ import torch from torch import Tensor -from typing import Optional +from typing import Optional, Tuple # --- Malicious model definitions --- @@ -110,6 +110,78 @@ def forward(self, x: Tensor) -> Tensor: return torch.from_file("/tmp/secret", size=10) +# --- Load-time execution models (__setstate__) --- +# +# TorchScript serialises a module's __setstate__ method and runs it *during* +# torch::jit::load(), before the loaded Module is ever handed to +# CModelGraphValidator. Because the validator only walks the (inlined) forward +# graph, a forbidden op hidden in __setstate__ is invisible to it AND has +# already executed by the time validation would run. These fixtures reproduce +# that gap: their forward graphs are benign, so the current validator accepts +# them, yet loading them triggers aten::from_file at deserialization time. + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__ — executes during load. + + The forbidden op never appears in forward(), so the post-load graph + validator (which inspects only the forward graph) does not see it. Worse, + torch::jit::load() runs __setstate__ before validation happens at all, so + the file read has already fired. This is the load-time execution gap + identified in security review. + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + # Runs at deserialization time, before any graph validation. + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderInSubmodule(torch.nn.Module): + """Same load-time attack, but hidden in a submodule's __setstate__. + + Confirms that a pre-load scan must recurse into every serialised + __setstate__ in the archive, not just the top-level module's. + """ + def __init__(self) -> None: + super().__init__() + self.child = _SetStateFileReaderChild() + + def forward(self, x: Tensor) -> Tensor: + return x + self.child(x) + + +class _SetStateFileReaderChild(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/tmp/secret", size=10) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + # --- Sandbox2 attack models (PR #2873) --- # # These reproduce real-world attack vectors that exploit torch.as_strided @@ -222,6 +294,23 @@ def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: # --- Generation logic --- +def _collect_setstate_ops(scripted: torch.jit.ScriptModule) -> set: + """Collect ops from every __setstate__ in the module tree. + + __setstate__ runs during torch::jit::load(), so ops here execute before + (and are invisible to) the forward-graph validator. Surfacing them makes + the load-time attack fixtures self-documenting. + """ + ops: set = set() + for submodule in scripted.modules(): + try: + graph = submodule._c._get_method("__setstate__").graph + except (RuntimeError, AttributeError): + continue + ops.update(node.kind() for node in graph.nodes()) + return ops + + MODELS = { "malicious_file_reader.pt": FileReaderModel, "malicious_mixed_file_reader.pt": MixedFileReaderModel, @@ -229,6 +318,8 @@ def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: "malicious_conditional.pt": ConditionalMalicious, "malicious_many_unrecognised.pt": ManyUnrecognisedOps, "malicious_file_reader_in_submodule.pt": FileReaderInSubmodule, + "malicious_setstate_file_reader.pt": SetStateFileReaderModel, + "malicious_setstate_file_reader_in_submodule.pt": SetStateFileReaderInSubmodule, "malicious_heap_leak.pt": HeapLeakModel, "malicious_rop_exploit.pt": RopExploitModel, } @@ -254,7 +345,12 @@ def generate(output_dir: Path): graph = scripted.forward.graph.copy() torch._C._jit_pass_inline(graph) ops = sorted(set(n.kind() for n in graph.nodes())) - print(f" ops: {ops}") + print(f" forward ops: {ops}") + + # Surface ops hidden in __setstate__ (executed at load time). + setstate_ops = _collect_setstate_ops(scripted) + if setstate_ops: + print(f" __setstate__ ops (run at load): {sorted(setstate_ops)}") succeeded.append(filename) except Exception as exc: diff --git a/docs/changelog/3078.yaml b/docs/changelog/3078.yaml new file mode 100644 index 000000000..f2ef914ac --- /dev/null +++ b/docs/changelog/3078.yaml @@ -0,0 +1,5 @@ +area: Machine Learning +issues: [] +pr: 3078 +summary: Reject TorchScript custom state hooks before load +type: bug diff --git a/test/test_pytorch_inference_evil_models.py b/test/test_pytorch_inference_evil_models.py index 64b81e2d9..442d70a75 100644 --- a/test/test_pytorch_inference_evil_models.py +++ b/test/test_pytorch_inference_evil_models.py @@ -37,6 +37,7 @@ import torch from torch import Tensor +from typing import Tuple # --------------------------------------------------------------------------- @@ -156,6 +157,68 @@ def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: return torch.zeros(0) +# --------------------------------------------------------------------------- +# Load-time execution models (__setstate__) +# --------------------------------------------------------------------------- +# +# TorchScript runs a module's __setstate__ during torch::jit::load(), before the +# loaded Module reaches CModelGraphValidator (which only walks the forward +# graph). A forbidden op in __setstate__ is therefore invisible to the current +# validator AND has already executed by load time. Pre-fix, feeding these models +# to the binary crashes/errors at load with no validator rejection message +# (INCONCLUSIVE below). Once a pre-load static scan is added, they must be +# cleanly REJECTED like the forward-graph attacks. + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__; forward() is benign.""" + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class _SetStateFileReaderChild(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/tmp/secret", size=10) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderInSubmodule(torch.nn.Module): + """Load-time aten::from_file hidden in a submodule's __setstate__.""" + def __init__(self) -> None: + super().__init__() + self.child = _SetStateFileReaderChild() + + def forward(self, x: Tensor) -> Tensor: + return x + self.child(x) + + # --------------------------------------------------------------------------- # Binary discovery # --------------------------------------------------------------------------- @@ -242,6 +305,18 @@ def find_pytorch_inference() -> str: "description": "ROP-chain file-write via aten::as_strided", "expect_stderr_contains": "Unrecognised operations", }, + "setstate_file_reader": { + "class": SetStateFileReaderModel, + "expect_rejected": True, + "description": "aten::from_file in __setstate__ (runs at load time)", + "expect_stderr_contains": "forbidden operations", + }, + "setstate_file_reader_submodule": { + "class": SetStateFileReaderInSubmodule, + "expect_rejected": True, + "description": "aten::from_file in a submodule's __setstate__", + "expect_stderr_contains": "forbidden operations", + }, } diff --git a/test/test_setstate_load_timing.py b/test/test_setstate_load_timing.py new file mode 100644 index 000000000..a2e6daeca --- /dev/null +++ b/test/test_setstate_load_timing.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python3 +# +# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +# or more contributor license agreements. Licensed under the Elastic License +# 2.0 and the following additional limitation. Functionality enabled by the +# files subject to the Elastic License 2.0 may only be used in production when +# invoked by an Elasticsearch process with a license key installed that permits +# use of machine learning features. You may not use this file except in +# compliance with the Elastic License 2.0 and the foregoing additional +# limitation. +# +"""Reproduce the __setstate__ load-time execution gap in model validation. + +TorchScript serialises a module's __setstate__ method and executes it *during* +torch::jit::load(). CModelGraphValidator, however, only inspects the (inlined) +forward graph of an already-loaded module. A forbidden operation hidden in +__setstate__ is therefore: + + 1. invisible to the validator (it never appears in the forward graph), and + 2. already executed by the time the validator would run (load happens first). + +This is the timing gap raised in security review: a malicious model can read +files / crash the process at load time, before any graph check. This test +demonstrates both facts *without needing a built binary* and only depends on +torch, so it can run in CI as a guard. + +It is deliberately safe: the load-time execution proof uses a benign probe +whose __setstate__ merely print()s a marker. The forbidden-op proofs use +static graph inspection on the saved archive and never load an attack model. + +Usage: + python3 test_setstate_load_timing.py +""" + +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Tuple + +import torch +from torch import Tensor + + +# Mirror of CSupportedOperations::FORBIDDEN_OPERATIONS (subset relevant here). +FORBIDDEN_OPERATIONS = { + "aten::from_file", + "aten::save", + "aten::as_strided", +} + +_LOAD_MARKER = "SETSTATE_EXECUTED_DURING_LOAD" + + +class BenignSetStateProbe(torch.nn.Module): + """Harmless probe: __setstate__ only print()s a marker (TorchScript-safe). + + Used to prove — without running any dangerous op — that __setstate__ code + executes during torch.jit.load(). + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + self.w = state[0] + self.training = state[1] + # Literal must be inlined: TorchScript cannot close over globals. + print("SETSTATE_EXECUTED_DURING_LOAD") + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__; forward() is benign. + + This is the attack shape: the forbidden op runs at load time and never + appears in the forward graph the validator inspects. We only ever inspect + this model statically (never load it), so no file is read here. + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +def _forward_ops(scripted: torch.jit.ScriptModule) -> set: + graph = scripted.forward.graph.copy() + torch._C._jit_pass_inline(graph) + return {node.kind() for node in graph.nodes()} + + +def _setstate_ops(scripted: torch.jit.ScriptModule) -> set: + ops: set = set() + for submodule in scripted.modules(): + try: + graph = submodule._c._get_method("__setstate__").graph + except (RuntimeError, AttributeError): + continue + ops.update(node.kind() for node in graph.nodes()) + return ops + + +def test_setstate_executes_during_load() -> bool: + """Prove that __setstate__ code runs during torch.jit.load().""" + print("--- 1. __setstate__ executes during load (benign probe) ---") + ok = True + with tempfile.TemporaryDirectory(prefix="setstate_timing_") as tmp: + path = Path(tmp) / "probe.pt" + scripted = torch.jit.script(BenignSetStateProbe()) + torch.jit.save(scripted, str(path)) + + # Load in a fresh subprocess; capture what is printed *during* load. + loader = ( + "import torch;" + f"torch.jit.load(r'{path}');" + "print('LOAD_RETURNED')" + ) + proc = subprocess.run( + [sys.executable, "-c", loader], + capture_output=True, text=True, timeout=120, + ) + stdout = proc.stdout + + # The benign probe must load *cleanly*: require a zero exit and that + # load() actually returned (LOAD_RETURNED). Otherwise a crash that + # happened to print the marker first would be a false positive — so we + # only accept the marker as proof when the whole load completed and the + # marker was printed strictly before load() returned. + loaded_cleanly = proc.returncode == 0 and "LOAD_RETURNED" in stdout + marker_before_return = ( + loaded_cleanly + and _LOAD_MARKER in stdout + and stdout.index(_LOAD_MARKER) < stdout.index("LOAD_RETURNED") + ) + + if marker_before_return: + print(f" OK: '{_LOAD_MARKER}' printed during load — " + "code ran before load() returned (and before validation).") + else: + print(" FAIL: marker not observed during a clean load.") + print(f" returncode: {proc.returncode}") + print(f" stdout: {stdout!r}") + if proc.stderr.strip(): + tail = "\n".join(proc.stderr.strip().splitlines()[-5:]) + print(f" stderr (tail):\n{tail}") + ok = False + print() + return ok + + +def test_forbidden_op_hidden_from_forward_graph() -> bool: + """The validator inspects only forward; the forbidden op lives in setstate.""" + print("--- 2. forbidden op is invisible to the forward-graph validator ---") + ok = True + scripted = torch.jit.script(SetStateFileReaderModel().eval()) + + forward = _forward_ops(scripted) + setstate = _setstate_ops(scripted) + + forward_forbidden = forward & FORBIDDEN_OPERATIONS + setstate_forbidden = setstate & FORBIDDEN_OPERATIONS + + print(f" forward ops : {sorted(forward)}") + print(f" __setstate__ ops: {sorted(setstate)}") + + if forward_forbidden: + print(f" FAIL: forbidden op unexpectedly in forward: {sorted(forward_forbidden)}") + ok = False + else: + print(" OK: forward graph contains no forbidden op — " + "the current validator would ACCEPT this model.") + + if "aten::from_file" in setstate_forbidden: + print(" OK: aten::from_file present in __setstate__ — " + "a correct pre-load scan must catch this.") + else: + print(" FAIL: aten::from_file not found in __setstate__ (fixture broken).") + ok = False + print() + return ok + + +def run_tests() -> bool: + print("=" * 72) + print("Repro: __setstate__ load-time execution gap (security review)") + print("=" * 72) + print(f"torch version: {torch.__version__}") + print() + + results = [ + test_setstate_executes_during_load(), + test_forbidden_op_hidden_from_forward_graph(), + ] + + all_passed = all(results) + print("=" * 72) + if all_passed: + print("ALL CHECKS PASSED — the load-time gap is reproduced and documented.") + print("A fix must refuse archives containing __setstate__/__getstate__") + print("BEFORE torch::jit::load() is called.") + else: + print("SOME CHECKS FAILED — see above.") + print("=" * 72) + return all_passed + + +if __name__ == "__main__": + sys.exit(0 if run_tests() else 1)