Forbid inductor::_reinterpret_tensor (as_strided heap-OOB bypass)

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,61 @@
 
 #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 <tuple>
+#include <utility>
+#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;
+};
+
+//! True for the serialised TorchScript source records (code/**/*.py) that hold
+//! the executable code of the module and its submodules — including
+//! __setstate__.  Other records (data.pkl, constants.pkl, version, tensors) are
+//! not executable source and are skipped.
+bool isSerialisedCodeRecord(const std::string& name) {
+    constexpr std::string_view PY_EXT{".py"};
+    if (name.size() < PY_EXT.size() ||
+        name.compare(name.size() - PY_EXT.size(), PY_EXT.size(), PY_EXT) != 0) {
+        return false;
+    }
+    // Archive entries are typically "<root>/code/....py"; accept a leading
+    // "code/" too in case the root prefix is absent.
+    return name.find("/code/") != std::string::npos || name.rfind("code/", 0) == 0;
+}
+}
 
 CModelGraphValidator::SResult CModelGraphValidator::validate(const ::torch::jit::Module& module) {
 
@@ -95,6 +144,59 @@ void CModelGraphValidator::collectBlockOps(const ::torch::jit::Block& block,
     }
 }
 
+CModelGraphValidator::TStringVec
+CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::size_t size) {
+    TStringSet found;
+
+    // Build the textual signatures we are looking for in serialised TorchScript
+    // source.  Known stable forms:
+    //   aten::foo            -> torch.foo(
+    //   inductor::_reinterpret_tensor -> ops.inductor._reinterpret_tensor(
+    // prim::CallFunction / prim::CallMethod describe unresolved calls in an
+    // inlined graph and have no stable textual form in the serialised source.
+    std::vector<std::pair<std::string, std::string>> signatures; // (needle, qualified op)
+    constexpr std::string_view ATEN{"aten::"};
+    constexpr std::string_view INDUCTOR{"inductor::"};
+    for (const auto& op : CSupportedOperations::FORBIDDEN_OPERATIONS) {
+        if (op.size() > ATEN.size() && op.substr(0, ATEN.size()) == ATEN) {
+            std::string shortName{op.substr(ATEN.size())};
+            signatures.emplace_back("torch." + shortName + "(", std::string{op});
+        } else if (op.size() > INDUCTOR.size() && op.substr(0, INDUCTOR.size()) == INDUCTOR) {
+            std::string shortName{op.substr(INDUCTOR.size())};
+            signatures.emplace_back("ops.inductor." + shortName + "(", std::string{op});
+        }
+    }
+
+    try {
+        auto adapter = std::make_shared<CMemoryReadAdapter>(data, size);
+        caffe2::serialize::PyTorchStreamReader reader{adapter};
+
+        for (const auto& name : reader.getAllRecords()) {
+            if (isSerialisedCodeRecord(name) == false) {
+                continue;
+            }
+            auto[recordData, recordSize] = reader.getRecord(name);
+            std::string_view code{static_cast<const char*>(recordData.get()), recordSize};
+            for (const auto & [ needle, qualifiedOp ] : signatures) {
+                if (code.find(needle) != std::string_view::npos) {
+                    found.emplace(qualifiedOp);
+                }
+            }
+        }
+    } 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 model code scan skipped (could not parse archive): "
+                 << e.what());
+        return {};
+    }
+
+    TStringVec result{found.begin(), found.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,26 @@ class CModelGraphValidator {
                             const std::unordered_set<std::string_view>& allowedOps,
                             const std::unordered_set<std::string_view>& forbiddenOps);
 
+    //! Statically scan the *serialised* TorchScript code inside a model archive
+    //! for forbidden operations, WITHOUT loading the model.
+    //!
+    //! This closes a load-time execution gap: torch::jit::load runs a module's
+    //! __setstate__ during deserialization, before validate() (which inspects an
+    //! already-loaded module) could ever run.  A forbidden op hidden in
+    //! __setstate__ would therefore execute at load time and never appear in the
+    //! forward graph.  Scanning the serialised code first lets us reject such a
+    //! model before any of its code runs.
+    //!
+    //! \p data / \p size are the raw bytes of the .pt (ZIP) archive.  Returns the
+    //! sorted, de-duplicated qualified names of any forbidden operations found in
+    //! the serialised code.  Returns empty if none are found, or if the archive
+    //! cannot be parsed (in which case torch::jit::load will surface the error).
+    //!
+    //! This is a textual, defense-in-depth check targeted at the curated
+    //! forbidden set; the post-load graph validation in validate() remains the
+    //! authoritative allowlist check.
+    static TStringVec scanSerialisedCodeForForbiddenOps(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
+    // (HackerOne / elastic/security#12242).
+    "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,21 @@ void verifySafeModel(const torch::jit::script::Module& module_) {
         HANDLE_FATAL(<< "Model graph validation failed: " << e.what());
     }
 }
