THRIFT-3821: make memory buffer size configurable so unit test does
not need 2GB to run; add unit test to prove THRIFT-3480
Client: C++
This closes #1369
diff --git a/lib/cpp/src/thrift/transport/TBufferTransports.cpp b/lib/cpp/src/thrift/transport/TBufferTransports.cpp
index ed5e927..1d2692d 100644
--- a/lib/cpp/src/thrift/transport/TBufferTransports.cpp
+++ b/lib/cpp/src/thrift/transport/TBufferTransports.cpp
@@ -364,9 +364,9 @@
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()) {
+ if (new_size > maxBufferSize_) {
throw TTransportException(TTransportException::BAD_ARGS,
- "Internal buffer size exceeded 2GB");
+ "Internal buffer size overflow");
}
avail = available_write() + (new_size - bufferSize_);
}
diff --git a/lib/cpp/src/thrift/transport/TBufferTransports.h b/lib/cpp/src/thrift/transport/TBufferTransports.h
index 37f4596..0a4cc6d 100644
--- a/lib/cpp/src/thrift/transport/TBufferTransports.h
+++ b/lib/cpp/src/thrift/transport/TBufferTransports.h
@@ -448,6 +448,9 @@
private:
// Common initialization done by all constructors.
void initCommon(uint8_t* buf, uint32_t size, bool owner, uint32_t wPos) {
+
+ maxBufferSize_ = std::numeric_limits<uint32_t>::max();
+
if (buf == NULL && size != 0) {
assert(owner);
buf = (uint8_t*)std::malloc(size);
@@ -673,6 +676,29 @@
*/
uint32_t readAll(uint8_t* buf, uint32_t len) { return TBufferBase::readAll(buf, len); }
+ //! \brief Get the current buffer size
+ //! \returns the current buffer size
+ uint32_t getBufferSize() const {
+ return bufferSize_;
+ }
+
+ //! \brief Get the current maximum buffer size
+ //! \returns the current maximum buffer size
+ uint32_t getMaxBufferSize() const {
+ return maxBufferSize_;
+ }
+
+ //! \brief Change the maximum buffer size
+ //! \param[in] maxSize the new maximum buffer size allowed to grow to
+ //! \throws TTransportException(BAD_ARGS) if maxSize is less than the current buffer size
+ void setMaxBufferSize(uint32_t maxSize) {
+ if (maxSize < bufferSize_) {
+ throw TTransportException(TTransportException::BAD_ARGS,
+ "Maximum buffer size would be less than current buffer size");
+ }
+ maxBufferSize_ = maxSize;
+ }
+
protected:
void swap(TMemoryBuffer& that) {
using std::swap;
@@ -705,6 +731,9 @@
// Allocated buffer size
uint32_t bufferSize_;
+ // Maximum allowed size
+ uint32_t maxBufferSize_;
+
// Is this object the owner of the buffer?
bool owner_;
diff --git a/lib/cpp/test/TMemoryBufferTest.cpp b/lib/cpp/test/TMemoryBufferTest.cpp
index aa44b16..0d1d14d 100644
--- a/lib/cpp/test/TMemoryBufferTest.cpp
+++ b/lib/cpp/test/TMemoryBufferTest.cpp
@@ -117,17 +117,52 @@
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);
+BOOST_AUTO_TEST_CASE(test_default_maximum_buffer_size)
+{
+ BOOST_CHECK_EQUAL(std::numeric_limits<uint32_t>::max(), TMemoryBuffer().getMaxBufferSize());
}
-#endif
+
+BOOST_AUTO_TEST_CASE(test_default_buffer_size)
+{
+ BOOST_CHECK_EQUAL(1024, TMemoryBuffer().getBufferSize());
+}
+
+BOOST_AUTO_TEST_CASE(test_error_set_max_buffer_size_too_small)
+{
+ TMemoryBuffer buf;
+ BOOST_CHECK_THROW(buf.setMaxBufferSize(buf.getBufferSize() - 1), TTransportException);
+}
+
+BOOST_AUTO_TEST_CASE(test_maximum_buffer_size)
+{
+ TMemoryBuffer buf;
+ buf.setMaxBufferSize(8192);
+ std::vector<uint8_t> small_buff(1);
+
+ for (size_t i = 0; i < 8192; ++i)
+ {
+ buf.write(&small_buff[0], 1);
+ }
+
+ BOOST_CHECK_THROW(buf.write(&small_buff[0], 1), TTransportException);
+}
+
+BOOST_AUTO_TEST_CASE(test_memory_buffer_to_get_sizeof_objects)
+{
+ // This is a demonstration of how to use TMemoryBuffer to determine
+ // the serialized size of a thrift object in the Binary protocol.
+ // See THRIFT-3480
+
+ shared_ptr<TMemoryBuffer> memBuffer(new TMemoryBuffer());
+ shared_ptr<TBinaryProtocol> binaryProtcol(new TBinaryProtocol(memBuffer));
+
+ thrift::test::Xtruct object;
+ object.i32_thing = 10;
+ object.i64_thing = 30;
+ object.string_thing = "who's your daddy?";
+
+ uint32_t size = object.write(binaryProtcol.get());
+ BOOST_CHECK_EQUAL(47, size);
+}
BOOST_AUTO_TEST_SUITE_END()