From c2543a2e21471ed93b03472bfadaa4f118768f2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alex=20=C5=9Euhan?= Date: Thu, 3 Aug 2017 12:28:17 -0700 Subject: [PATCH] THRIFT-3821 Check for overflow on buffer resize in TMemoryBuffer --- lib/cpp/src/thrift/transport/TBufferTransports.cpp | 6 +++++- lib/cpp/test/TMemoryBufferTest.cpp | 13 +++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/lib/cpp/src/thrift/transport/TBufferTransports.cpp b/lib/cpp/src/thrift/transport/TBufferTransports.cpp index 700bdd527a5..ed5e9272294 100644 --- a/lib/cpp/src/thrift/transport/TBufferTransports.cpp +++ b/lib/cpp/src/thrift/transport/TBufferTransports.cpp @@ -361,9 +361,13 @@ void TMemoryBuffer::ensureCanWrite(uint32_t len) { } // 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::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 1586609bf7d..aa44b16e90d 100644 --- a/lib/cpp/test/TMemoryBufferTest.cpp +++ b/lib/cpp/test/TMemoryBufferTest.cpp @@ -117,4 +117,17 @@ BOOST_AUTO_TEST_CASE(test_exceptions) { 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 small_buff(1); + std::vector 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()