diff --git a/onnxruntime/core/providers/cpu/ml/svmclassifier.cc b/onnxruntime/core/providers/cpu/ml/svmclassifier.cc index 4bfb0f673404a..6725c92b09f82 100644 --- a/onnxruntime/core/providers/cpu/ml/svmclassifier.cc +++ b/onnxruntime/core/providers/cpu/ml/svmclassifier.cc @@ -72,6 +72,40 @@ SVMClassifier::SVMClassifier(const OpKernelInfo& info) ORT_ENFORCE(classlabels_strings_.size() > 0 || classlabels_ints_.size() > 0); ORT_ENFORCE(proba_.size() == probb_.size()); ORT_ENFORCE(coefficients_.size() > 0); + + // Validate attribute array sizes against the declared dimensions to prevent + // 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(class_count_ - 1) * static_cast(vector_count_); + 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(class_count_) * static_cast(class_count_ - 1) / 2; + ORT_ENFORCE(rho_.size() >= num_classifiers, + "rho attribute size (", rho_.size(), + ") is smaller than expected (", num_classifiers, + ") for the given number of classes."); + + // prob_a and prob_b, when provided, need one entry per classifier pair + if (!proba_.empty()) { + ORT_ENFORCE(proba_.size() >= num_classifiers, + "prob_a attribute size (", proba_.size(), + ") is smaller than expected (", num_classifiers, + ") for the given number of classes."); + ORT_ENFORCE(probb_.size() >= num_classifiers, + "prob_b attribute size (", probb_.size(), + ") is smaller than expected (", num_classifiers, + ") for the given number of classes."); + } + } else { + // Linear mode: coefficients layout is [class_count, feature_count] + ORT_ENFORCE(rho_.size() >= 1, "rho attribute must have at least one entry."); + } + weights_are_all_positive_ = std::all_of(coefficients_.cbegin(), coefficients_.cend(), [](float value) { return value >= 0.f; }); } diff --git a/onnxruntime/test/providers/cpu/ml/svmclassifier_test.cc b/onnxruntime/test/providers/cpu/ml/svmclassifier_test.cc index 2fcf86d0447e5..5240c909d2878 100644 --- a/onnxruntime/test/providers/cpu/ml/svmclassifier_test.cc +++ b/onnxruntime/test/providers/cpu/ml/svmclassifier_test.cc @@ -263,5 +263,90 @@ TEST(MLOpTest, SVMClassifierLinear) { test.Run(); } +// 3 classes, 2 support vectors (1 each for first two classes), 4 features. +// Correctly sized attributes: +// coefficients: (class_count-1) * vector_count = 2*2 = 4 +// rho: class_count*(class_count-1)/2 = 3 +// prob_a/prob_b (if present): 3 + +TEST(MLOpTest, SVMClassifierUndersizedCoefficients) { + OpTester test("SVMClassifier", 1, onnxruntime::kMLDomain); + + std::vector coefficients = {1.f, 1.f}; // needs 4, only 2 provided + std::vector support_vectors = {0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f}; + std::vector rho = {0.1f, 0.1f, 0.1f}; // correct size + std::vector kernel_params = {0.01f, 0.f, 3.f}; + std::vector classes = {0, 1, 2}; + std::vector 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("X", {1, 4}, {0.f, 0.f, 0.f, 0.f}); + test.AddOutput("Y", {1}, {1}); + test.AddOutput("Z", {1, 3}, {0.f, 0.f, 0.f}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "coefficients attribute size"); +} + +TEST(MLOpTest, SVMClassifierUndersizedRho) { + OpTester test("SVMClassifier", 1, onnxruntime::kMLDomain); + + std::vector coefficients = {1.f, 1.f, 1.f, 1.f}; // correct size + std::vector support_vectors = {0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f}; + std::vector rho = {0.1f}; // needs 3, only 1 provided + std::vector kernel_params = {0.01f, 0.f, 3.f}; + std::vector classes = {0, 1, 2}; + std::vector 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("X", {1, 4}, {0.f, 0.f, 0.f, 0.f}); + test.AddOutput("Y", {1}, {1}); + test.AddOutput("Z", {1, 3}, {0.f, 0.f, 0.f}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "rho attribute size"); +} + +TEST(MLOpTest, SVMClassifierUndersizedProba) { + OpTester test("SVMClassifier", 1, onnxruntime::kMLDomain); + + std::vector coefficients = {1.f, 1.f, 1.f, 1.f}; // correct size + std::vector support_vectors = {0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f}; + std::vector rho = {0.1f, 0.1f, 0.1f}; // correct size + std::vector proba = {0.5f}; // needs 3, only 1 provided + std::vector probb = {0.5f}; // needs 3, only 1 provided + std::vector kernel_params = {0.01f, 0.f, 3.f}; + std::vector classes = {0, 1, 2}; + std::vector 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("prob_a", proba); + test.AddAttribute("prob_b", probb); + test.AddAttribute("kernel_params", kernel_params); + test.AddAttribute("classlabels_ints", classes); + + test.AddInput("X", {1, 4}, {0.f, 0.f, 0.f, 0.f}); + test.AddOutput("Y", {1}, {1}); + test.AddOutput("Z", {1, 3}, {0.f, 0.f, 0.f}); + + test.Run(OpTester::ExpectResult::kExpectFailure, "prob_a attribute size"); +} + } // namespace test } // namespace onnxruntime