Check the return value from malloc / realloc.

Summary:
https://github.com/facebook/folly/issues/7

Wrappers: checkedMalloc / checkedRealloc / checkedCalloc

Test Plan: all folly tests

FB internal diff: D486841

diff --git a/folly/Arena.h b/folly/Arena.h
index 8db3ca5dd0b19f28e05bda4868230219e8e9bdbe..eef6fab73b04d02c25fc00e10dc114be335da2f7 100644 (file)
--- a/folly/Arena.h
+++ b/folly/Arena.h
@@ -193,9 +193,7 @@ struct ArenaAllocatorTraits {
 class SysAlloc {
  public:
   void* allocate(size_t size) {
-    void* mem = malloc(size);
-    if (!mem) throw std::bad_alloc();
-    return mem;
+    return checkedMalloc(size);
   }
 
   void deallocate(void* p) {

diff --git a/folly/FBString.h b/folly/FBString.h
index a2c7835d99b3de5dfa17d856f3a4e611a420f8d8..fc3d40a45f63ccb43ee88176e0b34806ee9fdf4b 100644 (file)
--- a/folly/FBString.h
+++ b/folly/FBString.h
@@ -301,7 +301,7 @@ public:
       // one extra Char for the null terminator.
       auto const allocSize =
            goodMallocSize((1 + rhs.ml_.size_) * sizeof(Char));
-      ml_.data_ = static_cast<Char*>(malloc(allocSize));
+      ml_.data_ = static_cast<Char*>(checkedMalloc(allocSize));
       fbstring_detail::pod_copy(rhs.ml_.data_,
                                 // 1 for terminator
                                 rhs.ml_.data_ + rhs.ml_.size_ + 1,
@@ -364,7 +364,7 @@ public:
       // Medium strings are allocated normally. Don't forget to
       // allocate one extra Char for the terminating null.
       auto const allocSize = goodMallocSize((1 + size) * sizeof(Char));
-      ml_.data_ = static_cast<Char*>(malloc(allocSize));
+      ml_.data_ = static_cast<Char*>(checkedMalloc(allocSize));
       fbstring_detail::pod_copy(data, data + size, ml_.data_);
       ml_.size_ = size;
       ml_.capacity_ = (allocSize / sizeof(Char) - 1) | isMedium;
@@ -586,7 +586,7 @@ public:
         // Don't forget to allocate one extra Char for the terminating null
         auto const allocSizeBytes =
           goodMallocSize((1 + minCapacity) * sizeof(Char));
-        auto const data = static_cast<Char*>(malloc(allocSizeBytes));
+        auto const data = static_cast<Char*>(checkedMalloc(allocSizeBytes));
         auto const size = smallSize();
         fbstring_detail::pod_copy(small_, small_ + size + 1, data);
         // No need for writeTerminator(), we wrote it above with + 1.
@@ -742,7 +742,7 @@ private:
       // struct.
       const size_t allocSize = goodMallocSize(
         sizeof(RefCounted) + *size * sizeof(Char));
-      auto result = static_cast<RefCounted*>(malloc(allocSize));
+      auto result = static_cast<RefCounted*>(checkedMalloc(allocSize));
       result->refCount_.store(1, std::memory_order_release);
       *size = (allocSize - sizeof(RefCounted)) / sizeof(Char);
       return result;
@@ -2184,7 +2184,7 @@ getline(
   for (;;) {
     // This looks quadratic but it really depends on realloc
     auto const newSize = size + 128;
-    buf = static_cast<char*>(realloc(buf, newSize));
+    buf = static_cast<char*>(checkedRealloc(buf, newSize));
     is.getline(buf + size, newSize - size, delim);
     if (is.bad() || is.eof() || !is.fail()) {
       // done by either failure, end of file, or normal read

diff --git a/folly/FBVector.h b/folly/FBVector.h
index b9cecd93db68538e46eda6c2deb2897ecffc692e..e09eeda32447f147fb3c0e1fcd799571b913968e 100644 (file)
--- a/folly/FBVector.h
+++ b/folly/FBVector.h
@@ -280,7 +280,7 @@ public:
     }
 
     auto const nBytes = goodMallocSize(n * sizeof(T));
-    b_ = static_cast<T*>(malloc(nBytes));
+    b_ = static_cast<T*>(checkedMalloc(nBytes));
     fbvector_detail::uninitializedFillDefaultOrFree(b_, n);
     e_ = b_ + n;
     z_ = b_ + nBytes / sizeof(T);
@@ -293,7 +293,7 @@ public:
     }
 
     auto const nBytes = goodMallocSize(n * sizeof(T));
-    b_ = static_cast<T*>(malloc(nBytes));
+    b_ = static_cast<T*>(checkedMalloc(nBytes));
     fbvector_detail::uninitializedFillOrFree(b_, n, value);
     e_ = b_ + n;
     z_ = b_ + nBytes / sizeof(T);
@@ -396,7 +396,7 @@ private:
 
       // Must reallocate - just do it on the side
       auto const nBytes = goodMallocSize(newSize * sizeof(T));
-      auto const b = static_cast<T*>(malloc(nBytes));
+      auto const b = static_cast<T*>(checkedMalloc(nBytes));
       std::uninitialized_copy(first, last, b);
       this->fbvector::~fbvector();
       b_ = b;
@@ -590,7 +590,7 @@ private:
     assert(crtCapacity < n); // reserve_in_place should have taken
                              // care of this
     auto const newCapacityBytes = goodMallocSize(n * sizeof(T));
-    auto b = static_cast<T*>(malloc(newCapacityBytes));
+    auto b = static_cast<T*>(checkedMalloc(newCapacityBytes));
     auto const oldSize = size();
     memcpy(b, b_, oldSize * sizeof(T));
     // Done with the old chunk. Free but don't call destructors!

diff --git a/folly/Malloc.h b/folly/Malloc.h
index 52c4a30a887b5c3fb739a4f5aac687c0c440e7c7..946060abb2a18d63eca842c92c8b68e53d2036e4 100644 (file)
--- a/folly/Malloc.h
+++ b/folly/Malloc.h
@@ -133,6 +133,28 @@ inline size_t goodMallocSize(size_t minSize) {
 // expanded in place, and this constant reflects that.
 static const size_t jemallocMinInPlaceExpandable = 4096;
 
+/**
+ * Trivial wrappers around malloc, calloc, realloc that check for allocation
+ * failure and throw std::bad_alloc in that case.
+ */
+inline void* checkedMalloc(size_t size) {
+  void* p = malloc(size);
+  if (!p) throw std::bad_alloc();
+  return p;
+}
+
+inline void* checkedCalloc(size_t n, size_t size) {
+  void* p = calloc(n, size);
+  if (!p) throw std::bad_alloc();
+  return p;
+}
+
+inline void* checkedRealloc(void* ptr, size_t size) {
+  void* p = realloc(ptr, size);
+  if (!p) throw std::bad_alloc();
+  return p;
+}
+
 /**
  * This function tries to reallocate a buffer of which only the first
  * currentSize bytes are used. The problem with using realloc is that
@@ -162,7 +184,7 @@ inline void* smartRealloc(void* p,
       return p;
     }
     // Cannot expand; must move
-    auto const result = malloc(newCapacity);
+    auto const result = checkedMalloc(newCapacity);
     std::memcpy(result, p, currentSize);
     free(p);
     return result;
@@ -172,13 +194,13 @@ inline void* smartRealloc(void* p,
   auto const slack = currentCapacity - currentSize;
   if (slack * 2 > currentSize) {
     // Too much slack, malloc-copy-free cycle:
-    auto const result = malloc(newCapacity);
+    auto const result = checkedMalloc(newCapacity);
     std::memcpy(result, p, currentSize);
     free(p);
     return result;
   }
   // If there's not too much slack, we realloc in hope of coalescing
-  return realloc(p, newCapacity);
+  return checkedRealloc(p, newCapacity);
 }
 
 #ifdef _LIBSTDCXX_FBSTRING

diff --git a/folly/small_vector.h b/folly/small_vector.h
index 0a75d3ad2ea545ed4673139742cce1f34ddb30e9..5f7e27333dc6be1146bcbd7b9aa8e96d26770156 100644 (file)
--- a/folly/small_vector.h
+++ b/folly/small_vector.h
@@ -994,10 +994,7 @@ private:
       needBytes += kHeapifyCapacitySize;
     }
     auto const sizeBytes = goodMallocSize(needBytes);
-    void* newh = std::malloc(sizeBytes);
-    if (!newh) {
-      throw std::bad_alloc();
-    }
+    void* newh = checkedMalloc(sizeBytes);
     // We expect newh to be at least 2-aligned, because we want to
     // use its least significant bit as a flag.
     assert(!detail::pointerFlagGet(newh));