THRIFT-3821 Check for overflow on buffer resize in TMemoryBuffer
Client: C++

This closes #1326
diff --git a/lib/cpp/src/thrift/transport/TBufferTransports.cpp b/lib/cpp/src/thrift/transport/TBufferTransports.cpp
index 700bdd5..ed5e927 100644
--- a/lib/cpp/src/thrift/transport/TBufferTransports.cpp
+++ b/lib/cpp/src/thrift/transport/TBufferTransports.cpp
@@ -361,9 +361,13 @@
   }
 
   // Grow the buffer as necessary.
-  uint32_t new_size = bufferSize_;
+  uint64_t new_size = bufferSize_;
   while (len > avail) {
     new_size = new_size > 0 ? new_size * 2 : 1;
+    if (new_size > std::numeric_limits<uint32_t>::max()) {
+      throw TTransportException(TTransportException::BAD_ARGS,
+                                "Internal buffer size exceeded 2GB");
+    }
     avail = available_write() + (new_size - bufferSize_);
   }
 
diff --git a/lib/cpp/test/TMemoryBufferTest.cpp b/lib/cpp/test/TMemoryBufferTest.cpp
index 1586609..aa44b16 100644
--- a/lib/cpp/test/TMemoryBufferTest.cpp
+++ b/lib/cpp/test/TMemoryBufferTest.cpp
@@ -117,4 +117,17 @@
   BOOST_CHECK_NO_THROW(buf2.write((const uint8_t*)"bar", 3));
 }
 
+#ifndef _WIN32
+// We can't allocate 1 GB of memory in 32-bit environments.
+BOOST_AUTO_TEST_CASE(test_over_two_gb) {
+  TMemoryBuffer buf;
+  std::vector<uint8_t> small_buff(1);
+  std::vector<uint8_t> one_gb(1073741824);
+
+  buf.write(&small_buff[0], small_buff.size());
+  buf.write(&one_gb[0], one_gb.size());
+  BOOST_CHECK_THROW(buf.write(&one_gb[0], one_gb.size()), TTransportException);
+}
+#endif
+
 BOOST_AUTO_TEST_SUITE_END()