From 0560e904ea1b4bda6033ab5f59e1fbd661851540 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 22 Apr 2026 11:21:16 -0700 Subject: [PATCH 1/5] Harden ORT FlatBuffer model loader against malformed buffers - Add ValidateRequiredTableOffsets() to reject null table offsets in FlatBuffer vectors before dereferencing - Validate initializer raw_data byte size matches shape/type expectation in LoadInitializerOrtFormat - Compute required node slot count from actual node/edge indices instead of trusting the attacker-controlled max_node_index field - Reject dangling NodeEdge references to missing nodes - Reject duplicate NodeArg names and duplicate node indices - Fail on unknown NodeArg references in graph inputs/outputs - Add regression tests for malformed .ort buffers: raw_data size mismatch, null node arg entry, dangling node edge --- .../core/flatbuffers/flatbuffers_utils.cc | 1 + .../core/flatbuffers/flatbuffers_utils.h | 18 ++++ onnxruntime/core/graph/graph.cc | 52 +++++++++-- .../core/graph/graph_flatbuffers_utils.cc | 9 ++ onnxruntime/core/graph/model.cc | 1 + .../test/framework/ort_model_only_test.cc | 89 +++++++++++++++++++ 6 files changed, 164 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/flatbuffers/flatbuffers_utils.cc b/onnxruntime/core/flatbuffers/flatbuffers_utils.cc index 42dff12eaa2db..876c533073a96 100644 --- a/onnxruntime/core/flatbuffers/flatbuffers_utils.cc +++ b/onnxruntime/core/flatbuffers/flatbuffers_utils.cc @@ -282,6 +282,7 @@ Status LoadValueInfoOrtFormat(const fbs::ValueInfo& fbs_value_info, Status LoadOpsetImportOrtFormat(const flatbuffers::Vector>* fbs_op_set_ids, std::unordered_map& domain_to_version) { ORT_RETURN_IF(nullptr == fbs_op_set_ids, "Model must have opset imports. Invalid ORT format model."); + ORT_RETURN_IF_ERROR(ValidateRequiredTableOffsets(fbs_op_set_ids, "opset import")); domain_to_version.clear(); domain_to_version.reserve(fbs_op_set_ids->size()); diff --git a/onnxruntime/core/flatbuffers/flatbuffers_utils.h b/onnxruntime/core/flatbuffers/flatbuffers_utils.h index aed0c201a2dd5..b0fbc34e4fa26 100644 --- a/onnxruntime/core/flatbuffers/flatbuffers_utils.h +++ b/onnxruntime/core/flatbuffers/flatbuffers_utils.h @@ -50,6 +50,24 @@ onnxruntime::common::Status LoadOpsetImportOrtFormat( const flatbuffers::Vector>* fbs_op_set_ids, std::unordered_map& domain_to_version); +template +inline onnxruntime::common::Status ValidateRequiredTableOffsets( + const flatbuffers::Vector>* fbs_entries, + const char* entry_description) { + if (fbs_entries == nullptr) { + return onnxruntime::common::Status::OK(); + } + + const auto* raw_offsets = reinterpret_cast(fbs_entries->Data()); + for (flatbuffers::uoffset_t i = 0; i < fbs_entries->size(); ++i) { + const auto entry_offset = + flatbuffers::ReadScalar(raw_offsets + i * sizeof(flatbuffers::uoffset_t)); + ORT_RETURN_IF(entry_offset == 0, "Null ", entry_description, " entry. Invalid ORT format model."); + } + + return onnxruntime::common::Status::OK(); +} + // check if filename ends in .ort bool IsOrtFormatModel(const PathString& filename); diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 1f05dfd38968a..ff400f8a23f81 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -6520,8 +6521,10 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph // Initializers auto fbs_initializers = fbs_graph.initializers(); + ORT_RETURN_IF_ERROR(fbs::utils::ValidateRequiredTableOffsets(fbs_initializers, "initializer")); #if !defined(DISABLE_SPARSE_TENSORS) auto fbs_sparse_initializers = fbs_graph.sparse_initializers(); + ORT_RETURN_IF_ERROR(fbs::utils::ValidateRequiredTableOffsets(fbs_sparse_initializers, "sparse initializer")); flatbuffers::uoffset_t map_size = (fbs_initializers != nullptr ? fbs_initializers->size() : 0U) + (fbs_sparse_initializers != nullptr ? fbs_sparse_initializers->size() : 0U); #else @@ -6591,6 +6594,7 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph // NodeArgs auto fbs_node_args = fbs_graph.node_args(); if (fbs_node_args) { + ORT_RETURN_IF_ERROR(fbs::utils::ValidateRequiredTableOffsets(fbs_node_args, "node arg")); node_args_.reserve(fbs_node_args->size()); for (const auto* fbs_value_info : *fbs_node_args) { ORT_RETURN_IF(nullptr == fbs_value_info, "NodeArg is missing. Invalid ORT format model."); @@ -6598,7 +6602,8 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph ORT_RETURN_IF_ERROR(fbs::utils::LoadValueInfoOrtFormat(*fbs_value_info, node_arg_info)); const auto* name = fbs_value_info->name(); ORT_RETURN_IF(name == nullptr, "NodeArg name is missing. Invalid ORT format model."); - node_args_[name->str()] = std::make_unique(std::move(node_arg_info)); + const auto inserted = node_args_.emplace(name->str(), std::make_unique(std::move(node_arg_info))); + ORT_RETURN_IF(!inserted.second, "Duplicate NodeArg name '", name->str(), "'. Invalid ORT format model."); } } @@ -6606,8 +6611,37 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph // // Since we access a node using its index, we need to have nodes_ with size max_node_index to avoid // out of bounds access. - nodes_.resize(fbs_graph.max_node_index()); auto* fbs_nodes = fbs_graph.nodes(); + ORT_RETURN_IF_ERROR(fbs::utils::ValidateRequiredTableOffsets(fbs_nodes, "node")); + auto* fbs_node_edges = fbs_graph.node_edges(); + ORT_RETURN_IF_ERROR(fbs::utils::ValidateRequiredTableOffsets(fbs_node_edges, "node edge")); + + size_t required_node_slot_count = 0; + const auto update_required_node_slot_count = [&required_node_slot_count](uint32_t node_index) -> Status { + ORT_RETURN_IF(node_index == std::numeric_limits::max(), + "Node index is out of range. Invalid ORT format model."); + const auto node_slot_count = static_cast(node_index) + 1U; + required_node_slot_count = std::max(required_node_slot_count, node_slot_count); + return Status::OK(); + }; + + if (fbs_nodes != nullptr) { + for (const auto* fbs_node : *fbs_nodes) { + ORT_RETURN_IF(nullptr == fbs_node, "Node is missing. Invalid ORT format model."); + ORT_RETURN_IF_ERROR(update_required_node_slot_count(fbs_node->index())); + } + } + + if (fbs_node_edges != nullptr) { + for (const auto* fbs_node_edge : *fbs_node_edges) { + ORT_RETURN_IF(nullptr == fbs_node_edge, "NodeEdge is missing. Invalid ORT format model."); + ORT_RETURN_IF_ERROR(update_required_node_slot_count(fbs_node_edge->node_index())); + } + } + + ORT_RETURN_IF(fbs_graph.max_node_index() < required_node_slot_count, + "Serialized max node index is smaller than the required node slot count. Invalid ORT format model."); + nodes_.resize(required_node_slot_count); // It is possible to have no nodes in the model. Most likely scenario is the subgraph of an If Node // where the subgraph returns a Constant node. The Constant node will be lifted to an initializer by ORT @@ -6617,18 +6651,22 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph ORT_RETURN_IF(nullptr == fbs_node, "Node is missing. Invalid ORT format model."); std::unique_ptr node; ORT_RETURN_IF_ERROR(Node::LoadFromOrtFormat(*fbs_node, *this, load_options, logger_, node)); - ORT_RETURN_IF(node->Index() >= fbs_graph.max_node_index(), "Node index is out of range"); + ORT_RETURN_IF(node->Index() >= nodes_.size(), "Node index is out of range"); + ORT_RETURN_IF(nodes_[node->Index()] != nullptr, + "Duplicate node index ", node->Index(), ". Invalid ORT format model."); nodes_[node->Index()] = std::move(node); ++num_of_nodes_; } } // NodeEdges - auto* fbs_node_edges = fbs_graph.node_edges(); if (fbs_node_edges != nullptr) { for (const auto* fbs_node_edge : *fbs_node_edges) { ORT_RETURN_IF(nullptr == fbs_node_edge, "NodeEdge is missing. Invalid ORT format model."); - ORT_RETURN_IF(fbs_node_edge->node_index() >= fbs_graph.max_node_index(), "Node index is out of range"); + ORT_RETURN_IF(fbs_node_edge->node_index() >= nodes_.size(), "Node index is out of range"); + ORT_RETURN_IF(nodes_[fbs_node_edge->node_index()] == nullptr, + "NodeEdge references missing node ", fbs_node_edge->node_index(), + ". Invalid ORT format model."); ORT_RETURN_IF_ERROR(nodes_[fbs_node_edge->node_index()]->LoadEdgesFromOrtFormat(*fbs_node_edge, *this)); } } @@ -6640,7 +6678,9 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph node_args.reserve(fbs_node_args->size()); for (const auto* fbs_node_arg_name : *fbs_node_args) { ORT_RETURN_IF(nullptr == fbs_node_arg_name, "NodeArg Name is missing. Invalid ORT format model."); - gsl::not_null node_arg = GetNodeArg(fbs_node_arg_name->str()); + auto* node_arg = GetNodeArg(fbs_node_arg_name->str()); + ORT_RETURN_IF(node_arg == nullptr, "Graph references unknown NodeArg '", fbs_node_arg_name->str(), + "'. Invalid ORT format model."); node_args.push_back(node_arg); } } diff --git a/onnxruntime/core/graph/graph_flatbuffers_utils.cc b/onnxruntime/core/graph/graph_flatbuffers_utils.cc index c51f24229f145..fbe2b24ac3ca3 100644 --- a/onnxruntime/core/graph/graph_flatbuffers_utils.cc +++ b/onnxruntime/core/graph/graph_flatbuffers_utils.cc @@ -372,6 +372,15 @@ Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor, TensorProto& init } else { const auto* fbs_raw_data = fbs_tensor.raw_data(); if (fbs_raw_data) { + size_t expected_num_bytes = 0; + ORT_RETURN_IF_ERROR(GetSizeInBytesFromFbsTensor(fbs_tensor, expected_num_bytes)); + const char* tensor_name = initializer.has_name() ? initializer.name().c_str() : ""; + ORT_RETURN_IF( + fbs_raw_data->size() != expected_num_bytes, + "Initializer raw data size mismatch for tensor '", tensor_name, + "'. Expected ", expected_num_bytes, " bytes but found ", fbs_raw_data->size(), + ". Invalid ORT format model."); + if (load_options.can_use_flatbuffer_for_initializers && fbs_raw_data->size() > 127) { static_assert(sizeof(void*) <= sizeof(ExternalDataInfo::OFFSET_TYPE)); const void* data_offset = fbs_raw_data->Data(); diff --git a/onnxruntime/core/graph/model.cc b/onnxruntime/core/graph/model.cc index 59878052e7499..b51b66f05322f 100644 --- a/onnxruntime/core/graph/model.cc +++ b/onnxruntime/core/graph/model.cc @@ -983,6 +983,7 @@ common::Status Model::LoadFromOrtFormat(const fbs::Model& fbs_model, // Load the model metadata if (const auto* fbs_metadata_props = fbs_model.metadata_props()) { + ORT_RETURN_IF_ERROR(fbs::utils::ValidateRequiredTableOffsets(fbs_metadata_props, "metadata property")); model->model_metadata_.reserve(fbs_metadata_props->size()); for (const auto* prop : *fbs_metadata_props) { ORT_RETURN_IF(nullptr == prop, "Null entry in metadata_props. Invalid ORT format model."); diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index ec4f8967fd2a3..9a6d8d23df53b 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "core/flatbuffers/schema/ort.fbs.h" #include "core/framework/data_types.h" #include "core/framework/tensorprotoutils.h" @@ -41,6 +43,45 @@ struct OrtModelTestInfo { TransformerLevel optimization_level = TransformerLevel::Level3; }; +namespace { + +flatbuffers::Offset CreateFloatTensorTypeInfo(flatbuffers::FlatBufferBuilder& builder, + int64_t dim_value) { + const auto dim_value_off = + fbs::CreateDimensionValue(builder, fbs::DimensionValueType::VALUE, dim_value); + std::vector> dims{ + fbs::CreateDimension(builder, dim_value_off)}; + const auto shape = fbs::CreateShapeDirect(builder, &dims); + const auto tensor_type = + fbs::CreateTensorTypeAndShape(builder, fbs::TensorDataType::FLOAT, shape); + return fbs::CreateTypeInfoDirect(builder, nullptr, fbs::TypeInfoValue::tensor_type, tensor_type.Union()); +} + +std::vector BuildOrtModelBuffer( + const std::function(flatbuffers::FlatBufferBuilder&)>& create_graph) { + flatbuffers::FlatBufferBuilder builder; + + const auto graph = create_graph(builder); + std::vector> opset_imports{ + fbs::CreateOperatorSetIdDirect(builder, "", 18)}; + const auto model = fbs::CreateModelDirect(builder, 8, &opset_imports, "ort-model-test", "1", "", + 1, "", graph, ""); + const auto session = fbs::CreateInferenceSessionDirect(builder, ORT_VERSION, model); + fbs::FinishInferenceSessionBuffer(builder, session); + + return std::vector(builder.GetBufferPointer(), builder.GetBufferPointer() + builder.GetSize()); +} + +Status LoadOrtBuffer(const std::vector& buffer) { + SessionOptions so; + ORT_RETURN_IF_ERROR(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigLoadModelFormat, "ORT")); + + InferenceSessionWrapper session_object{so, GetEnvironment()}; + return session_object.Load(buffer.data(), static_cast(buffer.size())); +} + +} // namespace + static void RunOrtModel(const OrtModelTestInfo& test_info) { SessionOptions so; so.session_logid = test_info.logid; @@ -85,6 +126,54 @@ static void RunOrtModel(const OrtModelTestInfo& test_info) { test_info.output_verifier(fetches); } +TEST(OrtModelTest, RejectsInitializerRawDataSizeMismatch) { + const auto buffer = BuildOrtModelBuffer([](flatbuffers::FlatBufferBuilder& builder) { + std::vector dims{1}; + std::vector raw_data(sizeof(float) * 2, 0); + std::vector> initializers{ + fbs::CreateTensorDirect(builder, "bad_initializer", "", &dims, fbs::TensorDataType::FLOAT, &raw_data)}; + return fbs::CreateGraphDirect(builder, &initializers); + }); + + const auto status = LoadOrtBuffer(buffer); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("raw data size mismatch")); +} + +TEST(OrtModelTest, RejectsNullNodeArgTableEntry) { + auto buffer = BuildOrtModelBuffer([](flatbuffers::FlatBufferBuilder& builder) { + std::vector> node_args{ + fbs::CreateValueInfoDirect(builder, "X", "", CreateFloatTensorTypeInfo(builder, 1))}; + return fbs::CreateGraphDirect(builder, nullptr, &node_args); + }); + + const auto* fbs_session = fbs::GetInferenceSession(buffer.data()); + ASSERT_NE(fbs_session, nullptr); + const auto* fbs_node_args = fbs_session->model()->graph()->node_args(); + ASSERT_NE(fbs_node_args, nullptr); + + auto* raw_offsets = const_cast(reinterpret_cast(fbs_node_args->Data())); + std::fill_n(raw_offsets, sizeof(flatbuffers::uoffset_t), 0); + + const auto status = LoadOrtBuffer(buffer); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), + testing::AnyOf(testing::HasSubstr("Null node arg entry"), + testing::HasSubstr("verification failed"))); +} + +TEST(OrtModelTest, RejectsDanglingNodeEdge) { + const auto buffer = BuildOrtModelBuffer([](flatbuffers::FlatBufferBuilder& builder) { + std::vector> node_edges{ + fbs::CreateNodeEdgeDirect(builder, 0)}; + return fbs::CreateGraphDirect(builder, nullptr, nullptr, nullptr, 1, &node_edges); + }); + + const auto status = LoadOrtBuffer(buffer); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("references missing node")); +} + #if !defined(ORT_MINIMAL_BUILD) // Keep the CompareTypeProtos in case we need debug the difference /* From 28cf322ed75a52ff135cdcac9f51d3122758a391 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 22 Apr 2026 16:38:42 -0700 Subject: [PATCH 2/5] Address review feedback: add node index amplification guard and test - Add sanity bound on required_node_slot_count to reject crafted buffers where a single node with a huge index could cause multi-GB allocation. Cap at max(1024, 2 * total_entries). - Read tensor name from fbs_tensor.name() instead of the partially-populated protobuf initializer for the raw data size mismatch error message. - Add RejectsAdversarialLargeNodeIndex test to exercise the new bound. --- onnxruntime/core/graph/graph.cc | 13 +++++++++++++ onnxruntime/core/graph/graph_flatbuffers_utils.cc | 3 ++- onnxruntime/test/framework/ort_model_only_test.cc | 14 ++++++++++++++ 3 files changed, 29 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index ff400f8a23f81..0a257a34d1a1e 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -6639,6 +6639,19 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph } } + // Sanity bound: the slot count must not wildly exceed the number of actual entries. + // A legitimate model has max_node_index close to the number of nodes. Allow a reasonable + // multiplier to tolerate sparse index ranges, but reject buffers where a single node entry + // with a huge index would cause a multi-gigabyte allocation. + const size_t total_entries = (fbs_nodes != nullptr ? fbs_nodes->size() : 0U) + + (fbs_node_edges != nullptr ? fbs_node_edges->size() : 0U); + constexpr size_t kMinSlotCap = 1024; // allow small models without penalty + const size_t slot_cap = std::max(kMinSlotCap, total_entries * 2U); + ORT_RETURN_IF(required_node_slot_count > slot_cap, + "Node index ", required_node_slot_count - 1, + " is unreasonably large relative to the number of entries (", + total_entries, "). Invalid ORT format model."); + ORT_RETURN_IF(fbs_graph.max_node_index() < required_node_slot_count, "Serialized max node index is smaller than the required node slot count. Invalid ORT format model."); nodes_.resize(required_node_slot_count); diff --git a/onnxruntime/core/graph/graph_flatbuffers_utils.cc b/onnxruntime/core/graph/graph_flatbuffers_utils.cc index fbe2b24ac3ca3..16e4104fb76e9 100644 --- a/onnxruntime/core/graph/graph_flatbuffers_utils.cc +++ b/onnxruntime/core/graph/graph_flatbuffers_utils.cc @@ -374,7 +374,8 @@ Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor, TensorProto& init if (fbs_raw_data) { size_t expected_num_bytes = 0; ORT_RETURN_IF_ERROR(GetSizeInBytesFromFbsTensor(fbs_tensor, expected_num_bytes)); - const char* tensor_name = initializer.has_name() ? initializer.name().c_str() : ""; + const auto* fbs_name = fbs_tensor.name(); + const char* tensor_name = fbs_name ? fbs_name->c_str() : ""; ORT_RETURN_IF( fbs_raw_data->size() != expected_num_bytes, "Initializer raw data size mismatch for tensor '", tensor_name, diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index 9a6d8d23df53b..ee24f07ec6092 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -174,6 +174,20 @@ TEST(OrtModelTest, RejectsDanglingNodeEdge) { EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("references missing node")); } +TEST(OrtModelTest, RejectsAdversarialLargeNodeIndex) { + // A single node with a huge index should be rejected to prevent memory amplification. + const uint32_t huge_index = 100'000'000; + const auto buffer = BuildOrtModelBuffer([huge_index](flatbuffers::FlatBufferBuilder& builder) { + std::vector> nodes{ + fbs::CreateNodeDirect(builder, "n", "", "", 1, huge_index, "Identity")}; + return fbs::CreateGraphDirect(builder, nullptr, nullptr, &nodes, huge_index + 1); + }); + + const auto status = LoadOrtBuffer(buffer); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("unreasonably large")); +} + #if !defined(ORT_MINIMAL_BUILD) // Keep the CompareTypeProtos in case we need debug the difference /* From e469e755407db4c0cb65a272441d746bb5bb919c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 23 Apr 2026 01:54:08 -0700 Subject: [PATCH 3/5] fix: resolve CI failures in ort_model_only_test - Move huge_index variable inside lambda to fix -Wunused-lambda-capture error on Clang (macOS, wasm, Android NDK builds) - Use kOrtModelVersion instead of ORT_VERSION for the ORT format version field so minimal builds accept the test model (version 6 vs "1.26.0") --- onnxruntime/test/framework/ort_model_only_test.cc | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index ee24f07ec6092..02f4b013c53a2 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -3,6 +3,7 @@ #include +#include "core/flatbuffers/ort_format_version.h" #include "core/flatbuffers/schema/ort.fbs.h" #include "core/framework/data_types.h" #include "core/framework/tensorprotoutils.h" @@ -66,7 +67,8 @@ std::vector BuildOrtModelBuffer( fbs::CreateOperatorSetIdDirect(builder, "", 18)}; const auto model = fbs::CreateModelDirect(builder, 8, &opset_imports, "ort-model-test", "1", "", 1, "", graph, ""); - const auto session = fbs::CreateInferenceSessionDirect(builder, ORT_VERSION, model); + const auto session = fbs::CreateInferenceSessionDirect(builder, + std::to_string(kOrtModelVersion).c_str(), model); fbs::FinishInferenceSessionBuffer(builder, session); return std::vector(builder.GetBufferPointer(), builder.GetBufferPointer() + builder.GetSize()); @@ -176,8 +178,8 @@ TEST(OrtModelTest, RejectsDanglingNodeEdge) { TEST(OrtModelTest, RejectsAdversarialLargeNodeIndex) { // A single node with a huge index should be rejected to prevent memory amplification. - const uint32_t huge_index = 100'000'000; - const auto buffer = BuildOrtModelBuffer([huge_index](flatbuffers::FlatBufferBuilder& builder) { + const auto buffer = BuildOrtModelBuffer([](flatbuffers::FlatBufferBuilder& builder) { + const uint32_t huge_index = 100'000'000; std::vector> nodes{ fbs::CreateNodeDirect(builder, "n", "", "", 1, huge_index, "Identity")}; return fbs::CreateGraphDirect(builder, nullptr, nullptr, &nodes, huge_index + 1); From 95172db0d29e2ff2c087305687642782195fa625 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 2 May 2026 00:54:26 -0700 Subject: [PATCH 4/5] Address review feedback: widen sparsity cap and validate EdgeEnd node indices - Increase slot_cap multiplier from 2x to 64x with absolute cap of 10M to accommodate legitimately sparse node index ranges produced by ORT graph optimizations (ReleaseNode leaves holes). - Add bounds + nullptr checks for EdgeEnd.node_index in LoadEdgesFromOrtFormat to prevent ORT_ENFORCE crash or nullptr dereference on malformed buffers. - Add RejectsInvalidEdgeEndNodeIndex regression test. - Update stale comment at nodes_.resize() site. --- onnxruntime/core/graph/graph.cc | 29 +++++++++++++------ .../test/framework/ort_model_only_test.cc | 23 ++++++++++++++- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 0a257a34d1a1e..38b9b8a6b7fb6 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -872,7 +872,15 @@ Status Node::LoadEdgesFromOrtFormat(const onnxruntime::fbs::NodeEdge& fbs_node_e if (fbs_edges) { for (const auto* fbs_edge : *fbs_edges) { ORT_RETURN_IF(nullptr == fbs_edge, "Node::LoadEdgesFromOrtFormat, edge is missing for ", dst_name); - edge_set.emplace(*graph.GetNode(fbs_edge->node_index()), fbs_edge->src_arg_index(), fbs_edge->dst_arg_index()); + const auto edge_node_index = fbs_edge->node_index(); + ORT_RETURN_IF(edge_node_index >= static_cast(graph.MaxNodeIndex()), + "Node::LoadEdgesFromOrtFormat, ", dst_name, " has out-of-range node index ", + edge_node_index, ". Invalid ORT format model."); + const auto* edge_node = graph.GetNode(edge_node_index); + ORT_RETURN_IF(edge_node == nullptr, + "Node::LoadEdgesFromOrtFormat, ", dst_name, " references missing node ", + edge_node_index, ". Invalid ORT format model."); + edge_set.emplace(*edge_node, fbs_edge->src_arg_index(), fbs_edge->dst_arg_index()); } } return Status::OK(); @@ -6609,8 +6617,9 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph // Nodes // - // Since we access a node using its index, we need to have nodes_ with size max_node_index to avoid - // out of bounds access. + // Since we access a node using its index, we need to have nodes_ with a size that covers all + // referenced indices. We compute the required slot count from actual node and edge data rather + // than trusting the serialized max_node_index field. auto* fbs_nodes = fbs_graph.nodes(); ORT_RETURN_IF_ERROR(fbs::utils::ValidateRequiredTableOffsets(fbs_nodes, "node")); auto* fbs_node_edges = fbs_graph.node_edges(); @@ -6639,14 +6648,16 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph } } - // Sanity bound: the slot count must not wildly exceed the number of actual entries. - // A legitimate model has max_node_index close to the number of nodes. Allow a reasonable - // multiplier to tolerate sparse index ranges, but reject buffers where a single node entry - // with a huge index would cause a multi-gigabyte allocation. + // Sanity bound: reject buffers where a crafted node index would cause a multi-gigabyte + // allocation. After graph optimizations, ORT preserves original node indices (leaving holes + // in the nodes_ vector), so max_node_index can legitimately be much larger than the number + // of remaining nodes. Use a generous multiplier plus an absolute cap to accommodate this + // sparsity while still blocking adversarial amplification. const size_t total_entries = (fbs_nodes != nullptr ? fbs_nodes->size() : 0U) + (fbs_node_edges != nullptr ? fbs_node_edges->size() : 0U); - constexpr size_t kMinSlotCap = 1024; // allow small models without penalty - const size_t slot_cap = std::max(kMinSlotCap, total_entries * 2U); + constexpr size_t kMinSlotCap = 1024; // allow small models without penalty + constexpr size_t kAbsoluteSlotCap = 10000000; // ~80 MB of unique_ptr on 64-bit + const size_t slot_cap = std::min(kAbsoluteSlotCap, std::max(kMinSlotCap, total_entries * 64U)); ORT_RETURN_IF(required_node_slot_count > slot_cap, "Node index ", required_node_slot_count - 1, " is unreasonably large relative to the number of entries (", diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index 02f4b013c53a2..a84db21f869b5 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -68,7 +68,7 @@ std::vector BuildOrtModelBuffer( const auto model = fbs::CreateModelDirect(builder, 8, &opset_imports, "ort-model-test", "1", "", 1, "", graph, ""); const auto session = fbs::CreateInferenceSessionDirect(builder, - std::to_string(kOrtModelVersion).c_str(), model); + std::to_string(kOrtModelVersion).c_str(), model); fbs::FinishInferenceSessionBuffer(builder, session); return std::vector(builder.GetBufferPointer(), builder.GetBufferPointer() + builder.GetSize()); @@ -190,6 +190,27 @@ TEST(OrtModelTest, RejectsAdversarialLargeNodeIndex) { EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("unreasonably large")); } +TEST(OrtModelTest, RejectsInvalidEdgeEndNodeIndex) { + // An EdgeEnd referencing a non-existent node should be rejected gracefully + // rather than crashing via ORT_ENFORCE or nullptr dereference. + const auto buffer = BuildOrtModelBuffer([](flatbuffers::FlatBufferBuilder& builder) { + // Create a valid node at index 0 + std::vector> nodes{ + fbs::CreateNodeDirect(builder, "n0", "", "", 1, 0, "Identity")}; + // Create a NodeEdge for node 0 with an input edge referencing non-existent node 99 + std::vector input_edges{fbs::EdgeEnd(99, 0, 0)}; + std::vector> node_edges{ + fbs::CreateNodeEdgeDirect(builder, 0, &input_edges)}; + return fbs::CreateGraphDirect(builder, nullptr, nullptr, &nodes, 100, &node_edges); + }); + + const auto status = LoadOrtBuffer(buffer); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), + testing::AnyOf(testing::HasSubstr("out-of-range node index"), + testing::HasSubstr("references missing node"))); +} + #if !defined(ORT_MINIMAL_BUILD) // Keep the CompareTypeProtos in case we need debug the difference /* From 89409b8441a68792c7825857939174e2e4c30b33 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 2 May 2026 08:26:13 -0700 Subject: [PATCH 5/5] fix(test): provide empty node args in RejectsInvalidEdgeEndNodeIndex test The test created a node with null inputs/outputs/implicit_inputs fields, which caused Node::LoadFromOrtFormat to fail with 'fbs_node_arg_names cannot be null' before reaching the edge validation logic under test. Provide explicit empty vectors so the node loads successfully and the edge referencing non-existent node 99 triggers the intended 'out-of-range node index' error. --- onnxruntime/test/framework/ort_model_only_test.cc | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index a84db21f869b5..3cd949201cd3d 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -194,9 +194,14 @@ TEST(OrtModelTest, RejectsInvalidEdgeEndNodeIndex) { // An EdgeEnd referencing a non-existent node should be rejected gracefully // rather than crashing via ORT_ENFORCE or nullptr dereference. const auto buffer = BuildOrtModelBuffer([](flatbuffers::FlatBufferBuilder& builder) { - // Create a valid node at index 0 + // Create a valid node at index 0 with empty inputs/outputs so it passes node-loading validation. + std::vector> empty_args; + std::vector empty_arg_counts; std::vector> nodes{ - fbs::CreateNodeDirect(builder, "n0", "", "", 1, 0, "Identity")}; + fbs::CreateNodeDirect(builder, "n0", "", "", 1, 0, "Identity", + fbs::NodeType::Primitive, nullptr, + &empty_args, &empty_args, nullptr, + &empty_arg_counts, &empty_args)}; // Create a NodeEdge for node 0 with an input edge referencing non-existent node 99 std::vector input_edges{fbs::EdgeEnd(99, 0, 0)}; std::vector> node_edges{