Fix heap-buffer-overflow in constant_pad_nd (pytorch#18018)

Summary:
Fix write-heap-buffer-overflow in set_all_to_value triggered via
apply_padding_to_dim, reported by fuzzer (T258811544).

Root causes:
1. Negative padding values silently cast to huge size_t, causing massive
out-of-bounds writes.
2. When out_data advances past out_data_end, the remaining computation
(out_data_end - out_data) wraps around to a huge size_t, causing bounds
checks to incorrectly pass.
3. No error propagation after recursive apply_padding_to_dim calls,
allowing the loop to continue writing after a child call has failed.

Fixes:
- Validate all padding values are non-negative in
check_constant_pad_args.
- Read padding as int64_t and explicitly check >= 0 before casting to
size_t.
- Guard remaining computation with out_data <= out_data_end check at all
three bounds-check sites to prevent size_t wraparound.
- Check ctx.failure_state() after recursive calls and bail out early.
- Remove dead pad_i >= 0 check (always true for size_t).


Differential Revision: D95762335

Author: psiddh

kernels/portable/cpu/op_constant_pad_nd.cpp

@@ -51,9 +51,17 @@ void apply_padding_to_dim(
 
   size_t pad_before = 0;
   size_t pad_after = 0;
-  if (pad_i >= 0 && pad_i < pad.size() / 2) {
-    pad_before = pad[2 * pad_i];
-    pad_after = pad[2 * pad_i + 1];
+  if (pad_i < pad.size() / 2) {
+    int64_t pb = pad[2 * pad_i];
+    int64_t pa = pad[2 * pad_i + 1];
+    ET_KERNEL_CHECK_MSG(
+        ctx,
+        pb >= 0 && pa >= 0,
+        InvalidArgument,
+        /* void */,
+        "Padding values must be non-negative.");
+    pad_before = static_cast<size_t>(pb);
+    pad_after = static_cast<size_t>(pa);
   }
 
   size_t out_step_len = out_strides[dim];
@@ -62,6 +70,12 @@ void apply_padding_to_dim(
   // Do not copy padding beyond the out tensor bounds.
   // Use division to avoid potential overflow in multiplication.
   if (pad_before > 0) {
+    ET_KERNEL_CHECK_MSG(
+        ctx,
+        out_data <= out_data_end,
+        InvalidArgument,
+        /* void */,
+        "Out data pointer exceeds buffer bounds.");
     size_t remaining = out_data_end - out_data;
     ET_KERNEL_CHECK_MSG(
         ctx,
@@ -92,7 +106,12 @@ void apply_padding_to_dim(
           /* void */,
           "Out tensor overlaps with the input tensor. This is not supported.");
       // Bounds check before memcpy
-      // Use overflow-safe check for remaining >= copy_len
+      ET_KERNEL_CHECK_MSG(
+          ctx,
+          out_data <= out_data_end,
+          InvalidArgument,
+          /* void */,
+          "Out data pointer exceeds buffer bounds.");
       size_t remaining = out_data_end - out_data;
       ET_KERNEL_CHECK_MSG(
           ctx,
@@ -123,6 +142,10 @@ void apply_padding_to_dim(
           last_padded_dim,
           dim + 1);
 
+      if (ctx.failure_state() != Error::Ok) {
+        return;
+      }
+
       out_data += out_step_len;
       self_data += in_step_len;
     }
@@ -131,6 +154,12 @@ void apply_padding_to_dim(
   // Do not copy padding beyond the out tensor bounds.
   // Use division to avoid potential overflow in multiplication.
   if (pad_after > 0) {
+    ET_KERNEL_CHECK_MSG(
+        ctx,
+        out_data <= out_data_end,
+        InvalidArgument,
+        /* void */,
+        "Out data pointer exceeds buffer bounds.");
     size_t remaining = out_data_end - out_data;
     ET_KERNEL_CHECK_MSG(
         ctx,

kernels/portable/cpu/util/kernel_ops_util.cpp

@@ -564,6 +564,14 @@ bool check_constant_pad_args(
       pad.size() / 2,
       in.dim());
 
+  for (size_t i = 0; i < pad.size(); ++i) {
+    ET_CHECK_OR_RETURN_FALSE(
+        pad[i] >= 0,
+        "Padding values must be non-negative, but got pad[%zu] = %" PRId64,
+        i,
+        pad[i]);
+  }
+
   return true;
 }