From af0aea3a1e018c055f749177b1688d73877630bf Mon Sep 17 00:00:00 2001 From: lucylq Date: Mon, 21 Jul 2025 17:08:47 -0700 Subject: [PATCH] Integer overflow in HierarchicalAllocator::get_offset_address() ^ Test on top of D78676341. Differential Revision: [D78703809](https://our.internmc.facebook.com/intern/diff/D78703809/) [ghstack-poisoned] --- 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,