From 623b56db9d1caf53a4c160a1677c7c2eea5fa1b8 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Mon, 22 Jan 2024 15:43:03 -0800
Subject: [PATCH] Check malloc returning null

---
 ConflictSet.cpp | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/ConflictSet.cpp b/ConflictSet.cpp
index ca6de00..e9a616c 100644
--- a/ConflictSet.cpp
+++ b/ConflictSet.cpp
@@ -13,6 +13,13 @@
 #include <utility>
 #include <vector>
 
+__attribute__((always_inline)) void *safe_malloc(size_t s) {
+  if (void *p = malloc(s)) {
+    return p;
+  }
+  abort();
+}
+
 // ==================== BEGIN ARENA IMPL ====================
 
 /// Group allocations with similar lifetimes to amortize the cost of malloc/free
@@ -72,7 +79,6 @@ constexpr inline uint32_t nextPowerOfTwo(uint32_t x) {
   return x <= 1 ? 1 : 1 << (32 - __builtin_clz(x - 1));
 }
 
-/// \private
 struct Arena::ArenaImpl {
   Arena::ArenaImpl *prev;
   int capacity;
@@ -86,7 +92,7 @@ static_assert(alignof(Arena::ArenaImpl) == 8);
 Arena::Arena(int initialSize) : impl(nullptr) {
   if (initialSize > 0) {
     auto allocationSize = align_up(initialSize + sizeof(ArenaImpl), 16);
-    impl = (Arena::ArenaImpl *)malloc(allocationSize);
+    impl = (Arena::ArenaImpl *)safe_malloc(allocationSize);
     impl->prev = nullptr;
     impl->capacity = allocationSize - sizeof(ArenaImpl);
     impl->used = 0;
@@ -122,7 +128,7 @@ void *operator new(size_t size, std::align_val_t align, Arena &arena) {
                                                       arena.impl->capacity * 2)
                                       : 0)),
         16);
-    auto *impl = (Arena::ArenaImpl *)malloc(allocationSize);
+    auto *impl = (Arena::ArenaImpl *)safe_malloc(allocationSize);
     impl->prev = arena.impl;
     impl->capacity = allocationSize - sizeof(Arena::ArenaImpl);
     impl->used = 0;
@@ -158,8 +164,6 @@ template <class T> struct ArenaAlloc {
   }
 
   void deallocate(T *, size_t) noexcept {}
-
-private:
 };
 
 template <class T> using Vector = std::vector<T, ArenaAlloc<T>>;
@@ -393,9 +397,7 @@ private:
   std::span<const uint8_t> bytecode;
 };
 
-inline Arbitrary gArbitrary;
-
-void initFuzz(const uint8_t *data, size_t size);
+Arbitrary gArbitrary;
 
 uint32_t Arbitrary::bounded(uint32_t s) {
   if (s == 1) {
@@ -621,7 +623,7 @@ struct Node {
 Node *createNode(const Key &key, Node *parent, int64_t pointVersion,
                  Random &rand) {
   assert(key.len <= std::numeric_limits<int>::max());
-  Node *result = (Node *)malloc(sizeof(Node) + key.len);
+  Node *result = (Node *)safe_malloc(sizeof(Node) + key.len);
   result->maxVersion = pointVersion;
   result->pointVersion = pointVersion;
   result->child[0] = nullptr;
@@ -1155,7 +1157,7 @@ void ConflictSet::setOldestVersion(int64_t oldestVersion) {
 }
 
 ConflictSet::ConflictSet(int64_t oldestVersion, uint64_t seed)
-    : impl(new(malloc(sizeof(Impl))) Impl{oldestVersion, seed}) {}
+    : impl(new(safe_malloc(sizeof(Impl))) Impl{oldestVersion, seed}) {}
 
 ConflictSet::~ConflictSet() {
   if (impl) {
@@ -1194,7 +1196,7 @@ ConflictSet_setOldestVersion(void *cs, int64_t oldestVersion) {
 }
 __attribute__((__visibility__("default"))) void *
 ConflictSet_create(int64_t oldestVersion, uint64_t seed) {
-  return new (malloc(sizeof(ConflictSet::Impl)))
+  return new (safe_malloc(sizeof(ConflictSet::Impl)))
       ConflictSet::Impl{oldestVersion, seed};
 }
 __attribute__((__visibility__("default"))) void ConflictSet_destroy(void *cs) {