From f2b60a1bf57ef1a5a759cb39abe683cd466b0cfc Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Thu, 2 Apr 2026 17:39:23 -0500 Subject: [PATCH 1/2] Validate coefficients size in LinearRegressor to prevent OOB read LinearRegressor treats the coefficients attribute as a [num_targets, num_features] matrix and passes it directly to MLAS SGEMM. However, num_features is derived from the input tensor at runtime, and no validation ensured coefficients.size() == num_targets * num_features. A malformed model could provide fewer coefficients than expected, causing MlasSgemmTransposePackB to read past the buffer boundary. Add a size check after num_features is computed but before the GEMM dispatch to reject mismatched coefficients with a clear error message. Files changed: - onnxruntime/core/providers/cpu/ml/linearregressor.cc - onnxruntime/test/providers/cpu/ml/linearregressor_test.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/providers/cpu/ml/linearregressor.cc | 10 ++++++++++ .../test/providers/cpu/ml/linearregressor_test.cc | 13 +++++++++++++ 2 files changed, 23 insertions(+) diff --git a/onnxruntime/core/providers/cpu/ml/linearregressor.cc b/onnxruntime/core/providers/cpu/ml/linearregressor.cc index 3f877ec7e6d0e..1ad5619466321 100644 --- a/onnxruntime/core/providers/cpu/ml/linearregressor.cc +++ b/onnxruntime/core/providers/cpu/ml/linearregressor.cc @@ -86,6 +86,16 @@ Status LinearRegressor::Compute(OpKernelContext* ctx) const { ptrdiff_t num_batches = input_shape.NumDimensions() <= 1 ? 1 : narrow(input_shape[0]); ptrdiff_t num_features = input_shape.NumDimensions() <= 1 ? narrow(input_shape.Size()) : narrow(input_shape[1]); + + // Coefficients are treated as a [num_targets, num_features] matrix. + // Validate size to prevent out-of-bounds reads in the GEMM backend. + if (coefficients_.size() != static_cast(num_targets_) * static_cast(num_features)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "LinearRegressor: coefficients attribute size (", coefficients_.size(), + ") does not match targets (", num_targets_, + ") * input features (", num_features, ")"); + } + Tensor& Y = *ctx->Output(0, {num_batches, num_targets_}); concurrency::ThreadPool* tp = ctx->GetOperatorThreadPool(); diff --git a/onnxruntime/test/providers/cpu/ml/linearregressor_test.cc b/onnxruntime/test/providers/cpu/ml/linearregressor_test.cc index c6eb3f5f5777f..bd70051b65f06 100644 --- a/onnxruntime/test/providers/cpu/ml/linearregressor_test.cc +++ b/onnxruntime/test/providers/cpu/ml/linearregressor_test.cc @@ -85,5 +85,18 @@ INSTANTIATE_TEST_SUITE_P( LinearRegressorParam("SOFTMAX_ZERO", {3.442477e-14f, 1.f, 1.670142e-05f, 1.f, 1.0f, 0.f}, 2) )); + +// Regression test: coefficients size must match targets * num_features. +// A mismatch previously caused an out-of-bounds read in MLAS SGEMM packing. +TEST(LinearRegressorTest, CoefficientsSizeMismatch) { + OpTester test("LinearRegressor", 1, onnxruntime::kMLDomain); + // 1 coefficient but input has 2 features and targets=1 → expects 2 coefficients + test.AddAttribute("coefficients", std::vector{1.0f}); + test.AddAttribute("targets", int64_t{1}); + test.AddInput("X", {1, 2}, {1.f, 2.f}); + test.AddOutput("Y", {1, 1}, {0.f}); + test.Run(OpTester::ExpectResult::kExpectFailure, "coefficients attribute size"); +} + } // namespace test } // namespace onnxruntime From 494c07086fed9ce2725289ba6398bf96c9fda9c4 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Sat, 4 Apr 2026 01:45:40 -0500 Subject: [PATCH 2/2] Use non-throwing overflow check for coefficients validation Address review feedback: ORT_TRY/ORT_CATCH is a no-op in ORT_NO_EXCEPTIONS builds, so SafeInt overflow would abort the process. Replace with an explicit non-throwing overflow check using std::numeric_limits that works consistently across all build configs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/providers/cpu/ml/linearregressor.cc | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/cpu/ml/linearregressor.cc b/onnxruntime/core/providers/cpu/ml/linearregressor.cc index 1ad5619466321..e5baf0aa1b1ab 100644 --- a/onnxruntime/core/providers/cpu/ml/linearregressor.cc +++ b/onnxruntime/core/providers/cpu/ml/linearregressor.cc @@ -5,6 +5,8 @@ #include "core/common/narrow.h" #include "core/providers/cpu/math/gemm.h" +#include + namespace onnxruntime { namespace ml { @@ -89,7 +91,21 @@ Status LinearRegressor::Compute(OpKernelContext* ctx) const { // Coefficients are treated as a [num_targets, num_features] matrix. // Validate size to prevent out-of-bounds reads in the GEMM backend. - if (coefficients_.size() != static_cast(num_targets_) * static_cast(num_features)) { + if (num_targets_ < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "LinearRegressor: targets attribute must be non-negative, got ", num_targets_); + } + + const auto targets = static_cast(num_targets_); + const auto features = static_cast(num_features); + // Non-throwing overflow check: targets * features would overflow if features > max / targets. + if (targets != 0 && features > std::numeric_limits::max() / targets) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "LinearRegressor: coefficients size overflow for targets (", num_targets_, + ") * input features (", num_features, ")"); + } + + if (coefficients_.size() != targets * features) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "LinearRegressor: coefficients attribute size (", coefficients_.size(), ") does not match targets (", num_targets_,