[folly] Add additional overflow checks to IOBuf - CVE-2021-24036

Summary:
As per title

CVE-2021-24036

Reviewed By: jan

Differential Revision: D27938605

fbshipit-source-id: 7481c54ae6fbb7b67b15b3631d5357c2f7043f9c

folly/io/IOBuf.cpp

@@ -23,6 +23,7 @@
 #include <cassert>
 #include <cstdint>
 #include <cstdlib>
+#include <limits>
 #include <stdexcept>
 
 #include <folly/Conv.h>
@@ -32,6 +33,7 @@
 #include <folly/hash/SpookyHashV2.h>
 #include <folly/io/Cursor.h>
 #include <folly/lang/Align.h>
+#include <folly/lang/CheckedMath.h>
 #include <folly/lang/Exception.h>
 #include <folly/memory/Malloc.h>
 #include <folly/memory/SanitizeAddress.h>
@@ -84,7 +86,8 @@ enum : std::size_t {
   // benchmarks of real applications to see if adjusting this number makes a
   // difference.  Callers that know their exact use case can also explicitly
   // call createCombined() or createSeparate().)
-  kDefaultCombinedBufSize = 1024
+  kDefaultCombinedBufSize = 1024,
+  kMaxIOBufSize = std::numeric_limits<size_t>::max() >> 1,
 };
 
 // Helper function for IOBuf::takeOwnership()
@@ -187,6 +190,9 @@ void IOBuf::SharedInfo::releaseStorage(SharedInfo* info) noexcept {
 }
 
 void* IOBuf::operator new(size_t size) {
+  if (size > kMaxIOBufSize) {
+    throw_exception<std::bad_alloc>();
+  }
   size_t fullSize = offsetof(HeapStorage, buf) + size;
   auto storage = static_cast<HeapStorage*>(checkedMalloc(fullSize));
 
@@ -297,6 +303,10 @@ IOBuf::IOBuf(
     : IOBuf(op, br.data(), br.size(), headroom, minTailroom) {}
 
 unique_ptr<IOBuf> IOBuf::create(std::size_t capacity) {
+  if (capacity > kMaxIOBufSize) {
+    throw_exception<std::bad_alloc>();
+  }
+
   // For smaller-sized buffers, allocate the IOBuf, SharedInfo, and the buffer
   // all with a single allocation.
   //
@@ -328,6 +338,10 @@ unique_ptr<IOBuf> IOBuf::create(std::size_t capacity) {
 }
 
 unique_ptr<IOBuf> IOBuf::createCombined(std::size_t capacity) {
+  if (capacity > kMaxIOBufSize) {
+    throw_exception<std::bad_alloc>();
+  }
+
   // To save a memory allocation, allocate space for the IOBuf object, the
   // SharedInfo struct, and the data itself all with a single call to malloc().
   size_t requiredStorage = offsetof(HeapFullStorage, align) + capacity;
@@ -456,6 +470,10 @@ unique_ptr<IOBuf> IOBuf::takeOwnership(
     void* userData,
     bool freeOnError,
     TakeOwnershipOption option) {
+  if (capacity > kMaxIOBufSize) {
+    throw_exception<std::bad_alloc>();
+  }
+
   // do not allow only user data without a freeFn
   // since we use that for folly::sizedFree
 
@@ -1006,8 +1024,13 @@ void IOBuf::decrementRefcount() noexcept {
 }
 
 void IOBuf::reserveSlow(std::size_t minHeadroom, std::size_t minTailroom) {
-  size_t newCapacity = (size_t)length_ + minHeadroom + minTailroom;
-  DCHECK_LT(newCapacity, UINT32_MAX);
+  size_t newCapacity = length_;
+  if (!checked_add(&newCapacity, newCapacity, minHeadroom) ||
+      !checked_add(&newCapacity, newCapacity, minTailroom) ||
+      newCapacity > kMaxIOBufSize) {
+    // overflow
+    throw_exception<std::bad_alloc>();
+  }
 
   // reserveSlow() is dangerous if anyone else is sharing the buffer, as we may
   // reallocate and free the original buffer.  It should only ever be called if
@@ -1158,6 +1181,10 @@ void IOBuf::allocExtBuffer(
     uint8_t** bufReturn,
     SharedInfo** infoReturn,
     std::size_t* capacityReturn) {
+  if (minCapacity > kMaxIOBufSize) {
+    throw_exception<std::bad_alloc>();
+  }
+
   size_t mallocSize = goodExtBufferSize(minCapacity);
   auto buf = static_cast<uint8_t*>(checkedMalloc(mallocSize));
   initExtBuffer(buf, mallocSize, infoReturn, capacityReturn);
@@ -1173,6 +1200,10 @@ void IOBuf::allocExtBuffer(
 }
 
 size_t IOBuf::goodExtBufferSize(std::size_t minCapacity) {
+  if (minCapacity > kMaxIOBufSize) {
+    throw_exception<std::bad_alloc>();
+  }
+
   // Determine how much space we should allocate.  We'll store the SharedInfo
   // for the external buffer just after the buffer itself.  (We store it just
   // after the buffer rather than just before so that the code can still just