From c325d759ea41d642dd8207d96e8ed91d865722bd Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 5 Nov 2025 09:44:34 -0800 Subject: [PATCH] More overflow checks in memory allocator. (#15581) Summary: - Safer overflow math for high alignments - Only over-allocate up to alignment - 1 bytes when a higher-than-malloc alignment is requested - Move EXECUTORCH_TRACK_ALLOCATION to after a successful allocation and use the final size Reviewed By: JacobSzwejbka Differential Revision: D86228757 --- .github/workflows/trunk.yml | 2 +- .../malloc_memory_allocator.h | 21 +++--- .../test/malloc_memory_allocator_test.cpp | 2 +- runtime/core/hierarchical_allocator.h | 5 +- runtime/core/memory_allocator.h | 66 +++++++++++-------- 5 files changed, 52 insertions(+), 44 deletions(-) diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml index 9b801801e08..03a13e3717b 100644 --- a/.github/workflows/trunk.yml +++ b/.github/workflows/trunk.yml @@ -347,7 +347,7 @@ jobs: elif [[ ${{ matrix.os}} == "zephyr-preset" ]]; then setup_script_args="--target-toolchain zephyr" toolchain_prefix=arm-zephyr-eabi- - threshold="135240" # 132 KiB + threshold="135656" # 132 KiB toolchain_cmake=examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake else echo "Fail unsupport OS selection ${{ matrix.os }}" diff --git a/extension/memory_allocator/malloc_memory_allocator.h b/extension/memory_allocator/malloc_memory_allocator.h index a5e4c91378f..3dede4ac6fd 100644 --- a/extension/memory_allocator/malloc_memory_allocator.h +++ b/extension/memory_allocator/malloc_memory_allocator.h @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -45,8 +46,6 @@ class MallocMemoryAllocator : public executorch::runtime::MemoryAllocator { * memory alignment size. */ void* allocate(size_t size, size_t alignment = kDefaultAlignment) override { - EXECUTORCH_TRACK_ALLOCATION(prof_id(), size); - if (!isPowerOf2(alignment)) { ET_LOG(Error, "Alignment %zu is not a power of 2", alignment); return nullptr; @@ -56,30 +55,30 @@ class MallocMemoryAllocator : public executorch::runtime::MemoryAllocator { static constexpr size_t kMallocAlignment = alignof(std::max_align_t); if (alignment > kMallocAlignment) { // To get higher alignments, allocate extra and then align the returned - // pointer. This will waste an extra `alignment` bytes every time, but + // pointer. This will waste an extra `alignment - 1` bytes every time, but // this is the only portable way to get aligned memory from the heap. - - // Check for overflow before adding alignment to size - if (size > SIZE_MAX - alignment) { + const size_t extra = alignment - 1; + if ET_UNLIKELY (extra >= SIZE_MAX - size) { ET_LOG( - Error, "Size %zu + alignment %zu would overflow", size, alignment); + Error, "Malloc size overflow: size=%zu + extra=%zu", size, extra); return nullptr; } - size += alignment; + size += extra; } void* mem_ptr = std::malloc(size); - if (mem_ptr == nullptr) { - ET_LOG(Error, "Failed to allocate %zu bytes", size); + if (!mem_ptr) { + ET_LOG(Error, "Malloc failed to allocate %zu bytes", size); return nullptr; } mem_ptrs_.emplace_back(mem_ptr); + EXECUTORCH_TRACK_ALLOCATION(prof_id(), size); return alignPointer(mem_ptrs_.back(), alignment); } // Free up each hosted memory pointer. The memory was created via malloc. void reset() override { for (auto mem_ptr : mem_ptrs_) { - free(mem_ptr); + std::free(mem_ptr); } mem_ptrs_.clear(); } diff --git a/extension/memory_allocator/test/malloc_memory_allocator_test.cpp b/extension/memory_allocator/test/malloc_memory_allocator_test.cpp index 27261eebd92..b54650beda4 100644 --- a/extension/memory_allocator/test/malloc_memory_allocator_test.cpp +++ b/extension/memory_allocator/test/malloc_memory_allocator_test.cpp @@ -66,7 +66,7 @@ TEST_F(MallocMemoryAllocatorTest, SimpleAllocateSucceeds) { EXPECT_ALIGNED(p, kDefaultAlignment); auto p2 = allocator.allocate(16); - EXPECT_NE(p, nullptr); + EXPECT_NE(p2, nullptr); EXPECT_NE(p2, p); EXPECT_ALIGNED(p2, kDefaultAlignment); diff --git a/runtime/core/hierarchical_allocator.h b/runtime/core/hierarchical_allocator.h index b5031fa38e5..95b17704861 100644 --- a/runtime/core/hierarchical_allocator.h +++ b/runtime/core/hierarchical_allocator.h @@ -9,13 +9,10 @@ #pragma once #include + #include #include #include -#include -#include -#include -#include namespace executorch { namespace runtime { diff --git a/runtime/core/memory_allocator.h b/runtime/core/memory_allocator.h index 70cf93b076b..001ebd7ac4f 100644 --- a/runtime/core/memory_allocator.h +++ b/runtime/core/memory_allocator.h @@ -8,16 +8,12 @@ #pragma once -#include #include -#include #include #include #include -#include -#include #include namespace executorch { @@ -59,9 +55,26 @@ class MemoryAllocator { */ MemoryAllocator(uint32_t size, uint8_t* base_address) : begin_(base_address), - end_(base_address + size), + end_( + base_address + ? (UINTPTR_MAX - reinterpret_cast(base_address) >= + size + ? base_address + size + : nullptr) + : nullptr), cur_(base_address), - size_(size) {} + size_(size) { + ET_CHECK_MSG( + base_address || size == 0, + "Base address is null but size=%" PRIu32, + size); + ET_CHECK_MSG( + !base_address || size == 0 || + (UINTPTR_MAX - reinterpret_cast(base_address) >= size), + "Address space overflow in allocator"); + } + + virtual ~MemoryAllocator() = default; /** * Allocates `size` bytes of memory. @@ -74,6 +87,10 @@ class MemoryAllocator { * @retval nullptr Not enough memory, or `alignment` was not a power of 2. */ virtual void* allocate(size_t size, size_t alignment = kDefaultAlignment) { + if ET_UNLIKELY (!begin_ || !end_) { + ET_LOG(Error, "allocate() on zero-capacity allocator"); + return nullptr; + } if (!isPowerOf2(alignment)) { ET_LOG(Error, "Alignment %zu is not a power of 2", alignment); return nullptr; @@ -82,18 +99,17 @@ class MemoryAllocator { // The allocation will occupy [start, end), where the start is the next // position that's a multiple of alignment. uint8_t* start = alignPointer(cur_, alignment); - uint8_t* end = start + size; - - // If the end of this allocation exceeds the end of this allocator, print - // error messages and return nullptr - if (end > end_ || end < start) { + size_t padding = static_cast(start - cur_); + size_t available = static_cast(end_ - cur_); + if ET_UNLIKELY (padding > available || size > available - padding) { ET_LOG( Error, "Memory allocation failed: %zuB requested (adjusted for alignment), %zuB available", - static_cast(end - cur_), - static_cast(end_ - cur_)); + padding + size, + available); return nullptr; } + uint8_t* end = start + size; // Otherwise, record how many bytes were used, advance cur_ to the new end, // and then return start. Note that the number of bytes used is (end - cur_) @@ -144,8 +160,9 @@ class MemoryAllocator { if (overflow) { ET_LOG( Error, - "Failed to allocate list of type %zu: size * sizeof(T) overflowed", - size); + "Failed to allocate list: size(%zu) * sizeof(T)(%zu) overflowed", + size, + sizeof(T)); return nullptr; } return static_cast(this->allocate(bytes_size, alignment)); @@ -171,8 +188,6 @@ class MemoryAllocator { prof_id_ = EXECUTORCH_TRACK_ALLOCATOR(name); } - virtual ~MemoryAllocator() {} - protected: /** * Returns the profiler ID for this allocator. @@ -184,21 +199,18 @@ class MemoryAllocator { /** * Returns true if the value is an integer power of 2. */ - static bool isPowerOf2(size_t value) { - return value > 0 && (value & ~(value - 1)) == value; + static constexpr bool isPowerOf2(size_t value) { + return value && !(value & (value - 1)); } /** * Returns the next alignment for a given pointer. */ - static uint8_t* alignPointer(void* ptr, size_t alignment) { - intptr_t addr = reinterpret_cast(ptr); - if ((addr & (alignment - 1)) == 0) { - // Already aligned. - return reinterpret_cast(ptr); - } - addr = (addr | (alignment - 1)) + 1; - return reinterpret_cast(addr); + static inline uint8_t* alignPointer(void* ptr, size_t alignment) { + uintptr_t address = reinterpret_cast(ptr); + uintptr_t mask = static_cast(alignment - 1); + address = (address + mask) & ~mask; + return reinterpret_cast(address); } private: