fix(ml): use SafeInt checked arithmetic in ML operator coefficient size validation

onnxruntime/core/providers/cpu/ml/linearclassifier.cc

@@ -3,6 +3,7 @@
 
 #include "core/providers/cpu/ml/linearclassifier.h"
 #include "core/common/narrow.h"
+#include "core/common/safeint.h"
 #include "core/providers/cpu/math/gemm.h"
 
 namespace onnxruntime {
@@ -151,14 +152,16 @@ Status LinearClassifier::Compute(OpKernelContext* ctx) const {
   ptrdiff_t num_features = input_shape.NumDimensions() == 1 ? narrow<ptrdiff_t>(
                                                                   input_shape[0])
                                                             : narrow<ptrdiff_t>(input_shape[1]);
-
-  // Validate coefficients are large enough to prevent OOB read in GEMM.
-  const size_t expected_coefficients_size = SafeInt<size_t>(class_count_) * SafeInt<size_t>(num_features);
-  if (coefficients_.size() < expected_coefficients_size) {
+  size_t expected_coefficients_size = 0;
+  if (!SafeMultiply(static_cast<size_t>(class_count_), static_cast<size_t>(num_features),
+                    expected_coefficients_size)) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                           "LinearClassifier: coefficients length (", coefficients_.size(),
-                           ") is less than classes (", class_count_, ") * features (", num_features, ")");
+                           "class_count (", class_count_, ") * num_features (", num_features,
+                           ") overflows size_t");
   }
+  ORT_RETURN_IF_NOT(coefficients_.size() >= expected_coefficients_size,
+                    "coefficients size (", coefficients_.size(), ") is less than class_count (", class_count_,
+                    ") * num_features (", num_features, ")");
 
   Tensor* Y = ctx->Output(0, {num_batches});
 

onnxruntime/core/providers/cpu/ml/linearregressor.cc

@@ -3,6 +3,7 @@
 
 #include "core/providers/cpu/ml/linearregressor.h"
 #include "core/common/narrow.h"
+#include "core/common/safeint.h"
 #include "core/providers/cpu/math/gemm.h"
 
 namespace onnxruntime {
@@ -86,6 +87,17 @@ Status LinearRegressor::Compute(OpKernelContext* ctx) const {
   ptrdiff_t num_batches = input_shape.NumDimensions() <= 1 ? 1 : narrow<ptrdiff_t>(input_shape[0]);
   ptrdiff_t num_features = input_shape.NumDimensions() <= 1 ? narrow<ptrdiff_t>(input_shape.Size())
                                                             : narrow<ptrdiff_t>(input_shape[1]);
+  size_t expected_coefficients_size = 0;
+  if (!SafeMultiply(static_cast<size_t>(num_targets_), static_cast<size_t>(num_features),
+                    expected_coefficients_size)) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "num_targets (", num_targets_, ") * num_features (", num_features,
+                           ") overflows size_t");
+  }
+  ORT_RETURN_IF_NOT(coefficients_.size() >= expected_coefficients_size,
+                    "coefficients size (", coefficients_.size(), ") is less than num_targets (", num_targets_,
+                    ") * num_features (", num_features, ")");
+
   Tensor& Y = *ctx->Output(0, {num_batches, num_targets_});
   concurrency::ThreadPool* tp = ctx->GetOperatorThreadPool();
 

onnxruntime/core/providers/cpu/ml/svmclassifier.cc

@@ -77,14 +77,24 @@ SVMClassifier::SVMClassifier(const OpKernelInfo& info)
   // out-of-bounds reads from crafted models.
   if (mode_ == SVM_TYPE::SVM_SVC) {
     // SVC mode: coefficients layout is [class_count - 1, vector_count]
-    const size_t expected_coefficients = static_cast<size_t>(class_count_ - 1) * static_cast<size_t>(vector_count_);
+    size_t expected_coefficients = 0;
+    if (!SafeMultiply(static_cast<size_t>(class_count_ - 1), static_cast<size_t>(vector_count_),
+                      expected_coefficients)) {
+      ORT_THROW("class_count - 1 (", class_count_ - 1, ") * vector_count (", vector_count_,
+                ") overflows size_t");
+    }
     ORT_ENFORCE(coefficients_.size() >= expected_coefficients,
                 "coefficients attribute size (", coefficients_.size(),
                 ") is smaller than expected (", expected_coefficients,
                 ") for the given class_count and vector_count.");
 
     // rho needs one entry per classifier pair: class_count * (class_count - 1) / 2
-    const size_t num_classifiers = static_cast<size_t>(class_count_) * static_cast<size_t>(class_count_ - 1) / 2;
+    size_t num_classifiers = 0;
+    if (!SafeMultiply(static_cast<size_t>(class_count_), static_cast<size_t>(class_count_ - 1), num_classifiers)) {
+      ORT_THROW("class_count (", class_count_, ") * (class_count - 1) (", class_count_ - 1,
+                ") overflows size_t");
+    }
+    num_classifiers /= 2;
     ORT_ENFORCE(rho_.size() >= num_classifiers,
                 "rho attribute size (", rho_.size(),
                 ") is smaller than expected (", num_classifiers,
@@ -193,7 +203,40 @@ Status SVMClassifier::ComputeImpl(OpKernelContext& ctx,
                                   gsl::span<const float> x_data, const TensorShape& x_shape) const {
   concurrency::ThreadPool* threadpool = ctx.GetOperatorThreadPool();
 
-  const auto num_batches = SafeInt<int32_t>(x_shape.NumDimensions() == 1 ? 1 : x_shape[0]);
+  const auto input_rank = x_shape.NumDimensions();
+  ORT_RETURN_IF_NOT(input_rank > 0 && input_rank <= 2, "Input shape must have 1 or 2 dimensions. Dims=", input_rank);
+
+  const ptrdiff_t num_batches = SafeInt<ptrdiff_t>(input_rank == 1 ? 1 : x_shape[0]);
+  const ptrdiff_t num_features = input_rank == 1 ? narrow<ptrdiff_t>(x_shape[0])
+                                                 : narrow<ptrdiff_t>(x_shape[1]);
+  ORT_RETURN_IF_NOT(num_features == feature_count_ && num_features >= 0 && num_batches >= 0,
+                    "Invalid input for SVMClassifier: expected feature_count=", feature_count_,
+                    ", actual num_features=", num_features,
+                    ", input_rank=", input_rank,
+                    ", num_batches=", num_batches);
+  if (mode_ == SVM_TYPE::SVM_LINEAR) {
+    size_t expected_linear_size = 0;
+    if (!SafeMultiply(static_cast<size_t>(class_count_), static_cast<size_t>(num_features),
+                      expected_linear_size)) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "class_count (", class_count_, ") * num_features (", num_features,
+                             ") overflows size_t");
+    }
+    ORT_RETURN_IF_NOT(coefficients_.size() >= expected_linear_size,
+                      "coefficients size (", coefficients_.size(), ") is less than class_count (", class_count_,
+                      ") * num_features (", num_features, ")");
+  } else {
+    size_t expected_sv_size = 0;
+    if (!SafeMultiply(static_cast<size_t>(vector_count_), static_cast<size_t>(num_features),
+                      expected_sv_size)) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "vector_count (", vector_count_, ") * num_features (", num_features,
+                             ") overflows size_t");
+    }
+    ORT_RETURN_IF_NOT(support_vectors_.size() >= expected_sv_size,
+                      "support_vectors size (", support_vectors_.size(), ") is less than vector_count (", vector_count_,
+                      ") * num_features (", num_features, ")");
+  }
 
   // Total number of classifiers comparing pairs between the classes
   // e.g. if you have A, B C and D classes, the number of classifiers to compare between each pair is 6
