Address Pad Reflect vulnerability

onnxruntime/core/providers/cpu/tensor/pad.cc

@@ -495,9 +495,8 @@ static Status PadImpl(OpKernelContext* ctx,
   }
 
   // Special case for Reflect mode: ensure all extents >= 2 after slicing
-  // otherwise reflection is not possible. Matches numpy behavior as ONNX only
-  // implies that this would be wrong as the start and end positions should be distinct
-  // values and with 0 there is not one, and with 1 reflection degenerates into ambiguity.
+  // otherwise reflection is not possible. Also validate that pads do not
+  // exceed extent - 1 on each side, as required by the ONNX spec.
   if (mode == Mode::Reflect) {
     for (size_t i = 0; i < new_dims_count; ++i) {
       const int64_t extent = effective_input_extents[i];  // length after slicing
@@ -508,6 +507,19 @@ static Status PadImpl(OpKernelContext* ctx,
                                "Pad reflect requires axis length >= 2 after slicing. Input shape:",
                                orig_input_shape);
       }
+      // ONNX spec: reflect pads must not exceed extent - 1 on each side
+      if (reshaped_pad[i] > extent - 1) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Pad reflect: pre-pad (", reshaped_pad[i],
+                               ") exceeds maximum allowed (", extent - 1,
+                               ") for axis ", i, ". Input shape:", orig_input_shape);
+      }
+      if (reshaped_pad[i + new_dims_count] > extent - 1) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Pad reflect: post-pad (", reshaped_pad[i + new_dims_count],
+                               ") exceeds maximum allowed (", extent - 1,
+                               ") for axis ", i, ". Input shape:", orig_input_shape);
+      }
     }
   }
 

onnxruntime/core/providers/cuda/tensor/pad.cc

@@ -195,10 +195,10 @@ Status Pad<T>::ComputeInternal(OpKernelContext* ctx) const {
                            "Pad: invalid mode: ", static_cast<int>(mode_), " with zero effective input extent");
   }
 
-  // Special case for Reflect mode: ensure all extents >= 2 after slicing
-  // otherwise reflection is not possible. Matches numpy behavior as ONNX only
-  // implies that this would be wrong as the start and end positions should be distinct
-  // values and with 0 there is not one, and with 1 reflection degenerates into ambiguity.
+  // Special case for Reflect mode: ensure all extents >= 2 after slicing;
+  // otherwise reflection is not possible. Also validate that pads do not
+  // exceed extent - 1 on each side, as required by the ONNX spec, which
+  // aligns with NumPy behavior where start and end positions must be distinct.
   if (mode_ == Mode::Reflect) {
     for (size_t i = 0; i < dimension_count; ++i) {
       const int64_t extent = effective_input_extents[i];  // length after slicing
@@ -209,6 +209,19 @@ Status Pad<T>::ComputeInternal(OpKernelContext* ctx) const {
                                "Pad reflect requires axis length >= 2 after slicing. Input shape:",
                                input_shape);
       }
+      // ONNX spec: reflect pads must not exceed extent - 1 on each side
+      if ((*p_pads)[i] > extent - 1) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Pad reflect: pre-pad (", (*p_pads)[i],
+                               ") exceeds maximum allowed (", extent - 1,
+                               ") for axis ", i, ". Input shape:", input_shape);
+      }
+      if ((*p_pads)[i + dimension_count] > extent - 1) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Pad reflect: post-pad (", (*p_pads)[i + dimension_count],
+                               ") exceeds maximum allowed (", extent - 1,
+                               ") for axis ", i, ". Input shape:", input_shape);
+      }
     }
   }
 

onnxruntime/test/providers/cpu/tensor/pad_test.cc

@@ -1404,5 +1404,186 @@ TEST(PadOpTest, Pad_Wrap_NegativeFront_PositiveBack) {
   test.RunWithConfig();
 }
 
+// =====================================================================
+// Regression tests for reflect-mode pad-size validation (CVE / heap OOB)
+// ONNX spec: reflect pads must not exceed extent - 1 on each side.
+// =====================================================================
+
+// Bug repro: data_shape=[3], pads=[10,0] — pre-pad 10 > extent-1 (2)
+TEST(PadOpTest, Pad_Reflect_PrePadExceedsExtent_1D) {
+  const std::vector<int64_t> input_shape = {3};
+  const std::vector<float> input_data = {1.0f, 2.0f, 3.0f};
+  const std::vector<int64_t> pads = {10, 0};  // pre=10 > extent-1=2
+
+  // Output dims don't matter because we expect failure before any computation.
+  const std::vector<int64_t> expected_shape = {13};
+  const std::vector<float> expected_data(13, 0.0f);
+
+  OpTester test("Pad", 18);
+  test.AddInput<float>("data", input_shape, input_data);
+  test.AddInput<int64_t>("pads", {static_cast<int64_t>(pads.size())}, pads, true);
+  test.AddOutput<float>("output", expected_shape, expected_data);
+  test.AddAttribute("mode", "reflect");
+  test.ConfigExcludeEps({kDmlExecutionProvider, kQnnExecutionProvider,
+                         kTensorrtExecutionProvider, kWebGpuExecutionProvider});
+  test.Config(OpTester::ExpectResult::kExpectFailure,
+              "Pad reflect: pre-pad");
+  test.RunWithConfig();
+}
+
+// Post-pad exceeds extent - 1
+TEST(PadOpTest, Pad_Reflect_PostPadExceedsExtent_1D) {
+  const std::vector<int64_t> input_shape = {3};
+  const std::vector<float> input_data = {1.0f, 2.0f, 3.0f};
+  const std::vector<int64_t> pads = {0, 10};  // post=10 > extent-1=2
+
+  const std::vector<int64_t> expected_shape = {13};
+  const std::vector<float> expected_data(13, 0.0f);
+
+  OpTester test("Pad", 18);
+  test.AddInput<float>("data", input_shape, input_data);
+  test.AddInput<int64_t>("pads", {static_cast<int64_t>(pads.size())}, pads, true);
+  test.AddOutput<float>("output", expected_shape, expected_data);
+  test.AddAttribute("mode", "reflect");
+  test.ConfigExcludeEps({kDmlExecutionProvider, kQnnExecutionProvider,
+                         kTensorrtExecutionProvider, kWebGpuExecutionProvider});
+  test.Config(OpTester::ExpectResult::kExpectFailure,
+              "Pad reflect: post-pad");
+  test.RunWithConfig();
+}
+
+// Both pre and post exceed extent - 1
+TEST(PadOpTest, Pad_Reflect_BothPadsExceedExtent_1D) {
+  const std::vector<int64_t> input_shape = {3};
+  const std::vector<float> input_data = {1.0f, 2.0f, 3.0f};
+  const std::vector<int64_t> pads = {5, 5};  // both > extent-1=2
+
+  const std::vector<int64_t> expected_shape = {13};
+  const std::vector<float> expected_data(13, 0.0f);
+
+  OpTester test("Pad", 18);
+  test.AddInput<float>("data", input_shape, input_data);
+  test.AddInput<int64_t>("pads", {static_cast<int64_t>(pads.size())}, pads, true);
+  test.AddOutput<float>("output", expected_shape, expected_data);
+  test.AddAttribute("mode", "reflect");
+  test.ConfigExcludeEps({kDmlExecutionProvider, kQnnExecutionProvider,
+                         kTensorrtExecutionProvider, kWebGpuExecutionProvider});
+  // Pre-pad is checked first, so expect that message
+  test.Config(OpTester::ExpectResult::kExpectFailure,
+              "Pad reflect: pre-pad");
+  test.RunWithConfig();
+}
+
+// 2D: pre-pad exceeds extent-1 on one axis only
+TEST(PadOpTest, Pad_Reflect_PrePadExceedsExtent_2D) {
+  const std::vector<int64_t> input_shape = {3, 3};
+  const std::vector<float> input_data = {1, 2, 3, 4, 5, 6, 7, 8, 9};
+  // pads: [start_dim0, start_dim1, end_dim0, end_dim1]
+  // dim0 extent=3 → max pad=2, but we request 5
+  const std::vector<int64_t> pads = {5, 0, 0, 0};
+
+  const std::vector<int64_t> expected_shape = {8, 3};
+  const std::vector<float> expected_data(24, 0.0f);
+
+  OpTester test("Pad", 18);
+  test.AddInput<float>("data", input_shape, input_data);
+  test.AddInput<int64_t>("pads", {static_cast<int64_t>(pads.size())}, pads, true);
+  test.AddOutput<float>("output", expected_shape, expected_data);
+  test.AddAttribute("mode", "reflect");
+  test.ConfigExcludeEps({kDmlExecutionProvider, kQnnExecutionProvider,
+                         kTensorrtExecutionProvider, kWebGpuExecutionProvider});
+  test.Config(OpTester::ExpectResult::kExpectFailure,
+              "Pad reflect: pre-pad");
+  test.RunWithConfig();
+}
+
+// 2D: post-pad exceeds extent-1 on second axis
+TEST(PadOpTest, Pad_Reflect_PostPadExceedsExtent_2D) {
+  const std::vector<int64_t> input_shape = {3, 3};
+  const std::vector<float> input_data = {1, 2, 3, 4, 5, 6, 7, 8, 9};
+  // dim1 extent=3 → max pad=2, but post-pad on dim1 is 5
+  const std::vector<int64_t> pads = {0, 0, 0, 5};
+
+  const std::vector<int64_t> expected_shape = {3, 8};
+  const std::vector<float> expected_data(24, 0.0f);
+
+  OpTester test("Pad", 18);
+  test.AddInput<float>("data", input_shape, input_data);
+  test.AddInput<int64_t>("pads", {static_cast<int64_t>(pads.size())}, pads, true);
+  test.AddOutput<float>("output", expected_shape, expected_data);
+  test.AddAttribute("mode", "reflect");
+  test.ConfigExcludeEps({kDmlExecutionProvider, kQnnExecutionProvider,
+                         kTensorrtExecutionProvider, kWebGpuExecutionProvider});
+  test.Config(OpTester::ExpectResult::kExpectFailure,
+              "Pad reflect: post-pad");
+  test.RunWithConfig();
+}
+
+// Boundary: pad == extent - 1 should SUCCEED (max legal value)
+TEST(PadOpTest, Pad_Reflect_PadEqualsExtentMinus1_Succeeds) {
+  const std::vector<int64_t> input_shape = {3};
+  const std::vector<float> input_data = {1.0f, 2.0f, 3.0f};
+  // extent=3, extent-1=2 -> pad=2 is the maximum legal value
+  const std::vector<int64_t> pads = {2, 2};
+
+  const std::vector<int64_t> expected_shape = {7};
+  // reflect of [1,2,3] with pre=2, post=2:  3,2, 1,2,3, 2,1
+  const std::vector<float> expected_data = {3.0f, 2.0f, 1.0f, 2.0f, 3.0f, 2.0f, 1.0f};
+
+  OpTester test("Pad", 18);
+  test.AddInput<float>("data", input_shape, input_data);
+  test.AddInput<int64_t>("pads", {static_cast<int64_t>(pads.size())}, pads, true);
+  test.AddOutput<float>("output", expected_shape, expected_data);
+  test.AddAttribute("mode", "reflect");
+  test.ConfigExcludeEps({kDmlExecutionProvider, kQnnExecutionProvider,
+                         kTensorrtExecutionProvider, kWebGpuExecutionProvider});
+  test.RunWithConfig();
+}
+
+// Boundary: pad == extent is one past the legal limit → should FAIL
+TEST(PadOpTest, Pad_Reflect_PadEqualsExtent_Fails) {
+  const std::vector<int64_t> input_shape = {3};
+  const std::vector<float> input_data = {1.0f, 2.0f, 3.0f};
+  // extent=3 -> pad=3 exceeds the limit of 2
+  const std::vector<int64_t> pads = {3, 0};
+
+  const std::vector<int64_t> expected_shape = {6};
+  const std::vector<float> expected_data(6, 0.0f);
+
+  OpTester test("Pad", 18);
+  test.AddInput<float>("data", input_shape, input_data);
+  test.AddInput<int64_t>("pads", {static_cast<int64_t>(pads.size())}, pads, true);
+  test.AddOutput<float>("output", expected_shape, expected_data);
+  test.AddAttribute("mode", "reflect");
+  test.ConfigExcludeEps({kDmlExecutionProvider, kQnnExecutionProvider,
+                         kTensorrtExecutionProvider, kWebGpuExecutionProvider});
+  test.Config(OpTester::ExpectResult::kExpectFailure,
+              "Pad reflect: pre-pad");
+  test.RunWithConfig();
+}
+
+// Negative slice + positive pad: extent after slicing is 2, pad=2 > extent-1=1 → FAIL
+TEST(PadOpTest, Pad_Reflect_SlicedExtentExceeded) {
+  const std::vector<int64_t> input_shape = {4};
+  const std::vector<float> input_data = {1.0f, 2.0f, 3.0f, 4.0f};
+  // slice -2 from start → effective extent = 2, extent-1 = 1
+  // then pre-pad 2 > 1 → must fail
+  const std::vector<int64_t> pads = {-2, 4};  // net: -2 + 4 = +2, but reflect pad 4 > extent-1=1
+
+  const std::vector<int64_t> expected_shape = {6};
+  const std::vector<float> expected_data(6, 0.0f);
+
+  OpTester test("Pad", 18);
+  test.AddInput<float>("data", input_shape, input_data);
+  test.AddInput<int64_t>("pads", {static_cast<int64_t>(pads.size())}, pads, true);
+  test.AddOutput<float>("output", expected_shape, expected_data);
+  test.AddAttribute("mode", "reflect");
+  test.ConfigExcludeEps({kDmlExecutionProvider, kQnnExecutionProvider,
+                         kTensorrtExecutionProvider, kWebGpuExecutionProvider});
+  test.Config(OpTester::ExpectResult::kExpectFailure,
+              "Pad reflect: post-pad");
+  test.RunWithConfig();
+}
+
 }  // namespace test
 }  // namespace onnxruntime