Log and return early if std::malloc returns nullptr

Differential Revision: D86130282

Pull Request resolved: #15532

Author: JacobSzwejbka

extension/memory_allocator/malloc_memory_allocator.h

@@ -58,9 +58,21 @@ class MallocMemoryAllocator : public executorch::runtime::MemoryAllocator {
       // To get higher alignments, allocate extra and then align the returned
       // pointer. This will waste an extra `alignment` 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) {
+        ET_LOG(
+            Error, "Size %zu + alignment %zu would overflow", size, alignment);
+        return nullptr;
+      }
       size += alignment;
     }
-    mem_ptrs_.emplace_back(std::malloc(size));
+    void* mem_ptr = std::malloc(size);
+    if (mem_ptr == nullptr) {
+      ET_LOG(Error, "Failed to allocate %zu bytes", size);
+      return nullptr;
+    }
+    mem_ptrs_.emplace_back(mem_ptr);
     return alignPointer(mem_ptrs_.back(), alignment);
   }
 

extension/memory_allocator/test/malloc_memory_allocator_test.cpp

@@ -131,3 +131,15 @@ TEST_F(MallocMemoryAllocatorTest, ResetSucceeds) {
   EXPECT_NE(p, nullptr);
   EXPECT_ALIGNED(p, kDefaultAlignment);
 }
+
+TEST_F(MallocMemoryAllocatorTest, OverflowDetectionOnSizePlusAlignment) {
+  MallocMemoryAllocator allocator = MallocMemoryAllocator();
+
+  constexpr size_t kLargeAlignment = kDefaultAlignment * 64;
+  constexpr size_t kSizeThatWouldOverflow = SIZE_MAX - kLargeAlignment + 1;
+
+  auto p = allocator.allocate(kSizeThatWouldOverflow, kLargeAlignment);
+
+  // Should return nullptr due to overflow detection.
+  EXPECT_EQ(p, nullptr);
+}