@@ -229,7 +272,7 @@ Status SVMClassifier::ComputeImpl(OpKernelContext& ctx,
   std::vector<float> probsp2_data;
 
   if (mode_ == SVM_TYPE::SVM_SVC && have_proba) {
-    probsp2_data.resize(num_batches * class_count_squared, 0.f);
+    probsp2_data.resize(SafeInt<size_t>(num_batches) * class_count_squared, 0.f);
   }
 
   int write_additional_scores = -1;
@@ -261,7 +304,7 @@ Status SVMClassifier::ComputeImpl(OpKernelContext& ctx,
     if (have_proba) {
       // we will write num_batches * num_classifiers scores first, and transform those to num_batches * class_count_,
       // so need to use a separate buffer for the first scoring.
-      classifier_scores_data.resize(num_batches * num_classifiers);
+      classifier_scores_data.resize(SafeInt<size_t>(num_batches) * num_classifiers);
       classifier_scores = gsl::make_span<float>(classifier_scores_data.data(), classifier_scores_data.size());
     } else {
       // we will write directly to the final scores buffer

onnxruntime/core/providers/cpu/ml/svmregressor.cc

@@ -56,9 +56,17 @@ template <typename T>
 Status SVMRegressor<T>::Compute(OpKernelContext* ctx) const {
   const auto* X = ctx->Input<Tensor>(0);
 
-  ptrdiff_t num_features = X->Shape().NumDimensions() == 1 ? narrow<ptrdiff_t>(X->Shape()[0]) : narrow<ptrdiff_t>(X->Shape()[1]);
-  ptrdiff_t num_batches = X->Shape().NumDimensions() == 1 ? 1 : narrow<ptrdiff_t>(X->Shape()[0]);
-  ORT_RETURN_IF_NOT(num_features == feature_count_ && num_features >= 0 && num_batches >= 0, "Invalid argument");
+  const auto& x_shape = X->Shape();
+  const auto input_rank = x_shape.NumDimensions();
+  ORT_RETURN_IF_NOT(input_rank > 0 && input_rank <= 2, "Input shape must have 1 or 2 dimensions. Dims=", input_rank);
+
+  ptrdiff_t num_features = input_rank == 1 ? narrow<ptrdiff_t>(x_shape[0]) : narrow<ptrdiff_t>(x_shape[1]);
+  ptrdiff_t num_batches = input_rank == 1 ? 1 : narrow<ptrdiff_t>(x_shape[0]);
+  ORT_RETURN_IF_NOT(num_features == feature_count_ && num_features >= 0 && num_batches >= 0,
+                    "Invalid input for SVMRegressor: expected feature_count=", feature_count_,
+                    ", actual num_features=", num_features,
+                    ", input_rank=", input_rank,
+                    ", num_batches=", num_batches);
 
   // X: [num_batches, feature_count_] where features could be coefficients or support vectors
   // coefficients_: [vector_count_]

onnxruntime/test/providers/cpu/ml/linearclassifer_test.cc

@@ -126,6 +126,20 @@ TEST(MLOpTest, LinearClassifierBinaryWithLabels) {
   test.Run();
 }
 
+TEST(MLOpTest, LinearClassifierInvalidCoefficientsSize) {
+  OpTester test("LinearClassifier", 1, onnxruntime::kMLDomain);
+
+  test.AddAttribute("coefficients", std::vector<float>{1.f, 2.f});
+  test.AddAttribute("intercepts", std::vector<float>{0.f, 0.f});
+  test.AddAttribute("classlabels_ints", std::vector<int64_t>{0, 1});
+
+  test.AddInput<float>("X", {1, 2}, {1.f, 2.f});
+  test.AddOutput<int64_t>("Y", {1}, {0});
+  test.AddOutput<float>("Z", {1, 2}, {0.f, 0.f});
+
+  test.Run(OpTester::ExpectResult::kExpectFailure, "coefficients size");
+}
+
 template <typename T>
 void LinearClassifierMulticlass() {
   OpTester test("LinearClassifier", 1, onnxruntime::kMLDomain);
@@ -185,7 +199,7 @@ TEST(MLOpTest, LinearClassifierInvalidCoefficientsSizeFails) {
   test.AddOutput<float>("Z", {1, 3}, {0.f, 0.f, 0.f});
 
   test.Run(OpTester::ExpectResult::kExpectFailure,
-           "LinearClassifier: coefficients length (3) is less than classes (3) * features (2)");
+           "coefficients size (3) is less than class_count (3) * num_features (2)");
 }
 
 // Regression test: coefficients not divisible by class_count.

onnxruntime/test/providers/cpu/ml/linearregressor_test.cc

@@ -72,6 +72,19 @@ TEST_P(LinearRegressorTest, LinearRegressorUniTarget) {
   test.Run();
 }
 
+TEST(MLOpTest, LinearRegressorInvalidCoefficientsSize) {
+  OpTester test("LinearRegressor", 1, onnxruntime::kMLDomain);
+
+  test.AddAttribute("coefficients", std::vector<float>{1.f, 2.f});
+  test.AddAttribute("intercepts", std::vector<float>{0.f, 0.f});
+  test.AddAttribute("targets", static_cast<int64_t>(2));
+
+  test.AddInput<float>("X", {1, 2}, {1.f, 2.f});
+  test.AddOutput<float>("Y", {1, 2}, {0.f, 0.f});
+
+  test.Run(OpTester::ExpectResult::kExpectFailure, "coefficients size");
+}
+
 // For PROBIT, all the output values are NaN.
 INSTANTIATE_TEST_SUITE_P(
     LinearRegressorTest, LinearRegressorTest,

onnxruntime/test/providers/cpu/ml/svmclassifier_test.cc

@@ -4,6 +4,9 @@
 #include "gtest/gtest.h"
 #include "test/providers/provider_test_utils.h"
 
+#include <cstddef>
+#include <limits>
+
 namespace onnxruntime {
 namespace test {
 
@@ -294,6 +297,51 @@ TEST(MLOpTest, SVMClassifierUndersizedCoefficients) {
   test.Run(OpTester::ExpectResult::kExpectFailure, "coefficients attribute size");
 }
 
+TEST(MLOpTest, SVMClassifierInvalidInputFeatureCount) {
+  OpTester test("SVMClassifier", 1, onnxruntime::kMLDomain);
+
+  std::vector<float> coefficients = {1.f, 1.f, 1.f, 1.f};
+  std::vector<float> support_vectors = {0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f};
+  std::vector<float> rho = {0.1f, 0.1f, 0.1f};
+  std::vector<float> kernel_params = {0.01f, 0.f, 3.f};
+  std::vector<int64_t> classes = {0, 1, 2};
+  std::vector<int64_t> vectors_per_class = {1, 1, 0};
+
+  test.AddAttribute("kernel_type", std::string("RBF"));
+  test.AddAttribute("coefficients", coefficients);
+  test.AddAttribute("support_vectors", support_vectors);
+  test.AddAttribute("vectors_per_class", vectors_per_class);
+  test.AddAttribute("rho", rho);
+  test.AddAttribute("kernel_params", kernel_params);
+  test.AddAttribute("classlabels_ints", classes);
+
+  test.AddInput<float>("X", {1, 3}, {0.f, 0.f, 0.f});
+  test.AddOutput<int64_t>("Y", {1}, {1});
+  test.AddOutput<float>("Z", {1, 3}, {0.f, 0.f, 0.f});
+
+  test.Run(OpTester::ExpectResult::kExpectFailure, "Invalid input for SVMClassifier");
+}
+
+TEST(MLOpTest, SVMClassifierCoefficientsSizeOverflow) {
+  OpTester test("SVMClassifier", 1, onnxruntime::kMLDomain);
+
+  const int64_t max_vector_count = static_cast<int64_t>(std::numeric_limits<ptrdiff_t>::max());
+
+  test.AddAttribute("kernel_type", std::string("RBF"));
+  test.AddAttribute("coefficients", std::vector<float>{1.f});
+  test.AddAttribute("support_vectors", std::vector<float>{1.f});
+  test.AddAttribute("vectors_per_class", std::vector<int64_t>{max_vector_count, 0, 0, 0});
+  test.AddAttribute("rho", std::vector<float>{0.f});
+  test.AddAttribute("kernel_params", std::vector<float>{0.01f, 0.f, 3.f});
+  test.AddAttribute("classlabels_ints", std::vector<int64_t>{0, 1, 2, 3});
+
+  test.AddInput<float>("X", {1, 1}, {0.f});
+  test.AddOutput<int64_t>("Y", {1}, {0});
+  test.AddOutput<float>("Z", {1, 4}, {0.f, 0.f, 0.f, 0.f});
+
+  test.Run(OpTester::ExpectResult::kExpectFailure, "overflows size_t");
+}
+
 TEST(MLOpTest, SVMClassifierUndersizedRho) {
   OpTester test("SVMClassifier", 1, onnxruntime::kMLDomain);