pointer overflow in slim_tensor (#20134)

backends/aoti/slim/core/slim_tensor.h

@@ -497,15 +497,17 @@ class SlimTensor {
                   static_cast<size_t>(dst_offset), elem_size, &dst_byte_offset),
           "copy_: byte offset overflow");
 
-      // Copy elem_size bytes from src to dst
+      char* dst_byte_offset_ptr =
+          add_byte_offset_checked(dst_data, dst_byte_offset);
+      const char* src_byte_offset_ptr =
+          add_byte_offset_checked(src_data, src_byte_offset);
       if (this->device().is_cpu() && other.device().is_cpu()) {
-        std::memcpy(
-            dst_data + dst_byte_offset, src_data + src_byte_offset, elem_size);
+        std::memcpy(dst_byte_offset_ptr, src_byte_offset_ptr, elem_size);
       } else if (this->device().is_cuda() || other.device().is_cuda()) {
 #if defined(CUDA_AVAILABLE)
         DeviceTraits<c10::DeviceType::CUDA>::memcpy(
-            dst_data + dst_byte_offset,
-            src_data + src_byte_offset,
+            dst_byte_offset_ptr,
+            src_byte_offset_ptr,
             elem_size,
             device(), // dst device
             other.device() // src device
@@ -555,6 +557,17 @@ class SlimTensor {
   }
 
  private:
+  template <typename T>
+  static T* add_byte_offset_checked(T* data, size_t byte_offset) {
+    uintptr_t data_int = reinterpret_cast<uintptr_t>(data);
+    uintptr_t data_offset_int = 0;
+    ET_CHECK_MSG(
+        !::c10::add_overflows(data_int, byte_offset, &data_offset_int),
+        "copy_: data pointer overflow");
+    return reinterpret_cast<T*>( // NOLINT(performance-no-int-to-ptr)
+        data_offset_int);
+  }
+
   SlimTensor _clone_impl(
       c10::IntArrayRef sizes,
       c10::IntArrayRef strides,

backends/aoti/slim/core/test/test_slimtensor_copy.cpp

@@ -8,9 +8,12 @@
 
 #include <gtest/gtest.h>
 
+#include <limits>
+
 #include <executorch/backends/aoti/slim/core/slim_tensor.h>
 #include <executorch/backends/aoti/slim/core/storage.h>
 #include <executorch/backends/aoti/slim/factory/empty.h>
+#include <executorch/runtime/platform/platform.h>
 
 namespace executorch::backends::aoti::slim {
 
@@ -214,6 +217,42 @@ TEST(SlimTensorCopyTest, CopyNonContiguousDst) {
   EXPECT_FLOAT_EQ(dst_data[5], 5.0f);
 }
 
+// =============================================================================
+// Overflow Validation Tests
+// =============================================================================
+
+TEST(SlimTensorCopyTest, CopyFailsWhenDataPointerOverflows) {
+  std::vector<int64_t> sizes = {2};
+  std::vector<int64_t> src_strides = {1};
+  // For Short, INT64_MAX elements becomes UINTPTR_MAX - 1 bytes. That
+  // passes byte-offset validation and only fails at checked pointer addition.
+  std::vector<int64_t> dst_strides = {std::numeric_limits<int64_t>::max()};
+
+  Storage src_storage = make_cpu_storage(2 * sizeof(int16_t));
+  int16_t* src_data = static_cast<int16_t*>(src_storage->data());
+  src_data[0] = 1;
+  src_data[1] = 2;
+  SlimTensor src(
+      std::move(src_storage),
+      makeArrayRef(sizes),
+      makeArrayRef(src_strides),
+      c10::ScalarType::Short);
+
+  Storage dst_storage = make_cpu_storage(sizeof(int16_t));
+  SlimTensor dst(
+      std::move(dst_storage),
+      makeArrayRef(sizes),
+      makeArrayRef(dst_strides),
+      c10::ScalarType::Short);
+
+  EXPECT_DEATH(
+      {
+        et_pal_init();
+        dst.copy_(src);
+      },
+      "copy_: data pointer overflow");
+}
+
 // =============================================================================
 // Storage Offset Tests
 // =============================================================================