Fix integer overflow in GridSample GsReflect causing heap OOB read

Closed
GO
6mo ago
+425-80

Description

The reflection branch of GridSample's PixelAtGrid / PixelAtGrid3D could produce out-of-range indices when grid coordinates were extreme. Two minimal changes:

  1. In GsReflect, widen int n = static_cast<int>(dx / range) to int64_t n = static_cast<int64_t>(dx / range) in both branches. For large (but finite) inputs dx / range can exceed INT_MAX, making the narrow int cast undefined behavior.
  2. In PixelAtGrid / PixelAtGrid3D, clamp the int64_t index returned by GsReflect back into [0, W-1] / [0, H-1] / [0, D-1] before indexing the image buffer. This keeps the final load in bounds as a defense-in-depth measure for any remaining edge case (coordinates ≥ 2^63, NaN, Inf, etc.).

Motivation and Context

Reflection padding indexes the image with the result of GsReflect. The internal narrow int cast on a large floating-point value is undefined behavior and could yield a wild index, leading to an out-of-bounds read on the input image buffer.

Tests

Adds one regression test, test_grid_sample_20_4D_nearest_reflection_extreme_coords, that runs GridSample with mode=nearest and padding_mode=reflection on a 2x2 constant-valued image with grid coordinates of magnitude 1e+10 (well above INT_MAX). It is parameterized over float and MLFloat16 and asserts the kernel returns the constant value, which can only happen if every sampled index is in bounds.

All existing *GridSample* tests continue to pass (102/102 locally).

Files changed

  • onnxruntime/core/providers/cpu/tensor/grid_sample.cc
  • onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.inc