Fix linearregressor heap oob read by bmehta001 · Pull Request #27964 · microsoft/onnxruntime

bmehta001 · 2026-04-03T08:27:19Z

Description

Motivation and Context

The condition 'axis < -tensorRank && axis >= tensorRank' is always false (a number cannot simultaneously be less than -r and >= r). Changed to '||' so out-of-range axes are properly rejected. Affects both onnxjs and jsep utility implementations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

After normalizing negative axes, validate that each axis is within [0, input_rank). Without this check, out-of-range axes could cause out-of-bounds access on the input shape array. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

When split sizes come from a runtime tensor input (not an attribute), they were not validated for non-negative values. Crafted negative split sizes that sum to the correct total could pass the sum check and create tensors with negative dimensions. The attribute path already had this validation. Added non-negative check in both CPU PrepareForCompute and CUDA PrepareForComputeLocal. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add checks that output dimensions are not negative after computing them from user-supplied pads/slices. Without these checks, excessive negative padding or slicing could produce tensors with negative dimensions, leading to buffer underflows or undefined behavior. Zero-sized dimensions are valid (empty tensors) and are not rejected. Affected operators and providers: - CPU Pad: both reshaped and true output dimensions - CUDA Pad: output dimensions in compute loop - WebGPU Pad: output dimensions in compute loop - WebGPU Slice: output dimensions after ceil computation - JS onnxjs Slice (WebGL): output shape after ends - starts - JS onnxjs padShape utility - JS jsep padShape utility - JS jsep Slice (WebGPU): output shape after ceil computation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

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>

Copilot

Pull request overview

This PR hardens several operators against invalid shapes/attributes that could lead to out-of-bounds reads (notably LinearRegressor), and adds additional runtime validation in CPU/CUDA/WebGPU and JS implementations.

Changes:

Add runtime validation for LinearRegressor that coefficients matches targets * num_features, plus a regression test for the mismatch case.
Add additional validation in Slice/Pad (negative output dims, axis range) and Split (non-negative split sizes) across multiple execution providers/backends.
Fix JS normalizeAxis() logic and add extra shape validation in JS pad/slice utilities.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
onnxruntime/core/providers/cpu/ml/linearregressor.cc	Reject mismatched coefficients size to prevent GEMM backend OOB reads.
onnxruntime/test/providers/cpu/ml/linearregressor_test.cc	Add regression test covering invalid coefficients size.
onnxruntime/core/providers/webgpu/tensor/slice.cc	Add axis range checking and output-dimension validation in WebGPU Slice.
onnxruntime/core/providers/webgpu/tensor/pad.cc	Add output-dimension negative check in WebGPU Pad.
onnxruntime/core/providers/cuda/tensor/split.cc	Reject negative values in runtime `split` input.
onnxruntime/core/providers/cuda/tensor/pad.cc	Add output-dimension negative check in CUDA Pad.
onnxruntime/core/providers/cpu/tensor/split.h	Reject negative values in runtime `split` input (CPU).
onnxruntime/core/providers/cpu/tensor/pad.cc	Add output-dimension negative checks in CPU Pad (reshaped + final dims).
js/web/lib/wasm/jsep/webgpu/ops/slice.ts	Add JS WebGPU slice output-dimension validation.
js/web/lib/wasm/jsep/util.ts	Fix axis normalization condition; add negative-dimension check in pad shape.
js/web/lib/onnxjs/util.ts	Fix axis normalization condition; add negative-dimension check in pad shape.
js/web/lib/onnxjs/backends/webgl/ops/slice.ts	Add JS WebGL slice output-dimension validation.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-03T18:50:12Z

js/web/lib/wasm/jsep/webgpu/ops/slice.ts

    outputShape[axis] = Math.ceil((ends[axis] - starts[axis]) / steps[axis]);
+    if (outputShape[axis] < 0) {
+      throw new Error(
+        `Output dimension ${axis} is negative (${outputShape[axis]}). This may occur from excessive negative slicing.`,
+      );
+    }


The new negative-dimension check will throw for valid slices where end < start with a positive step (which should produce an empty dimension of 0 per Slice semantics). Instead of throwing, clamp the computed dimension to 0 (e.g., Math.max(0, ...)) so empty outputs are handled correctly.

Suggested change

outputShape[axis] = Math.ceil((ends[axis] - starts[axis]) / steps[axis]);

if (outputShape[axis] < 0) {

throw new Error(

`Output dimension ${axis} is negative (${outputShape[axis]}). This may occur from excessive negative slicing.`,

);

}

outputShape[axis] = Math.max(0, Math.ceil((ends[axis] - starts[axis]) / steps[axis]));

Copilot · 2026-04-03T18:50:13Z

js/web/lib/onnxjs/backends/webgl/ops/slice.ts

    outputShape[normalizedAxes[i]] = ends[i] - starts[i];
+    if (outputShape[normalizedAxes[i]] < 0) {
+      throw new Error(
+        `Output dimension ${normalizedAxes[i]} is negative (${outputShape[normalizedAxes[i]]}). This may occur from excessive negative slicing.`,
+      );
+    }


outputShape[axis] = ends[i] - starts[i] can be negative for cases where end < start, which should result in an empty dimension (0) rather than an exception. Please clamp the dimension to 0 so valid empty slices don’t fail.

Suggested change

outputShape[normalizedAxes[i]] = ends[i] - starts[i];

if (outputShape[normalizedAxes[i]] < 0) {

throw new Error(

`Output dimension ${normalizedAxes[i]} is negative (${outputShape[normalizedAxes[i]]}). This may occur from excessive negative slicing.`,

);

}

outputShape[normalizedAxes[i]] = Math.max(0, ends[i] - starts[i]);

Copilot · 2026-04-03T18:50:13Z

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

+  // 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<size_t>(num_targets_) * static_cast<size_t>(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, ")");
+  }


This check multiplies num_targets_ * num_features after casting to size_t. For very large dimensions this can overflow and potentially let an invalid coefficients buffer pass validation. Consider using SafeInt (or equivalent checked multiplication) when computing the expected size, and also validating num_targets_ is non-negative.

Copilot · 2026-04-03T18:50:13Z

onnxruntime/core/providers/cpu/tensor/split.h

+      ORT_RETURN_IF_NOT(std::all_of(split_sizes.cbegin(), split_sizes.cend(),
+                                    [](int64_t value) { return value >= 0; }),
+                        "Invalid value in 'split' input. All values must be >= 0");


This new std::all_of usage relies on transitive includes, but this header doesn’t include <algorithm>. Please add #include <algorithm> here to avoid build breaks (especially in configurations where common.h isn’t included, e.g. SHARED_PROVIDER).

Copilot · 2026-04-03T18:50:14Z

onnxruntime/core/providers/webgpu/tensor/slice.cc

+    ORT_RETURN_IF(val < 0 || val >= static_cast<int64_t>(input_rank),
+                  "Axis value ", axes_raw[i], " is out of range for input rank ", input_rank);


ORT_RETURN_IF returns an ONNXRUNTIME::FAIL status. Since an out-of-range axis is invalid user input, it would be more accurate to return INVALID_ARGUMENT (e.g., via ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, ...)) so callers can distinguish bad input from runtime failures.

Suggested change

ORT_RETURN_IF(val < 0 || val >= static_cast<int64_t>(input_rank),

"Axis value ", axes_raw[i], " is out of range for input rank ", input_rank);

if (val < 0 || val >= static_cast<int64_t>(input_rank)) {

return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,

"Axis value ", axes_raw[i], " is out of range for input rank ", input_rank);

}

Copilot · 2026-04-03T18:50:14Z

onnxruntime/core/providers/webgpu/tensor/slice.cc

    auto temp = static_cast<int64_t>(ceil(1.0 * (end - start) / static_cast<float>(step)));
    output_dims[axes[i]] = (temp > 0 && dim_value != 0) ? temp : 0;
+    ORT_RETURN_IF(output_dims[axes[i]] < 0,
+                  "Output dimension ", axes[i], " is negative (", output_dims[axes[i]],
+                  "). This may occur from excessive negative slicing.");


output_dims[axes[i]] is assigned using (temp > 0 && dim_value != 0) ? temp : 0, so it cannot be negative here. This makes the new negative-dimension ORT_RETURN_IF effectively unreachable; consider either removing it or validating/clamping temp before the assignment if there’s a real negative-dimension scenario to detect.

Copilot · 2026-04-03T18:50:14Z

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

+    ORT_RETURN_IF(reshaped_output_dims[i] < 0,
+                  "Output dimension ", i, " is negative (", reshaped_output_dims[i],
+                  "). This may occur from excessive negative padding/slicing.");
  }

  // Compute true output dimensions
  for (size_t i = 0; i < data_rank; i++) {
    output_dims[i] += SafeInt<int64_t>(pads[i]) + pads[i + data_rank] + slices[i] + slices[i + data_rank];
+    ORT_RETURN_IF(output_dims[i] < 0,
+                  "Output dimension ", i, " is negative (", output_dims[i],
+                  "). This may occur from excessive negative padding/slicing.");


Like the other new dimension validation, this uses ORT_RETURN_IF, which reports ONNXRUNTIME::FAIL. Since negative output dimensions come from invalid pad/slice inputs, consider returning INVALID_ARGUMENT to make the error classification consistent with other input-validation failures in Pad (e.g., reflect-mode checks).

Suggested change

ORT_RETURN_IF(reshaped_output_dims[i] < 0,

"Output dimension ", i, " is negative (", reshaped_output_dims[i],

"). This may occur from excessive negative padding/slicing.");

}

// Compute true output dimensions

for (size_t i = 0; i < data_rank; i++) {

output_dims[i] += SafeInt<int64_t>(pads[i]) + pads[i + data_rank] + slices[i] + slices[i + data_rank];

ORT_RETURN_IF(output_dims[i] < 0,

"Output dimension ", i, " is negative (", output_dims[i],

"). This may occur from excessive negative padding/slicing.");

if (reshaped_output_dims[i] < 0) {

return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,

"Output dimension ", i, " is negative (", reshaped_output_dims[i],

"). This may occur from excessive negative padding/slicing.");

}

}

// Compute true output dimensions

for (size_t i = 0; i < data_rank; i++) {

output_dims[i] += SafeInt<int64_t>(pads[i]) + pads[i + data_rank] + slices[i] + slices[i + data_rank];

if (output_dims[i] < 0) {

return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,

"Output dimension ", i, " is negative (", output_dims[i],

"). This may occur from excessive negative padding/slicing.");

}

Copilot · 2026-04-03T18:50:14Z

onnxruntime/core/providers/cuda/tensor/pad.cc

+    ORT_RETURN_IF(output_dims[i] < 0,
+                  "Output dimension ", i, " is negative (", output_dims[i],
+                  "). This may occur from excessive negative padding/slicing.");


This new negative-dimension validation uses ORT_RETURN_IF (which returns ONNXRUNTIME::FAIL). For invalid pad/slice inputs, returning INVALID_ARGUMENT would be more accurate and consistent with other input checks.

Suggested change

ORT_RETURN_IF(output_dims[i] < 0,

"Output dimension ", i, " is negative (", output_dims[i],

"). This may occur from excessive negative padding/slicing.");

if (output_dims[i] < 0) {

return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,

"Output dimension ", i, " is negative (", output_dims[i],

"). This may occur from excessive negative padding/slicing.");

}

Copilot · 2026-04-03T18:50:15Z

onnxruntime/core/providers/webgpu/tensor/pad.cc

+    ORT_RETURN_IF(output_dims[i] < 0,
+                  "Output dimension ", i, " is negative (", output_dims[i],
+                  "). This may occur from excessive negative padding/slicing.");


This new negative-dimension validation uses ORT_RETURN_IF (which returns ONNXRUNTIME::FAIL). Since this is invalid user input, consider returning INVALID_ARGUMENT for clearer error classification.

Suggested change

ORT_RETURN_IF(output_dims[i] < 0,

"Output dimension ", i, " is negative (", output_dims[i],

"). This may occur from excessive negative padding/slicing.");

if (output_dims[i] < 0) {

return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,

"Output dimension ", i, " is negative (", output_dims[i],

"). This may occur from excessive negative padding/slicing.");

}

bmehta001 and others added 5 commits April 2, 2026 16:01

tianleiwu requested a review from Copilot April 3, 2026 18:44

Copilot started reviewing on behalf of tianleiwu April 3, 2026 18:45 View session

Copilot AI reviewed Apr 3, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix linearregressor heap oob read#27964

Fix linearregressor heap oob read#27964
bmehta001 wants to merge 5 commits intomainfrom
fix/linearregressor-heap-oob-read

bmehta001 commented Apr 3, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Copilot AI Apr 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		ORT_RETURN_IF(val < 0 \|\| val >= static_cast<int64_t>(input_rank),
		"Axis value ", axes_raw[i], " is out of range for input rank ", input_rank);

Conversation

bmehta001 commented Apr 3, 2026

Description

Motivation and Context

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants