From 20691f1baec03787161bfa3144f48c257d0ab2cc Mon Sep 17 00:00:00 2001 From: lucylq Date: Mon, 21 Jul 2025 17:08:35 -0700 Subject: [PATCH 1/3] Check for buffer overflow in prim_ops::et_copy_index() ^ Test on top of D78676341. Differential Revision: [D78701418](https://our.internmc.facebook.com/intern/diff/D78701418/) [ghstack-poisoned] --- kernels/prim_ops/et_copy_index.cpp | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/kernels/prim_ops/et_copy_index.cpp b/kernels/prim_ops/et_copy_index.cpp index e3d9ae46e54..8ba79165729 100644 --- a/kernels/prim_ops/et_copy_index.cpp +++ b/kernels/prim_ops/et_copy_index.cpp @@ -111,10 +111,21 @@ 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 for overflow in the offset. + size_t offset = index * size_copy_from; + ET_CHECK_MSG( + offset / size_copy_from == static_cast(index), + "Integer overflow: index * size_copy_from calculation overflows."); + + // Check that the destination has enough space for the copy. + 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), - copy_from_ptr, - size_copy_from); + (void*)((uintptr_t)copy_to_ptr + offset), copy_from_ptr, size_copy_from); } } // namespace function From 59816a2b21ec6f6fe358efebeddc507dfc7314c5 Mon Sep 17 00:00:00 2001 From: lucylq Date: Tue, 22 Jul 2025 12:27:01 -0700 Subject: [PATCH 2/3] Update on "Check for buffer overflow in prim_ops::et_copy_index()" ^ Test on top of D78676341. Differential Revision: [D78701418](https://our.internmc.facebook.com/intern/diff/D78701418/) [ghstack-poisoned] --- kernels/prim_ops/et_copy_index.cpp | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/kernels/prim_ops/et_copy_index.cpp b/kernels/prim_ops/et_copy_index.cpp index 8ba79165729..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]; } @@ -112,20 +110,18 @@ void et_copy_index(KernelRuntimeContext& context, EValue** stack) { // successfully resized so we can now copy over the data from // copy_from into the copy_to tensor. - // Check for overflow in the offset. - size_t offset = index * size_copy_from; - ET_CHECK_MSG( - offset / size_copy_from == static_cast(index), - "Integer overflow: index * size_copy_from calculation overflows."); - // 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 + offset), copy_from_ptr, size_copy_from); + // NOLINTNEXTLINE(performance-no-int-to-ptr) + (void*)((uintptr_t)copy_to_ptr + offset), + copy_from_ptr, + size_copy_from); } } // namespace function From 3ce7f725417355c70ad873a1929687f0381f9b84 Mon Sep 17 00:00:00 2001 From: lucylq Date: Wed, 23 Jul 2025 15:43:42 -0700 Subject: [PATCH 3/3] Update on "Check for buffer overflow in prim_ops::et_copy_index()" ^ Test on top of D78676341. Differential Revision: [D78701418](https://our.internmc.facebook.com/intern/diff/D78701418/) [ghstack-poisoned] --- kernels/prim_ops/test/prim_ops_test.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 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.