WebGPU Pad kernel int64 / int32 truncation can read to oob read

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

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <limits>
 #include <string>
 #include <vector>
 
@@ -66,8 +67,12 @@ Status Pad::ComputeInternal(ComputeContext& context) const {
   for (size_t i = 0; i < dimension_count; i++) {
     int64_t lower_pad = (*p_pads)[i] + (*p_slices)[i];
     int64_t upper_pad = (*p_pads)[i + dimension_count] + (*p_slices)[i + dimension_count];
+    ORT_RETURN_IF_NOT(lower_pad >= std::numeric_limits<int32_t>::min() &&
+                          lower_pad <= std::numeric_limits<int32_t>::max(),
+                      "WebGPU Pad only supports lower pads in the int32 range. Got ", lower_pad,
+                      " for axis ", i);
     lower_pads[i] = static_cast<int32_t>(lower_pad);
-    output_dims[i] += lower_pad + upper_pad;
+    output_dims[i] += SafeInt<int64_t>(lower_pad) + upper_pad;
   }
   TensorShape output_shape(output_dims);
 

onnxruntime/core/providers/webgpu/tensor/pad.wgsl.template

@@ -62,6 +62,10 @@ $MAIN {
         in_coord = output_index - lower_pads;
     }
 
+#if pad_mode == PAD_MODE_WRAP
+    in_coord = ((in_coord % data_shape) + data_shape) % data_shape;
+#endif
+
     input_index += select(u32(in_coord)
 #if output.rank > 1
         * getElementAt(uniforms.data_stride, dim, output.rank - 1)

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

@@ -219,6 +219,41 @@ TYPED_TEST(PadOpTest, Pad_Wrap_1D) {
                                "wrap");
 }
 
+#ifdef USE_WEBGPU
+TEST(PadOpTest, Pad_Wrap_WebGpu_LowerPadExceedsInt32Fails) {
+  if (DefaultWebGpuExecutionProvider().get() == nullptr) {
+    GTEST_SKIP() << "WebGPU execution provider is not available";
+  }
+
+  OpTester test("Pad", 19);
+  test.AddAttribute("mode", "wrap");
+  test.AddInput<float>("data", {4}, {1.0f, 2.0f, 3.0f, 4.0f});
+  test.AddInput<int64_t>("pads", {2}, {2147483648LL, -2147483647LL}, true);
+  test.AddOutput<float>("output", {5}, {0.0f, 0.0f, 0.0f, 0.0f, 0.0f});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> eps;
+  eps.push_back(DefaultWebGpuExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectFailure,
+           "WebGPU Pad only supports lower pads in the int32 range", {}, nullptr, &eps);
+}
+
+TEST(PadOpTest, Pad_Wrap_WebGpu_PadGreaterThanInputDimension) {
+  if (DefaultWebGpuExecutionProvider().get() == nullptr) {
+    GTEST_SKIP() << "WebGPU execution provider is not available";
+  }
+
+  OpTester test("Pad", 19);
+  test.AddAttribute("mode", "wrap");
+  test.AddInput<float>("data", {3}, {1.0f, 2.0f, 3.0f});
+  test.AddInput<int64_t>("pads", {2}, {5, 0}, true);
+  test.AddOutput<float>("output", {8}, {2.0f, 3.0f, 1.0f, 2.0f, 3.0f, 1.0f, 2.0f, 3.0f});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> eps;
+  eps.push_back(DefaultWebGpuExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps);
+}
+#endif
+
 TYPED_TEST(PadOpTest, Pad_Edge_1D) {
   using T = TypeParam;
   RunAllOpsetAllDomainPadTests<T>({3, 2},