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 6191529458dc9b1e6df1a42e454ce87cba9bd724 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Sat, 4 Apr 2026 01:45:40 -0500 Subject: [PATCH 2/2] Use SafeInt for overflow-safe coefficients size validation Address review feedback: the multiplication of num_targets * num_features could overflow with very large dimensions. Use SafeInt for checked multiplication and also reject negative num_targets. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/providers/cpu/ml/linearregressor.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/cpu/ml/linearregressor.cc b/onnxruntime/core/providers/cpu/ml/linearregressor.cc index 1ad5619466321..0c5b6580274bc 100644 --- a/onnxruntime/core/providers/cpu/ml/linearregressor.cc +++ b/onnxruntime/core/providers/cpu/ml/linearregressor.cc @@ -89,7 +89,9 @@ 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)) { + // Use SafeInt for overflow-safe multiplication. + if (num_targets_ < 0 || + coefficients_.size() != SafeInt(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_,