More overflow checks in memory allocator. (pytorch#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

Author: shoumikhin

.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 }}"

extension/memory_allocator/malloc_memory_allocator.h

@@ -10,6 +10,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
 #include <vector>
 
 #include <executorch/runtime/core/memory_allocator.h>
@@ -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();
   }

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);
 

runtime/core/hierarchical_allocator.h

@@ -9,13 +9,10 @@
 #pragma once
 
 #include <c10/util/irange.h>
+
 #include <executorch/runtime/core/memory_allocator.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/core/span.h>
-#include <executorch/runtime/platform/assert.h>
-#include <executorch/runtime/platform/compiler.h>
-#include <executorch/runtime/platform/log.h>
-#include <cstdint>
 
 namespace executorch {
 namespace runtime {

runtime/core/memory_allocator.h

@@ -8,16 +8,12 @@
 
 #pragma once
 
-#include <stdio.h>
 #include <cinttypes>
-#include <cstdint>
 
 #include <c10/util/safe_numerics.h>
 
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/platform/assert.h>
-#include <executorch/runtime/platform/compiler.h>
-#include <executorch/runtime/platform/log.h>
 #include <executorch/runtime/platform/profiler.h>
 
 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<uintptr_t>(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<uintptr_t>(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<size_t>(start - cur_);
+    size_t available = static_cast<size_t>(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<size_t>(end - cur_),
-          static_cast<size_t>(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<T*>(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<intptr_t>(ptr);
-    if ((addr & (alignment - 1)) == 0) {
-      // Already aligned.
-      return reinterpret_cast<uint8_t*>(ptr);
-    }
-    addr = (addr | (alignment - 1)) + 1;
-    return reinterpret_cast<uint8_t*>(addr);
+  static inline uint8_t* alignPointer(void* ptr, size_t alignment) {
+    uintptr_t address = reinterpret_cast<uintptr_t>(ptr);
+    uintptr_t mask = static_cast<uintptr_t>(alignment - 1);
+    address = (address + mask) & ~mask;
+    return reinterpret_cast<uint8_t*>(address);
   }
 
  private: