fix(security): add SafeInt overflow protection in Expand and constant folding output size limit

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.

onnxruntime/core/optimizer/constant_folding.cc

@@ -2,15 +2,18 @@
 // Licensed under the MIT License.
 
 #include <limits>
+#include <string>
 
 #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<int64_t> 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<ONNX_NAMESPACE::TensorProto_DataType>(
+      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<int64_t>(num_elements) * static_cast<int64_t>(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<ONNX_NAMESPACE::TensorProto_DataType>(
+      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<int64_t> 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<int64_t>(input_num_elements) * static_cast<int64_t>(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<int64_t> 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,14 +494,59 @@ 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
 
       std::vector<OrtValue> 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<int64_t> 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<Tensor>();
+              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());

onnxruntime/core/providers/cpu/tensor/expand.cc

@@ -94,8 +94,8 @@ Status Expand<T>::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<int64_t>(input_count) * input_dim;
+    output_count = SafeInt<int64_t>(output_count) * output_dim;
 
     if (0 == input_count || 0 == output_count) {
       return Status::OK();
@@ -106,26 +106,26 @@ Status Expand<T>::Compute(OpKernelContext* context) const {
       input_dim_group[onnxruntime::narrow<size_t>(dim_group_start)] = input_count;
       output_dim_group[onnxruntime::narrow<size_t>(dim_group_start)] = output_count;
       expand_dim_size[onnxruntime::narrow<size_t>(dim_group_start)] = output_count / input_count / last_dim_size;
-      last_dim_size *= expand_dim_size[onnxruntime::narrow<size_t>(dim_group_start)];
+      last_dim_size = SafeInt<int64_t>(last_dim_size) * expand_dim_size[onnxruntime::narrow<size_t>(dim_group_start)];
     }
   }
 
   auto distribute_count = input_dim_group[onnxruntime::narrow<size_t>(dim_group_start)] / input_dim_group[SafeInt<size_t>(max_dims_size) - 1];
   std::vector<int64_t> output_offsets(onnxruntime::narrow<size_t>(distribute_count), 0);
   int64_t copy_len = input_dim_group[SafeInt<size_t>(max_dims_size) - 1];
-  auto copy_byte = copy_len * sizeof(T);
+  size_t copy_byte = SafeInt<size_t>(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<int64_t>(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<size_t>(j)];
-            output_offset += current_count * output_dim_group[onnxruntime::narrow<size_t>(j)];
+            output_offset = SafeInt<int64_t>(output_offset) + SafeInt<int64_t>(current_count) * output_dim_group[onnxruntime::narrow<size_t>(j)];
             remains = remains % input_dim_group[onnxruntime::narrow<size_t>(j)];
           }  // for j
-          memcpy(output_data + output_offset, input_data + input_offset, onnxruntime::narrow<size_t>(copy_byte));
+          memcpy(output_data + output_offset, input_data + input_offset, copy_byte);
           output_offsets[onnxruntime::narrow<size_t>(i)] = output_offset;
         }  // for i
       };  // distribute_fn