Validate Col2Im inputs to prevent heap over-read

[GopalakrishnanN]

Description

The CPU Col2Im kernel derived the output spatial extent from the image_shape and block_shape runtime inputs and dispatched to math::Col2im / math::Col2imNd without checking that the column tensor actually contained the implied number of sliding-block positions.

A crafted model whose image_shape implies more sliding-block positions than the column allocation holds caused the inner loop in math::Col2im to advance past the end of data_col and copy adjacent heap contents into the output tensor, leaking process memory to the caller.

Motivation and Context

Heap buffer over-read in Col2Im reported as an information-disclosure issue. Reproducer: col shape [1, 1, 4] with image_shape = [4, 4] and block_shape = [1, 1] produces a [1, 1, 4, 4] output whose first 4 elements are the input data and whose remaining 12 elements are uninitialized heap values that vary across runs.

Fix

In onnxruntime/core/providers/cpu/tensor/col2im.cc, validate inputs before dispatch:

• image_shape and block_shape are 1-D and have the same length, with at least one spatial dimension.
• All image_shape and block_shape values are positive.
• Padded image extent is at least as large as the dilated kernel in each spatial dimension.
• col (input 0) is rank 3.
• col_shape[1] is a positive multiple of prod(block_shape).
• col_shape[2] equals the expected sliding-block count derived from image_shape, block_shape, pads, strides, and dilations.

Existing ORT_ENFORCE attribute-size checks are converted to ORT_RETURN_IF_NOT so all validation surfaces as a clean Status failure rather than an exception.

Test

Added regression test Col2ImOpTest.ImageShapeLargerThanColumnTensor in onnxruntime/test/providers/cpu/tensor/col2im_test.cc that mirrors the reproducer and expects the kernel to fail with a message identifying the sliding-block mismatch.

Local run of onnxruntime_provider_test --gtest_filter=Col2ImOpTest.* on Windows / Release passes all 7 tests (6 existing + 1 new).

[Copilot]

Pull request overview

This PR hardens the CPU Col2Im kernel by validating runtime shape inputs before dispatching to the underlying math routines, preventing a column-buffer over-read when image_shape implies more sliding blocks than the input tensor contains.

Changes:

• Adds rank, dimensionality, positivity, padded extent, channel/block, and sliding-block count validation in col2im.cc.
• Corrects the N-D Col2imNd output-shape argument to use computed sliding-block extents.
• Adds regression and stride-based 5D coverage for the CPU Col2Im operator.

Reviewed changes

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

File	Description
onnxruntime/core/providers/cpu/tensor/col2im.cc	Adds input validation and passes correct sliding-block dimensions to N-D col2im.
onnxruntime/test/providers/cpu/tensor/col2im_test.cc	Adds tests for 5D strided behavior and the over-read regression case.

---

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

[tianleiwu]

Thanks for the well-targeted hardening fix. The core change correctly closes the heap over-read by validating the column tensor against the sliding-block count derived from image_shape/block_shape/pads/strides/dilations, and it also fixes the output_shape argument passed to Col2imNd (the old adjusted_image_shape_dims ignored pads and strides). Converting the attribute-size ORT_ENFORCE checks to ORT_RETURN_IF_NOT is the right move for untrusted models, and the col_shape[2] check guards both the 2-D and N-D dispatch paths.

One remaining crash vector (inline): the strides and dilations attribute values are never validated as positive, so a crafted model with strides = [0, 0] causes an integer division by zero at the new sliding-block computation. Adding positivity checks would close the last crafted-model DoS in this path.

[tianleiwu]

Thanks for the fix — bounding col_shape[2] against the derived sliding-block count closes the heap over-read cleanly, and the per-dimension strides/dilations positivity checks are now placed before they are used as divisors.

One concrete issue: two of the new tests assert error substrings that do not match the messages the kernel actually emits, so they will fail in CI (details inline).

Nit: the PR description mentions a single new test and "7 tests (6 existing + 1 new)", but the diff adds four tests (WithStrides5dNCHWD, ImageShapeLargerThanColumnTensor, StridesMustBePositive, DilationsMustBePositive). Please update the description and the local count.

[Copilot]

Pull request overview

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

[tianleiwu]

Re-reviewed at head 308d8dd. All previously raised concerns are addressed and the corresponding threads are resolved:

• strides[i]/dilations[i] are now validated as positive in the per-dimension loop before being used as a divisor / kernel-extent factor, closing the SIGFPE-on-crafted-model gap.
• The expected substrings in StridesMustBePositive and DilationsMustBePositive now match the kernel's emitted messages.

The core hardening is sound: the kernel now rejects inputs whose col_shape[2] does not equal the sliding-block count derived from image_shape/block_shape/pads/strides/dilations, plus rank/positivity/multiple-of-block checks, with overflow-safe SafeInt arithmetic. Switching the ND path from adjusted_image_shape_dims (stride=1/pad=0 assumption) to the correctly computed sliding_block_shape_dims is also a genuine correctness fix for strided/padded inputs, and is covered by the new WithStrides5dNCHWD test. The added negative tests give good coverage of each new validation path. Converting the attribute-size ORT_ENFORCE checks to ORT_RETURN_IF_NOT is the right call for surfacing a clean Status. LGTM.

[tianleiwu]

Re-reviewed at head 62a51b28. The final commit aligns the Compute() strides/dilations error messages (All stride values must be positive / All dilation values must be positive) with both the constructor guard and the expected substrings in the new tests, closing the last open thread.

Verification of the core fix:

• N-D path — Swapping adjusted_image_shape_dims (which assumed stride 1 / no pad) for the stride/pad/dilation-aware sliding_block_shape_dims is correct. Col2imNd forwards this as output_shape to Im2col, where it drives the column index (index_col *= output_shape[d_i] and NextPosition(rank, output_shape, ...)). The previous value produced an incorrect column index for strided/padded N-D inputs, so this is a genuine correctness fix in addition to the bounds hardening.
• 2-D path — col_shape[2] == expected_col_blocks exactly matches the output_h * output_w extent that math::Col2im reads internally, so the validated col tensor can no longer be over-read.
• Divisor / overflow safety — strides[i] > 0, dilations[i] > 0, and kernel_shape_size > 0 are all checked before being used as divisors, and padded_extent >= adjusted_kernel prevents a negative sliding-block count. Size arithmetic uses SafeInt.
• Converting the attribute-size ORT_ENFORCE checks to ORT_RETURN_IF_NOT gives clean Status failures consistent with the new validation.

The added regression tests cover the original reproducer plus the strides/dilations/padded-kernel/channel-multiple guards. All previously raised concerns are addressed and resolved. LGTM.

[tianleiwu]

@GopalakrishnanN, you do not need the last commit since the fix is in main branch. Simple fix is to revert it, and the PR is ready to merge (the failed CI is not required right now).