From c5c0457c6bac384cc706bd2e5629c08d6677f705 Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Wed, 23 Jul 2025 15:43:42 -0700
Subject: [PATCH] Check for buffer overflow in prim_ops::et_copy_index()

Pull Request resolved: https://github.com/pytorch/executorch/pull/12697

^
Test on top of D78676341.
ghstack-source-id: 298161226
@exported-using-ghexport

Differential Revision: [D78701418](https://our.internmc.facebook.com/intern/diff/D78701418/)
---
 kernels/prim_ops/et_copy_index.cpp      | 19 +++++++++++++------
 kernels/prim_ops/test/prim_ops_test.cpp |  2 +-
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/kernels/prim_ops/et_copy_index.cpp b/kernels/prim_ops/et_copy_index.cpp
index e3d9ae46e54..2b3f55bf5a0 100644
--- a/kernels/prim_ops/et_copy_index.cpp
+++ b/kernels/prim_ops/et_copy_index.cpp
@@ -86,11 +86,9 @@ void et_copy_index(KernelRuntimeContext& context, EValue** stack) {
     // If we're copying past the first index then the shape of
     // copy_from and copy_to without the leading dimension should be
     // the same. i.e. copy_to.size[1:] == copy_from.size[:].
-    if (index > 0) {
-      ET_CHECK_MSG(
-          copy_to.sizes()[i + 1] == copy_from.sizes()[i],
-          "Mismatch in shape between copy_to and copy_from tensors");
-    }
+    ET_CHECK_MSG(
+        copy_to.sizes()[i + 1] == copy_from.sizes()[i],
+        "Mismatch in shape between copy_to and copy_from tensors");
     expected_output_size[i + 1] = copy_from.sizes()[i];
   }
 
@@ -111,8 +109,17 @@ void et_copy_index(KernelRuntimeContext& context, EValue** stack) {
   // If we've reached here, it means the copy_to tensor has been
   // successfully resized so we can now copy over the data from
   // copy_from into the copy_to tensor.
+
+  // Check that the destination has enough space for the copy.
+  size_t offset = index * size_copy_from;
+  size_t copy_to_size = copy_to.element_size() * copy_to.numel();
+  ET_CHECK_MSG(
+      offset + size_copy_from <= copy_to_size,
+      "Buffer overflow: copy_to tensor is smaller than copy_from tensor.");
+
   memcpy(
-      (void*)((uintptr_t)copy_to_ptr + index * size_copy_from),
+      // NOLINTNEXTLINE(performance-no-int-to-ptr)
+      (void*)((uintptr_t)copy_to_ptr + offset),
       copy_from_ptr,
       size_copy_from);
 }
diff --git a/kernels/prim_ops/test/prim_ops_test.cpp b/kernels/prim_ops/test/prim_ops_test.cpp
index 5d7429bd0d7..58a247c2c7e 100644
--- a/kernels/prim_ops/test/prim_ops_test.cpp
+++ b/kernels/prim_ops/test/prim_ops_test.cpp
@@ -215,7 +215,7 @@ TEST_F(RegisterPrimOpsTest, TestETCopyIndex) {
   Tensor copy_to = tf.make({2, 2}, {0, 0, 0, 0});
 #else
   std::vector<int> buf(4);
-  SizesType expected_output_size[2] = {0, 0};
+  SizesType expected_output_size[2] = {0, 2};
   Tensor copy_to =
       tf.make({2, 2}, {0, 0, 0, 0}, {}, TensorShapeDynamism::DYNAMIC_BOUND);
   // Resize the tensor to 0 size for the tests.