diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 9941224258506..4572f4b68d108 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -111,6 +111,18 @@ static const char* const kOrtSessionOptionsMemoryOptimizerProbeConfig = "optimiz // Default is an empty string which means no optimizers are disabled. static const char* const kOrtSessionOptionsDisableSpecifiedOptimizers = "optimization.disable_specified_optimizers"; +// Maximum total output size in bytes that the constant folding optimizer is allowed to produce per node. +// Prevents malicious models from causing excessive memory allocation during optimization. +// If the estimated or actual output size of a constant-foldable node exceeds this limit, the node will +// not be constant folded and will instead be executed at runtime. +// +// Option values: +// - A positive integer (as string): Maximum allowed output size in bytes per constant-folded node. +// Default is "1073741824" (1 GB). +// - "0": Disable the size limit (not recommended for untrusted models). +static const char* const kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes = + "optimization.constant_folding_max_output_size_in_bytes"; + // It controls whether to run graph optimizations in loop or not. // // "0": disable. Graph Optimization Loop is disabled. diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index cb6d65342bc54..1b8bb57d6a74e 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -2,15 +2,18 @@ // Licensed under the MIT License. #include +#include #include "core/optimizer/constant_folding.h" #include "core/optimizer/initializer.h" #include "core/optimizer/utils.h" #include "core/graph/graph_utils.h" #include "core/optimizer/optimizer_execution_frame.h" -#include "core/optimizer/utils.h" #include "core/framework/op_kernel.h" #include "core/framework/tensorprotoutils.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "core/common/safeint.h" +#include "core/common/parse_string.h" using namespace onnxruntime::common; @@ -140,6 +143,154 @@ static Status ConstantFoldIfNode(Graph& graph, Node& if_node, const logging::Log return status; } +// Default maximum output size per constant-folded node: 1 GB. +// This prevents malicious models from causing excessive memory allocation during optimization. +static constexpr int64_t kDefaultConstantFoldingMaxOutputSizeInBytes = 1024 * 1024 * 1024; + +static size_t GetElementSizeForConstantFolding(ONNX_NAMESPACE::TensorProto_DataType elem_type) { + const size_t element_size = utils::GetElementSizeOfTensor(elem_type); + if (element_size != 0) { + return element_size; + } + + // String tensors allocate storage for std::string slots even though the payload size is variable. + return elem_type == ONNX_NAMESPACE::TensorProto_DataType_STRING ? sizeof(std::string) : 0; +} + +static int64_t EstimateTensorElementCount(const ONNX_NAMESPACE::TensorShapeProto& shape) { + SafeInt num_elements = 1; + for (int i = 0; i < shape.dim_size(); ++i) { + const auto& dim = shape.dim(i); + if (!utils::HasDimValue(dim)) { + return -1; // Symbolic dimension + } + int64_t dim_value = dim.dim_value(); + if (dim_value < 0) { + return -1; // Invalid dimension + } + num_elements *= dim_value; + } + + return num_elements; +} + +static int64_t EstimateTensorSizeInBytes(const NodeArg& node_arg) { + const auto* type_proto = node_arg.TypeAsProto(); + if (type_proto == nullptr || !utils::HasTensorType(*type_proto)) { + return -1; // Cannot estimate non-tensor or unknown types + } + + const auto* shape = node_arg.Shape(); + if (shape == nullptr) { + return -1; // Unknown shape + } + + auto elem_type = static_cast( + type_proto->tensor_type().elem_type()); + size_t element_size = GetElementSizeForConstantFolding(elem_type); + if (element_size == 0) { + return -1; // Unknown element type + } + + int64_t num_elements = EstimateTensorElementCount(*shape); + if (num_elements < 0) { + return -1; + } + + return SafeInt(num_elements) * static_cast(element_size); +} + +static int64_t EstimateUniqueOutputSizeInBytes(const Node& node) { + const auto& input_defs = node.InputDefs(); + if (input_defs.empty() || input_defs[0] == nullptr) { + return -1; + } + + const int64_t input_num_elements = input_defs[0]->Shape() != nullptr + ? EstimateTensorElementCount(*input_defs[0]->Shape()) + : -1; + if (input_num_elements < 0) { + return -1; + } + + const auto* input_type_proto = input_defs[0]->TypeAsProto(); + if (input_type_proto == nullptr || !utils::HasTensorType(*input_type_proto)) { + return -1; + } + + auto input_elem_type = static_cast( + input_type_proto->tensor_type().elem_type()); + const size_t input_element_size = GetElementSizeForConstantFolding(input_elem_type); + if (input_element_size == 0) { + return -1; + } + + SafeInt total_size = 0; + const auto& output_defs = node.OutputDefs(); + for (size_t output_idx = 0; output_idx < output_defs.size(); ++output_idx) { + const auto* output_def = output_defs[output_idx]; + if (!output_def->Exists()) { + continue; + } + + const size_t element_size = output_idx == 0 ? input_element_size : sizeof(int64_t); + total_size += SafeInt(input_num_elements) * static_cast(element_size); + } + + return total_size; +} + +static int64_t EstimateIdentityOutputSizeInBytes(const Node& node) { + const auto& input_defs = node.InputDefs(); + if (input_defs.empty() || input_defs[0] == nullptr) { + return -1; + } + + return EstimateTensorSizeInBytes(*input_defs[0]); +} + +// Estimate the total output size in bytes for a node using shape inference results. +// Returns -1 if the output size cannot be estimated (e.g., unknown shapes or types). +static int64_t EstimateNodeOutputSizeInBytes(const Node& node) { + if (node.OpType() == "Identity" && node.Domain().empty()) { + return EstimateIdentityOutputSizeInBytes(node); + } + + if (node.OpType() == "Unique" && node.Domain().empty()) { + return EstimateUniqueOutputSizeInBytes(node); + } + + SafeInt total_size = 0; + for (const auto* output_def : node.OutputDefs()) { + if (!output_def->Exists()) { + continue; + } + + const int64_t output_size = EstimateTensorSizeInBytes(*output_def); + if (output_size < 0) { + return -1; + } + + total_size += output_size; + } + + return total_size; +} + +// Get the configured max output size from session options, or use the default. +static int64_t GetConstantFoldingMaxOutputSize(const ConfigOptions& config_options) { + std::string max_size_str = config_options.GetConfigOrDefault( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, + std::to_string(kDefaultConstantFoldingMaxOutputSizeInBytes)); + + int64_t max_size = 0; + if (!TryParseStringWithClassicLocale(max_size_str, max_size) || max_size < 0) { + max_size = kDefaultConstantFoldingMaxOutputSizeInBytes; + } + + return max_size; +} + Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { bool have_updated_nodes = false; GraphViewer graph_viewer(graph); @@ -151,6 +302,8 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, }; #endif + const int64_t max_output_size = GetConstantFoldingMaxOutputSize(config_options_); + for (NodeIndex i : order) { auto* node = graph.GetNode(i); if (!node || !AllowConstantFolding(*node)) { @@ -233,6 +386,35 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, } } + // Check if the estimated output size exceeds the configured limit. + // This prevents malicious models from causing excessive memory allocation during constant folding. + if (max_output_size > 0) { + int64_t estimated_size = -1; + try { + estimated_size = EstimateNodeOutputSizeInBytes(*node); + } catch (const std::exception&) { + // SafeInt overflow means the size is astronomically large - definitely skip + LOGS(logger, WARNING) << "Integer overflow while estimating output size of " + << node->OpType() << " node '" << node->Name() + << "'. Skipping constant folding for this node."; + continue; + } + + if (estimated_size > max_output_size) { + LOGS(logger, WARNING) << "Skipping constant folding for " << node->OpType() + << " node '" << node->Name() + << "' because estimated output size (" << estimated_size + << " bytes) exceeds the limit (" << max_output_size << " bytes)."; + continue; + } + if (estimated_size < 0) { + LOGS(logger, INFO) << "Skipping constant folding for " << node->OpType() + << " node '" << node->Name() + << "' because output size could not be estimated before execution."; + continue; + } + } + #if !defined(DISABLE_SPARSE_TENSORS) // Create execution frame for executing constant nodes. OptimizerExecutionFrame::Info info({node}, constant_inputs, graph.ModelPath(), execution_provider_, @@ -312,7 +494,25 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, #pragma warning(disable : 6387) #endif OpKernelContext op_kernel_context(&frame, kernel.get(), /*stream*/ nullptr, nullptr, logger); - ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context)); + + // Skip the current node if Compute fails so one bad constant-fold candidate does not abort + // the entire constant folding pass. + Status compute_status = Status::OK(); + try { + compute_status = kernel->Compute(&op_kernel_context); + } catch (const std::exception& ex) { + LOGS(logger, WARNING) << "Exception during constant folding of " << node->OpType() + << " node '" << node->Name() << "': " << ex.what() + << ". Skipping constant folding for this node."; + continue; + } + + if (!compute_status.IsOK()) { + LOGS(logger, WARNING) << "Failure during constant folding of " << node->OpType() + << " node '" << node->Name() << "': " << compute_status.ErrorMessage() + << ". Skipping constant folding for this node."; + continue; + } #ifdef _WIN32 #pragma warning(pop) #endif @@ -320,6 +520,33 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, std::vector fetches; ORT_RETURN_IF_ERROR(frame.GetOutputs(fetches)); + // Post-execution size check: verify actual output sizes don't exceed the limit. + // This catches cases where pre-execution shape inference couldn't determine the output size. + if (max_output_size > 0) { + SafeInt actual_total_size = 0; + bool size_exceeded = false; + try { + for (size_t fetch_idx = 0; fetch_idx < fetches.size(); ++fetch_idx) { + if (fetches[fetch_idx].IsAllocated() && fetches[fetch_idx].IsTensor()) { + const auto& tensor = fetches[fetch_idx].Get(); + actual_total_size += tensor.SizeInBytes(); + } + } + size_exceeded = actual_total_size > max_output_size; + } catch (const std::exception&) { + // SafeInt overflow means total size is astronomically large + size_exceeded = true; + } + + if (size_exceeded) { + LOGS(logger, WARNING) << "Skipping constant folding for " << node->OpType() + << " node '" << node->Name() + << "' because actual output size exceeds the limit (" + << max_output_size << " bytes)."; + continue; + } + } + // Go over all output node args and substitute them with the newly computed tensors, which will be // added to the graph as initializers. ORT_ENFORCE(fetches.size() == fetch_to_output_idx.size()); diff --git a/onnxruntime/core/providers/cpu/tensor/expand.cc b/onnxruntime/core/providers/cpu/tensor/expand.cc index b0c636281bc7a..6d299282f3e60 100644 --- a/onnxruntime/core/providers/cpu/tensor/expand.cc +++ b/onnxruntime/core/providers/cpu/tensor/expand.cc @@ -94,8 +94,8 @@ Status Expand::Compute(OpKernelContext* context) const { auto input_dim = input_dims_iter > -1 ? input_dims[input_dims_iter] : 1; auto output_dim = output_dims[output_dims_iter]; - input_count *= input_dim; - output_count *= output_dim; + input_count = SafeInt(input_count) * input_dim; + output_count = SafeInt(output_count) * output_dim; if (0 == input_count || 0 == output_count) { return Status::OK(); @@ -106,26 +106,26 @@ Status Expand::Compute(OpKernelContext* context) const { input_dim_group[onnxruntime::narrow(dim_group_start)] = input_count; output_dim_group[onnxruntime::narrow(dim_group_start)] = output_count; expand_dim_size[onnxruntime::narrow(dim_group_start)] = output_count / input_count / last_dim_size; - last_dim_size *= expand_dim_size[onnxruntime::narrow(dim_group_start)]; + last_dim_size = SafeInt(last_dim_size) * expand_dim_size[onnxruntime::narrow(dim_group_start)]; } } auto distribute_count = input_dim_group[onnxruntime::narrow(dim_group_start)] / input_dim_group[SafeInt(max_dims_size) - 1]; std::vector output_offsets(onnxruntime::narrow(distribute_count), 0); int64_t copy_len = input_dim_group[SafeInt(max_dims_size) - 1]; - auto copy_byte = copy_len * sizeof(T); + size_t copy_byte = SafeInt(copy_len) * sizeof(T); auto distribute_fn = [&](ptrdiff_t i_start, ptrdiff_t i_end) { for (auto i = i_start; i < i_end; i++) { - auto input_offset = i * copy_len; + int64_t input_offset = SafeInt(i) * copy_len; int64_t output_offset = 0; for (auto j = dim_group_start + 1, remains = input_offset; j < max_dims_size; ++j) { auto current_count = remains / input_dim_group[onnxruntime::narrow(j)]; - output_offset += current_count * output_dim_group[onnxruntime::narrow(j)]; + output_offset = SafeInt(output_offset) + SafeInt(current_count) * output_dim_group[onnxruntime::narrow(j)]; remains = remains % input_dim_group[onnxruntime::narrow(j)]; } // for j - memcpy(output_data + output_offset, input_data + input_offset, onnxruntime::narrow(copy_byte)); + memcpy(output_data + output_offset, input_data + input_offset, copy_byte); output_offsets[onnxruntime::narrow(i)] = output_offset; } // for i }; // distribute_fn diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index ce53703eabf18..cf99e31d445ce 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1473,6 +1473,201 @@ TEST_F(GraphTransformationTests, ConstantFoldingIfConstantInliningEdgesWithMiddl ASSERT_TRUE(dest_edges.find(2) != dest_edges.end()); } +// Test that constant folding respects the output size limit and skips nodes +// whose output would exceed it. This is a security measure against malicious +// models that could cause memory exhaustion during optimization. +TEST_F(GraphTransformationTests, ConstantFoldingOutputSizeLimit) { + // Build a model with an Expand node: scalar input [1.0] expanded by shape [1024, 1024]. + // Output = 1024*1024 * 4 bytes = 4 MB of float data. + // With a 1 MB limit, this should NOT be constant folded. + // With a 8 MB limit, this SHOULD be constant folded. + + auto build_model = [&](ModelTestBuilder& builder) { + auto* input_data = builder.MakeInitializer({1}, {1.0f}); + auto* shape_data = builder.MakeInitializer({2}, {1024, 1024}); + auto* output_arg = builder.MakeOutput(); + + builder.AddNode("Expand", {input_data, shape_data}, {output_arg}); + }; + + // Test 1: With a 1 MB limit, the Expand node should NOT be folded (output is ~4 MB). + { + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Expand should remain because output is too large + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + ConfigOptions config_options; + // Set limit to 1 MB + ASSERT_STATUS_OK(config_options.AddConfigEntry( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, "1048576")); + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); + } + + // Test 2: With an 8 MB limit, the Expand node SHOULD be folded (output is ~4 MB). + { + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Expand should be folded since output is within limit + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 0); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + ConfigOptions config_options; + // Set limit to 8 MB + ASSERT_STATUS_OK(config_options.AddConfigEntry( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, "8388608")); + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); + } +} + +// Test that an explicitly configured constant folding output-size limit blocks +// folding a very large ConstantOfShape output. +TEST_F(GraphTransformationTests, ConstantFoldingConfiguredLimitBlocksLargeConstantOfShape) { + // Build a model with a ConstantOfShape node producing a huge output. + // Shape = [16384, 16384] = 268M elements * 4 bytes = 1 GB. + // Use an explicit 512 MB limit so the 1 GB output is not folded. + + auto build_model = [&](ModelTestBuilder& builder) { + auto* shape_data = builder.MakeInitializer({2}, {16384, 16384}); + auto* output_arg = builder.MakeOutput(); + + auto& node = builder.AddNode("ConstantOfShape", {shape_data}, {output_arg}); + // Default value is float 0.0 + ONNX_NAMESPACE::AttributeProto value_attr; + value_attr.set_name("value"); + value_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + auto* tensor = value_attr.mutable_t(); + tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor->add_dims(1); + tensor->add_float_data(0.0f); + node.AddAttributeProto(std::move(value_attr)); + }; + + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["ConstantOfShape"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // ConstantOfShape should remain because output is too large (1 GB > 512 MB limit) + TEST_RETURN_IF_NOT(op_to_count["ConstantOfShape"] == 1); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + ConfigOptions config_options; + // Set limit to 512 MB so the 1 GB output is blocked + ASSERT_STATUS_OK(config_options.AddConfigEntry( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, "536870912")); + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); +} + +// Test that small constant folding still works with the size limit. +TEST_F(GraphTransformationTests, ConstantFoldingSmallOutputAllowed) { + // Build a model with a small Expand: scalar -> [4, 4] = 16 * 4 = 64 bytes. + // This is well within even a small limit and should be folded. + + auto build_model = [&](ModelTestBuilder& builder) { + auto* input_data = builder.MakeInitializer({1}, {42.0f}); + auto* shape_data = builder.MakeInitializer({2}, {4, 4}); + auto* output_arg = builder.MakeOutput(); + + builder.AddNode("Expand", {input_data, shape_data}, {output_arg}); + }; + + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Small Expand should be constant folded + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 0); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + const ConfigOptions empty_config_options; + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, empty_config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); +} + +// Test that constant folding gracefully handles an Expand node whose output shape +// dimensions would cause integer overflow. This simulates the attack vector where +// a malicious model embeds constant initializers with extreme shape values, causing +// kernel Compute() to execute during graph optimization. The SafeInt-protected +// arithmetic in Expand::Compute (or TensorShape overflow) should be caught by the +// try/catch around Compute, and the node should NOT be constant folded. +TEST_F(GraphTransformationTests, ConstantFoldingExpandOverflowDimsSkipped) { + constexpr int64_t kLargeDim = int64_t(1) << 32; // 4294967296 + + auto build_model = [&](ModelTestBuilder& builder) { + auto* input_data = builder.MakeInitializer({1}, {1.0f}); + auto* shape_data = builder.MakeInitializer({2}, {kLargeDim, kLargeDim}); + auto* output_arg = builder.MakeOutput(); + + builder.AddNode("Expand", {input_data, shape_data}, {output_arg}); + }; + + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Expand should remain because the overflow prevents constant folding. + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + std::unique_ptr e = + std::make_unique(CPUExecutionProviderInfo()); + const ConfigOptions empty_config_options; + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, empty_config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); +} + // Check transformations in the case of a subgraph with constant inputs. TEST_F(GraphTransformationTests, SubgraphWithConstantInputs) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "constant-subgraph.onnx";