Harden model validation: reject __setstate__ hooks and forbid inductor::_reinterpret_tensor

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,60 @@ 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__"};
+
+    std::unique_ptr<caffe2::serialize::PyTorchStreamReader> reader;
+    try {
+        auto adapter = std::make_shared<CMemoryReadAdapter>(data, size);
+        reader = std::make_unique<caffe2::serialize::PyTorchStreamReader>(std::move(adapter));
+    } 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 {};
+    }
+
+    for (const auto& name : reader->getAllRecords()) {
+        try {
+            auto[recordData, recordSize] = reader->getRecord(name);
+            std::string_view bytes{static_cast<const char*>(recordData.get()), recordSize};
+
+            // Remediation for a privately reported finding: 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__");
+            }
+        } catch (const std::exception& e) {
+            // A single unreadable record (e.g. a deliberately bad CRC) must not
+            // abort the whole scan: an attacker could otherwise hide a
+            // __setstate__ hook in a later record behind a corrupt earlier one,
+            // slip past the scan, and have torch::jit::load run it at load time.
+            // Warn and keep scanning the remaining records.
+            LOG_WARN(<< "Pre-load state-hook scan: skipping unreadable record '"
+                     << name << "': " << e.what());
+            continue;
+        }
+        if (hooks.size() == 2) {
+            break;
+        }
+    }
+
+    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 a load-time execution gap from a privately reported finding:
+    //! torch::jit::load runs a module's __setstate__ during deserialization,
+    //! before post-load graph validation.  Matching the recommended 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/CSupportedOperations.cc

@@ -22,6 +22,10 @@ const CSupportedOperations::TStringViewSet CSupportedOperations::FORBIDDEN_OPERA
     "aten::as_strided"sv,
     "aten::from_file"sv,
     "aten::save"sv,
+    // Unchecked storage-offset reinterpret (TorchInductor). Same class of OOB
+    // heap read/write as as_strided; used to bypass the as_strided forbid
+    // (privately reported security finding).
+    "inductor::_reinterpret_tensor"sv,
     // After graph inlining, method and function calls should be resolved.
     // Their presence indicates an opaque call that cannot be validated.
     "prim::CallFunction"sv,

bin/pytorch_inference/Main.cc

@@ -73,6 +73,23 @@ 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 the recommended remediation for a
+//! privately reported finding, 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 +332,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,20 @@ 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);
+    // Fail loudly (naming the path) if a fixture goes missing, rather than
+    // returning an empty string that trips an opaque "bytes.empty() == false".
+    BOOST_REQUIRE_MESSAGE(file.is_open(), "Could not open fixture: " + path);
+    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 +453,84 @@ 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 a privately reported finding, 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 reported RCE uses allowlisted ops inside __setstate__; rejecting the
+    // hooks themselves is the recommended 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());
+}
+
+BOOST_AUTO_TEST_CASE(testMaliciousReinterpretTensorRejectedPostLoad) {
+    // inductor::_reinterpret_tensor is the as_strided heap-OOB bypass
+    // (privately reported finding).  This forward-only fixture carries no custom
+    // state hooks, so it loads cleanly; post-load validate must then report the
+    // op as forbidden (not merely unrecognised).
+    auto module = ::torch::jit::load(
+        "testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt");
+    auto result = CModelGraphValidator::validate(module);
+
+    BOOST_REQUIRE(result.s_IsValid == false);
+    BOOST_REQUIRE(hasForbiddenOp(result, "inductor::_reinterpret_tensor"));
+}
+
 // --- Prepacked model compatibility tests ---
 //
 // These load TorchScript models that mirror the ops used by Elasticsearch's