diff --git a/backends/aoti/slim/core/slim_tensor.h b/backends/aoti/slim/core/slim_tensor.h index 5a58508a4a2..623843fee92 100644 --- a/backends/aoti/slim/core/slim_tensor.h +++ b/backends/aoti/slim/core/slim_tensor.h @@ -497,15 +497,17 @@ class SlimTensor { static_cast(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::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 + static T* add_byte_offset_checked(T* data, size_t byte_offset) { + uintptr_t data_int = reinterpret_cast(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( // NOLINT(performance-no-int-to-ptr) + data_offset_int); + } + SlimTensor _clone_impl( c10::IntArrayRef sizes, c10::IntArrayRef strides, diff --git a/backends/aoti/slim/core/test/test_slimtensor_copy.cpp b/backends/aoti/slim/core/test/test_slimtensor_copy.cpp index 6c48689619d..36f95ae73ea 100644 --- a/backends/aoti/slim/core/test/test_slimtensor_copy.cpp +++ b/backends/aoti/slim/core/test/test_slimtensor_copy.cpp @@ -8,9 +8,12 @@ #include +#include + #include #include #include +#include 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 sizes = {2}; + std::vector 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 dst_strides = {std::numeric_limits::max()}; + + Storage src_storage = make_cpu_storage(2 * sizeof(int16_t)); + int16_t* src_data = static_cast(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 // =============================================================================