+
+//! Reject models that carry a forbidden operation in their serialised code
+//! BEFORE the model is loaded.  torch::jit::load executes a module's
+//! __setstate__ during deserialization, so a purely post-load check would run
+//! too late for an op hidden there.  \p modelData / \p modelSize are the raw
+//! bytes of the buffered .pt archive.
+void verifySafeModelBeforeLoad(const char* modelData, std::size_t modelSize) {
+    auto forbiddenOps = ml::torch::CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        modelData, modelSize);
+    if (forbiddenOps.empty() == false) {
+        std::string ops = ml::core::CStringUtils::join(forbiddenOps, ", ");
+        HANDLE_FATAL(<< "Model contains forbidden operations: " << ops
+                     << " (detected in serialised code before load)");
+    }
+}
 }
 
 torch::Tensor infer(torch::jit::script::Module& module_,
@@ -315,6 +330,12 @@ int main(int argc, char** argv) {
         if (readAdapter->init() == false) {
             return EXIT_FAILURE;
         }
+        if (skipModelValidation == false) {
+            // Scan the serialised code before loading so that a forbidden op
+            // hidden in __setstate__ (which torch::jit::load would execute during
+            // deserialization) is rejected before any of the model's code runs.
+            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,125 @@ BOOST_AUTO_TEST_CASE(testMaliciousRopExploit) {
     BOOST_REQUIRE(hasForbiddenOp(result, "aten::as_strided"));
 }
 
+// --- Pre-load __setstate__ scan tests ---
+//
+// torch::jit::load executes a module's __setstate__ during deserialization,
+// before CModelGraphValidator::validate (which inspects an already-loaded
+// module) could run.  A forbidden op hidden in __setstate__ therefore never
+// appears in the forward graph AND has already executed by load time.
+// scanSerialisedCodeForForbiddenOps closes this gap by inspecting the
+// serialised code textually, WITHOUT loading the model — so these tests never
+// call torch::jit::load on the attack fixtures.
+
+BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReader) {
+    // aten::from_file lives only in __setstate__; the forward graph is benign,
+    // so the post-load validator would miss it. The pre-load scan must catch it
+    // without loading (and thus without executing) the model.
+    std::string bytes = readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt");
+    BOOST_REQUIRE(bytes.empty() == false);
+
+    auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        bytes.data(), bytes.size());
+
+    BOOST_REQUIRE(scanContains(forbidden, "aten::from_file"));
+}
+
+BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReaderInSubmodule) {
+    // The forbidden op is hidden in a submodule's __setstate__; the scan must
+    // inspect every serialised code record, not just the top-level module's.
+    std::string bytes = readFileBytes(
+        "testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt");
+    BOOST_REQUIRE(bytes.empty() == false);
+
+    auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        bytes.data(), bytes.size());
+
+    BOOST_REQUIRE(scanContains(forbidden, "aten::from_file"));
+}
+
+BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsForwardForbiddenOps) {
+    // The scan is textual, so it also flags forbidden ops that appear in the
+    // forward graph (a superset of the load-time surface).
+    std::string fileReader = readFileBytes("testfiles/malicious_models/malicious_file_reader.pt");
+    BOOST_REQUIRE(fileReader.empty() == false);
+    auto forbiddenFileReader = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        fileReader.data(), fileReader.size());
+    BOOST_REQUIRE(scanContains(forbiddenFileReader, "aten::from_file"));
+
+    std::string heapLeak = readFileBytes("testfiles/malicious_models/malicious_heap_leak.pt");
+    BOOST_REQUIRE(heapLeak.empty() == false);
+    auto forbiddenHeapLeak = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        heapLeak.data(), heapLeak.size());
+    BOOST_REQUIRE(scanContains(forbiddenHeapLeak, "aten::as_strided"));
+}
+
+BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsBenignModel) {
+    // A legitimate model must not trip the scan (no false positives).
+    std::string bytes = readFileBytes("testfiles/e5_with_norm.pt");
+    BOOST_REQUIRE(bytes.empty() == false);
+
+    auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        bytes.data(), bytes.size());
+
+    BOOST_REQUIRE(forbidden.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 no forbidden ops.
+    std::string garbage = "this is not a valid .pt archive";
+    auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        garbage.data(), garbage.size());
+    BOOST_REQUIRE(forbidden.empty());
+}
+
+BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsReinterpretTensorOobRead) {
+    // Bypass of the as_strided forbid via inductor::_reinterpret_tensor
+    // (elastic/security#12242). Must be caught before load.
+    std::string bytes = readFileBytes(
+        "testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt");
+    BOOST_REQUIRE(bytes.empty() == false);
+
+    auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        bytes.data(), bytes.size());
+
+    BOOST_REQUIRE(scanContains(forbidden, "inductor::_reinterpret_tensor"));
+}
+
+BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsReinterpretTensorOobWrite) {
+    std::string bytes = readFileBytes(
+        "testfiles/malicious_models/malicious_reinterpret_tensor_oob_write.pt");
+    BOOST_REQUIRE(bytes.empty() == false);
+
+    auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        bytes.data(), bytes.size());
+
+    BOOST_REQUIRE(scanContains(forbidden, "inductor::_reinterpret_tensor"));
+}
+
+BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateReinterpretTensor) {
+    // Forbidden op only in __setstate__; must be rejected without loading.
+    std::string bytes = readFileBytes(
+        "testfiles/malicious_models/malicious_setstate_reinterpret_tensor.pt");
+    BOOST_REQUIRE(bytes.empty() == false);
+
+    auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps(
+        bytes.data(), bytes.size());
+
+    BOOST_REQUIRE(scanContains(forbidden, "inductor::_reinterpret_tensor"));
+}
+
+BOOST_AUTO_TEST_CASE(testMaliciousReinterpretTensorRejectedPostLoad) {
+    // Forward-only fixture: safe to load, then post-load validate must 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