Add missing rank/length validation to pooling operators (fix OOB reads)

[Ti-Tai Wang (titaiwangms)]

Summary

The pooling operators (MaxPool / AveragePool / LpPool and the contrib NhwcMaxPool) did not fully validate attribute lengths and input ranks before indexing into per-spatial-dimension attribute vectors. On malformed models this could lead to out-of-bounds reads. This PR adds the missing input-validation guards so such models are rejected with clear errors instead of reading past the end of the attribute containers.

Because all standard pooling execution providers (CPU / CUDA / ROCm / JS / WebGPU / XNNPACK / ACL) share PoolBase / PoolAttributes, the pool_attributes.h guards cover them all with a single change. NhwcMaxPool needs its own guard because it does not go through PoolAttributes::SetOutputSize and hand-writes its spatial loop.

What changed, by file

onnxruntime/core/providers/cpu/nn/pool_attributes.h

• Validate pads.size() == 2 * kernel_shape.size() before indexing pads (F1).
• Validate that the kernel_shape rank matches the input spatial rank in InferOutputSize before indexing strides / kernel_shape / dilations (F2).
• Validate input rank >= 2 in SetOutputSize before indexing the input shape (F5).
• Remove the int truncation of the spatial dimension; use int64 throughout (F3).
• Wrap the residual output-size arithmetic in SafeInt<int64_t> to guard against overflow and enforce out_size >= 0 (F4).
• Give the strides length check an explicit error message (previously a bare condition).
• Validate MaxPool storage_order is 0 (row-major) or 1 (column-major) at construction (F7).

onnxruntime/contrib_ops/cpu/quantization/nhwc_max_pool.cc

• NhwcMaxPool bypasses the shared SetOutputSize path and hand-writes its spatial loop, indexing kernel_shape / strides / dilations by dimension. Added a guard validating that the kernel_shape rank matches the input spatial rank before the loop. Compute returns Status, so this uses ORT_RETURN_IF_NOT (no-exception-build compatible), mirroring the existing guard in fp16_pool.cc. The PoolAttributes constructor already enforces that strides and dilations match kernel_shape, so guarding kernel_shape here covers all three indexed vectors.

onnxruntime/core/providers/cpu/fp16/fp16_pool.cc

• Fixed an incorrect loop bound in the malformed-kernel_shape diagnostic path: the error-message
  loop iterated with TensorShape::Size() (the element count, product of dims) while indexing the
  dimension array, which is bounded by the rank. Replaced it with NumDimensions() so the loop
  iterates the dims correctly. This PoolFp16 kernel is ARM64-only (built under
  MLAS_F16VEC_INTRINSICS_SUPPORTED, i.e. non-Apple ARM64/ARM64EC) and the affected code runs only
  on the already-failing error path. Because the kernel is not compiled on x86 and the existing
  fp16 pool tests are gated on USE_CUDA/USE_COREML (exercising those EPs' kernels rather than
  this one), no x86 CI negative test can reach this path; the fix is validated by inspection.

Tests

Added 11 negative tests total, all using kExpectFailure and compatible with no-exception builds:

• pool_op_test.cc: MaxPool_PadsTooShort, AveragePool_PadsTooShort, LpPool_PadsTooShort, LpPool_KernelRankMismatch, AveragePool_NegativeOutputDim, MaxPool_PadsTooLong, MaxPool_StridesLengthMismatch, MaxPool_DilationsLengthMismatch, MaxPool_InputRankTooLow, MaxPool_InvalidStorageOrder.
• nhwc_maxpool_op_test.cc: NhwcMaxPool KernelRankMismatch_S8.

Design note

The pads.size() == 2 * kernel_shape.size() check is intentionally unconditional (not gated on auto_pad). Making it conditional on auto_pad == NOTSET would reopen the out-of-bounds indexing path for models that set auto_pad together with a malformed pads list.

Testing

All pooling and NHWC max-pool tests pass: 105 passed / 2 skipped (DML not built) / 0 failed, with zero regression.

Follow-ups (out of scope, tracked separately)

• F6: auto_pad vs explicit pads mutual-exclusion per the ONNX spec — deferred due to backward-compatibility risk.
• CUDA narrow_cast<int> truncation on int64 dims — non-OOB.
• Optional integer-domain ceil to remove a float path in ComputeOutputSize.

[Copilot]

Pull request overview

This PR hardens pooling operator input/attribute validation to prevent out-of-bounds reads on malformed models by adding rank/length checks in the shared PoolAttributes path and in the contrib NhwcMaxPool kernel, and it adds negative tests to pin the new failure modes.

Changes:

• Added strict attribute-length and rank validation in PoolAttributes (e.g., pads length, kernel_shape rank vs input spatial rank), plus overflow-checked output-size arithmetic.
• Added a kernel-rank guard in contrib NhwcMaxPool prior to its per-spatial-dimension loop.
• Added multiple negative tests covering malformed pooling attributes and rank mismatches.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/core/providers/cpu/nn/pool_attributes.h	Adds missing pooling attribute/rank validation and safer output-size arithmetic.
onnxruntime/contrib_ops/cpu/quantization/nhwc_max_pool.cc	Adds kernel/input rank validation before indexing per-dimension attribute vectors.
onnxruntime/test/providers/cpu/nn/pool_op_test.cc	Adds negative tests for malformed pooling attributes and rank issues.
onnxruntime/test/contrib_ops/nhwc_maxpool_op_test.cc	Adds a negative test for NHWC max-pool kernel rank mismatch.

[Ti-Tai Wang (titaiwangms)]

Review synthesis (multi-reviewer)

Verdict: solid, mergeable. Deep spec/math review validated the core design; no true Critical/Major issues in the changed code. Findings below are deduped and verified against the code.

What's well done (verified)

• Unconditional pads.size() == 2 * kernel_shape.size() check is correct. Confirmed against the ONNX spec: gating it on auto_pad == NOTSET would reopen a genuine operator[] OOB read in the constructor loop (pads[dim + rank]). The design note is right.
• SafeInt rewrite is bit-for-bit identical to the original for valid inputs; the out_size >= 0 enforce provably cannot reject a valid model (the ceil-mode decrement cannot manufacture a negative dim).
• Shared SetOutputSize guard propagates protection to CUDA/WebGPU/ACL/QLinear with no parallel patches needed; NhwcMaxPool correctly gets its own guard because it bypasses InferOutputSize.

In-scope, worth fixing before merge (Minor/Nit)

1. Self-contradicting test comment — the MaxPool_StridesLengthMismatch comment says "the bare ORT_ENFORCE carries the stringified condition as its message," but this PR gave that check a real message ("Strides dimensions should match kernel shape"). The new comment is factually wrong.
2. Overstated SafeInt comment — "SafeInt guards against int64 overflow in the output-size arithmetic" overstates coverage: the auto_pad=SAME_* arithmetic in ComputeSizePadDilations and the numerator / stride division are not SafeInt-wrapped. Narrow the comment (or extend SafeInt to those sites).
3. Internal "F5" label leaks into two test comments (MaxPool_InputRankTooLow); one also references SetOutputSize while actually pinning the NhwcMaxPool::Compute message. Describe guards by location, not private labels.
4. (Optional) ORT_ENFORCE(input_dims.size() >= 2) in InferOutputSize has no message, unlike every other guard added here.

Follow-ups (out of scope — track separately, not blockers)

• onnxruntime/core/providers/cpu/fp16/fp16_pool.cc:80 — a real OOB read of the same class this PR targets (pre-existing; not touched here). The malformed-rank error path loops for (int64_t i = 0; i < input_shape.Size(); i++) (element count) and indexes input_shape[i], valid only up to the rank. Should be NumDimensions(). Worth a follow-up fix.
• ceil_mode uses float for int64 output-size arithmetic (precision loss beyond 2^24) — pre-existing, not worsened by this PR.
• XNNPACK constructs PoolAttributes in GetCapability (max_pool.cc, average_pool.cc, unwrapped), so the new enforces can throw during partitioning. This is pre-existing throw behavior, fires only on malformed models (which this PR intends to reject), is caught at Initialize, and replaces a prior OOB read with a clean failure. Low priority; optional defense-in-depth is to guard the XNNPACK constructors or make that construction exception-safe.

Note: an earlier reviewer flagged the negative-numerator integer-truncation-vs-floor path as a regression; on deeper review this is pre-existing behavior the PR does not change, and the new out_size >= 0 enforce is a strict improvement.

[github-actions]

You can commit the suggested changes from lintrunner.

[Ti-Tai Wang (titaiwangms)]

cc Akshay Sonawane (@apsonawane)