From 0830af8207240df8d7f35b984cdf8bc35d74fa73 Mon Sep 17 00:00:00 2001 From: lucylq Date: Thu, 24 Jul 2025 12:49:31 -0700 Subject: [PATCH] Integer overflow in HierarchicalAllocator::get_offset_address() Differential Revision: D78703809 Pull Request resolved: https://github.com/pytorch/executorch/pull/12699 (cherry picked from commit 2065924e52d025c892a54a6c2372030a742443ed) --- runtime/core/hierarchical_allocator.h | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/runtime/core/hierarchical_allocator.h b/runtime/core/hierarchical_allocator.h index b5031fa38e5..09f4da4daba 100644 --- a/runtime/core/hierarchical_allocator.h +++ b/runtime/core/hierarchical_allocator.h @@ -60,6 +60,15 @@ class HierarchicalAllocator final { uint32_t memory_id, size_t offset_bytes, size_t size_bytes) { + // Check for integer overflow in offset_bytes + size_bytes. + ET_CHECK_OR_RETURN_ERROR( + size_bytes <= SIZE_MAX - offset_bytes, + InvalidArgument, + "Integer overflow in offset_bytes (%" ET_PRIsize_t + ") + size_bytes (%" ET_PRIsize_t ")", + offset_bytes, + size_bytes); + ET_CHECK_OR_RETURN_ERROR( memory_id < buffers_.size(), InvalidArgument, @@ -67,6 +76,7 @@ class HierarchicalAllocator final { memory_id, buffers_.size()); Span buffer = buffers_[memory_id]; + ET_CHECK_OR_RETURN_ERROR( offset_bytes + size_bytes <= buffer.size(), MemoryAllocationFailed,