From 0241ba04cae22582cadb23af601c8f781051c72e Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 11 Jun 2026 12:12:56 -0700 Subject: [PATCH 1/2] Fix arbitrary memory read --- .../core/framework/tensorprotoutils.cc | 9 +++ onnxruntime/test/ir/graph_test.cc | 68 +++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 8685655420a38..db07d3497b06b 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -2125,6 +2125,15 @@ common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& n switch (constant_attribute.type()) { case AttributeProto_AttributeType_TENSOR: + // Defense-in-depth: reject ORT in-memory address markers on a Constant node's dense tensor + // attribute. These markers are an internal ORT sentinel for trusted in-memory buffers and must + // never appear in a deserialized protobuf. The Graph constructor also checks all initializers + // (including those converted from Constant nodes), but we block early here to prevent the + // crafted tensor from propagating further. + ORT_RETURN_IF(HasExternalDataInMemory(constant_attribute.t()), + "Constant node '", node.name(), + "' tensor attribute references an ORT in-memory address marker, " + "which is not allowed in a model protobuf."); tensor = constant_attribute.t(); break; case AttributeProto_AttributeType_FLOAT: diff --git a/onnxruntime/test/ir/graph_test.cc b/onnxruntime/test/ir/graph_test.cc index 019f15a46abc5..7c64d5d5ed84d 100644 --- a/onnxruntime/test/ir/graph_test.cc +++ b/onnxruntime/test/ir/graph_test.cc @@ -1482,6 +1482,74 @@ TEST_F(GraphTest, RejectInMemoryMarkerOnDenseInitializer) { } } +// Regression test: a Constant node with a dense tensor attribute carrying an ORT in-memory address +// marker must be rejected during model load. This is the attack vector described in the MSRC report +// where an attacker crafts a Constant node to make ORT dereference an attacker-supplied pointer. +TEST_F(GraphTest, RejectInMemoryMarkerOnConstantNodeTensorAttribute) { + Model model("RejectInMemoryMarkerOnConstantNode", false, *logger_); + auto model_proto = model.ToProto(); + auto* m_graph = model_proto.mutable_graph(); + + // Build a minimal graph: Constant -> Identity -> output + static std::vector backing(16, 0); + + auto* const_node = m_graph->add_node(); + const_node->set_op_type("Constant"); + const_node->set_name("malicious_constant"); + const_node->add_output("c"); + + auto* attr = const_node->add_attribute(); + attr->set_name("value"); + attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + auto* t = attr->mutable_t(); + t->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + t->add_dims(4); + t->set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + auto* loc = t->add_external_data(); + loc->set_key("location"); + loc->set_value(ToUTF8String(onnxruntime::utils::kTensorProtoLittleEndianMemoryAddressTag)); + auto* off = t->add_external_data(); + off->set_key("offset"); + off->set_value(std::to_string(reinterpret_cast(backing.data()))); + auto* len = t->add_external_data(); + len->set_key("length"); + len->set_value(std::to_string(backing.size())); + + auto* identity_node = m_graph->add_node(); + identity_node->set_op_type("Identity"); + identity_node->set_name("identity"); + identity_node->add_input("c"); + identity_node->add_output("output"); + + auto* output = m_graph->add_output(); + output->set_name("output"); + auto* type_proto = output->mutable_type()->mutable_tensor_type(); + type_proto->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + type_proto->mutable_shape()->add_dim()->set_dim_value(4); + + std::string serialized; + model_proto.SerializeToString(&serialized); + + ModelProto model_proto_roundtrip; + ASSERT_TRUE(model_proto_roundtrip.ParseFromString(serialized)); + + std::shared_ptr p_tmp_model; + ORT_TRY { + auto status = onnxruntime::Model::Load(model_proto_roundtrip, p_tmp_model, nullptr, *logger_); + EXPECT_FALSE(status.IsOK()) << "Loading a model with an in-memory marker on a Constant node tensor attribute must fail."; + if (!status.IsOK()) { + EXPECT_THAT(status.ErrorMessage(), + ::testing::HasSubstr("in-memory address marker")); + } + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + EXPECT_THAT(std::string(ex.what()), + ::testing::HasSubstr("in-memory address marker")); + }); + } +} + TEST_F(GraphTest, GraphConstruction_CheckIsNotAcyclic) { // A cyclic graph // SouceNode From a064a9be1e0da16c9ae02f8735c4c4eb2f468dc3 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 26 Jun 2026 08:46:58 -0700 Subject: [PATCH 2/2] Add targeted test isolating the ConstantNodeProtoToTensorProto guard --- .../test/framework/tensorutils_test.cc | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/onnxruntime/test/framework/tensorutils_test.cc b/onnxruntime/test/framework/tensorutils_test.cc index 38a0f15f2301d..e394a62e04635 100644 --- a/onnxruntime/test/framework/tensorutils_test.cc +++ b/onnxruntime/test/framework/tensorutils_test.cc @@ -1378,6 +1378,46 @@ TEST(SparseTensorProtoToDenseTensorProtoMarkerTest, RejectsInMemoryMarkerOnIndic #endif // !defined(DISABLE_SPARSE_TENSORS) +// Defense-in-depth: ConstantNodeProtoToTensorProto must reject ORT's in-memory address marker +// on a Constant node's dense tensor attribute. This isolates the guard added in +// ConstantNodeProtoToTensorProto from the pre-existing dense-initializer guard in the Graph +// constructor: callers such as Graph::AddConstantProtoAsInitializer and the ORT-format build +// path emplace directly into name_to_initial_tensor_ and bypass that constructor-side check, +// so this test exercises the new chokepoint directly. +TEST(ConstantNodeProtoToTensorProtoMarkerTest, RejectsInMemoryMarkerOnDenseTensorAttribute) { + ONNX_NAMESPACE::NodeProto node; + node.set_op_type("Constant"); + node.set_name("malicious_constant"); + node.add_output("c"); + + auto* attr = node.add_attribute(); + attr->set_name("value"); + attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + auto* t = attr->mutable_t(); + t->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + t->add_dims(4); + t->set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + + // Backing buffer is irrelevant — the guard must reject before any dereference. + static std::vector backing(16, 0); + + auto* loc = t->add_external_data(); + loc->set_key("location"); + loc->set_value(ToUTF8String(onnxruntime::utils::kTensorProtoLittleEndianMemoryAddressTag)); + auto* off = t->add_external_data(); + off->set_key("offset"); + off->set_value(std::to_string(reinterpret_cast(backing.data()))); + auto* len = t->add_external_data(); + len->set_key("length"); + len->set_value(std::to_string(backing.size())); + + ONNX_NAMESPACE::TensorProto tensor_out; + Status status = utils::ConstantNodeProtoToTensorProto(node, std::filesystem::path{}, tensor_out); + ASSERT_FALSE(status.IsOK()) + << "Constant node tensor attribute with an in-memory address marker must be rejected."; + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("in-memory address marker")); +} + // Defense-in-depth: GetExtDataFromTensorProto must reject absolute external paths even when // called with an empty model_path (e.g. from training checkpoint or custom-op init paths). // Previously, ValidateExternalDataPath was only invoked from Graph::ConvertInitializersIntoOrtValues,