Reject TorchScript custom state hooks before load (__setstate__ gap)

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;
 

bin/pytorch_inference/CModelGraphValidator.cc

@@ -15,12 +15,44 @@
 
 #include <core/CLogger.h>
 
+#include <caffe2/serialize/inline_container.h>
+#include <caffe2/serialize/read_adapter_interface.h>
 #include <torch/csrc/jit/passes/inliner.h>
 
 #include <algorithm>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <vector>
 
 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<std::size_t>(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<CMemoryReadAdapter>(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<const char*>(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) {

bin/pytorch_inference/CModelGraphValidator.h

@@ -14,6 +14,7 @@
 
 #include <torch/script.h>
 
+#include <cstddef>
 #include <string>
 #include <string_view>
 #include <unordered_set>
@@ -71,6 +72,23 @@ class CModelGraphValidator {
                             const std::unordered_set<std::string_view>& allowedOps,
                             const std::unordered_set<std::string_view>& 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,

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. "

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

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,13 +294,32 @@ 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,
     "malicious_hidden_in_submodule.pt": HiddenInSubmodule,
     "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: