onnxruntime
a91b0b49 - Add missing rank/length validation to pooling operators (fix OOB reads) (#29579)

Commit
3 days ago
Add missing rank/length validation to pooling operators (fix OOB reads) (#29579) ## 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`. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading