Skip to content

Fix heap out-of-bounds write in MaxPoolGrad via unchecked indices#27932

Closed
vraspar wants to merge 2 commits intomicrosoft:mainfrom
vraspar:vrparikh/fix-maxpoolgrad-oob-write
Closed

Fix heap out-of-bounds write in MaxPoolGrad via unchecked indices#27932
vraspar wants to merge 2 commits intomicrosoft:mainfrom
vraspar:vrparikh/fix-maxpoolgrad-oob-write

Conversation

@vraspar
Copy link
Copy Markdown
Contributor

@vraspar vraspar commented Apr 1, 2026

Description

Adds bounds validation for index values in MaxPoolGrad::Compute to prevent heap out-of-bounds writes when the indices tensor contains values outside the valid range [0, output_size).

Motivation

The MaxPoolGrad operator uses values from the indices input tensor as direct offsets into the output gradient buffer (dX_data) without validating they are within bounds. A maliciously crafted indices tensor with negative or out-of-range values can write to arbitrary heap memory.

Changes

  • File: orttraining/orttraining/training_ops/cpu/nn/pool_gradient_op.cc
  • Cache dX_shape.Size() into dX_size to avoid repeated calls
  • Validate each indices_data[i] is in [0, dX_size) using ORT_ENFORCE
  • Replace raw pointer arithmetic with array indexing for clarity

This follows the same pattern used in the recent RoiAlign OOB fix (#27543).

Add bounds validation for index values in MaxPoolGrad::Compute to prevent
heap out-of-bounds writes when the indices tensor contains values outside
the valid range [0, output_size). Each index is now validated with
ORT_ENFORCE before being used as an offset into the output buffer.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@vraspar vraspar force-pushed the vrparikh/fix-maxpoolgrad-oob-write branch from 5625230 to 2d90937 Compare April 1, 2026 20:46
@vraspar vraspar requested a review from Copilot April 1, 2026 20:54
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR hardens the training MaxPoolGrad CPU kernel against malicious or malformed indices tensors by adding explicit bounds validation before writing into the dX gradient buffer, and adds unit tests to cover negative and out-of-range indices.

Changes:

  • Add per-element range checks for indices in MaxPoolGrad::Compute and simplify the write path (dX_data[index] += ...).
  • Add CPU unit tests that verify failure on negative and out-of-range indices values.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File Description
orttraining/orttraining/training_ops/cpu/nn/pool_gradient_op.cc Adds index bounds validation to prevent heap out-of-bounds writes in MaxPoolGrad.
orttraining/orttraining/test/training_ops/cpu/nn/pool_gradient_op_test.cc Adds negative/out-of-range indices regression tests for MaxPoolGrad.

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

Return INVALID_ARGUMENT status instead of aborting to avoid DoS when
built with ORT_NO_EXCEPTIONS, where ORT_ENFORCE maps to abort().

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@vraspar
Copy link
Copy Markdown
Contributor Author

vraspar commented Apr 3, 2026

Closing in favor of #27903

@vraspar vraspar closed this Apr 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants