1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  *   http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
20 #include <array>
21 #include <boost/test/unit_test.hpp>
22 #include <climits>
23 #include <cstdlib>
24 #include <iostream>
25 #include <memory>
26 #include <numeric>
27 #include <thrift/protocol/TBinaryProtocol.h>
28 #include <thrift/transport/TBufferTransports.h>
29 #include <vector>
30 
31 #include "gen-cpp/ThriftTest_types.h"
32 
33 BOOST_AUTO_TEST_SUITE(TMemoryBufferTest)
34 
35 using apache::thrift::protocol::TBinaryProtocol;
36 using apache::thrift::transport::TMemoryBuffer;
37 using apache::thrift::transport::TTransportException;
38 using std::shared_ptr;
39 using std::cout;
40 using std::endl;
41 using std::string;
42 
BOOST_AUTO_TEST_CASE(test_read_write_grow)43 BOOST_AUTO_TEST_CASE(test_read_write_grow) {
44   // Added to test the fix for THRIFT-1248
45   TMemoryBuffer uut;
46   const int maxSize = 65536;
47   uint8_t verify[maxSize];
48   std::vector<uint8_t> buf;
49   buf.resize(maxSize);
50 
51   for (uint32_t i = 0; i < maxSize; ++i) {
52     buf[i] = static_cast<uint8_t>(i);
53   }
54 
55   for (uint32_t i = 1; i < maxSize; i *= 2) {
56     uut.write(&buf[0], i);
57   }
58 
59   for (uint32_t i = 1; i < maxSize; i *= 2) {
60     uut.read(verify, i);
61     BOOST_CHECK_EQUAL(0, ::memcmp(verify, &buf[0], i));
62   }
63 }
64 
BOOST_AUTO_TEST_CASE(test_roundtrip)65 BOOST_AUTO_TEST_CASE(test_roundtrip) {
66   shared_ptr<TMemoryBuffer> strBuffer(new TMemoryBuffer());
67   shared_ptr<TBinaryProtocol> binaryProtcol(new TBinaryProtocol(strBuffer));
68 
69   thrift::test::Xtruct a;
70   a.i32_thing = 10;
71   a.i64_thing = 30;
72   a.string_thing = "holla back a";
73 
74   a.write(binaryProtcol.get());
75   std::string serialized = strBuffer->getBufferAsString();
76 
77   shared_ptr<TMemoryBuffer> strBuffer2(new TMemoryBuffer());
78   shared_ptr<TBinaryProtocol> binaryProtcol2(new TBinaryProtocol(strBuffer2));
79 
80   strBuffer2->resetBuffer((uint8_t*)serialized.data(), static_cast<uint32_t>(serialized.length()));
81   thrift::test::Xtruct a2;
82   a2.read(binaryProtcol2.get());
83 
84   BOOST_CHECK(a == a2);
85 }
86 
BOOST_AUTO_TEST_CASE(test_readAppendToString)87 BOOST_AUTO_TEST_CASE(test_readAppendToString) {
88   string str1 = "abcd1234";
89   TMemoryBuffer buf((uint8_t*)str1.data(),
90                     static_cast<uint32_t>(str1.length()),
91                     TMemoryBuffer::COPY);
92 
93   string str3 = "wxyz", str4 = "6789";
94   buf.readAppendToString(str3, 4);
95   buf.readAppendToString(str4, INT_MAX);
96 
97   BOOST_CHECK(str3 == "wxyzabcd");
98   BOOST_CHECK(str4 == "67891234");
99 }
100 
BOOST_AUTO_TEST_CASE(test_exceptions)101 BOOST_AUTO_TEST_CASE(test_exceptions) {
102   char data[] = "foo\0bar";
103 
104   TMemoryBuffer buf1((uint8_t*)data, 7, TMemoryBuffer::OBSERVE);
105   string str = buf1.getBufferAsString();
106   BOOST_CHECK(str.length() == 7);
107 
108   buf1.resetBuffer();
109 
110   BOOST_CHECK_THROW(buf1.write((const uint8_t*)"foo", 3), TTransportException);
111 
112   TMemoryBuffer buf2((uint8_t*)data, 7, TMemoryBuffer::COPY);
113   BOOST_CHECK_NO_THROW(buf2.write((const uint8_t*)"bar", 3));
114 }
115 
BOOST_AUTO_TEST_CASE(test_default_maximum_buffer_size)116 BOOST_AUTO_TEST_CASE(test_default_maximum_buffer_size)
117 {
118   BOOST_CHECK_EQUAL((std::numeric_limits<uint32_t>::max)(), TMemoryBuffer().getMaxBufferSize());
119 }
120 
BOOST_AUTO_TEST_CASE(test_default_buffer_size)121 BOOST_AUTO_TEST_CASE(test_default_buffer_size)
122 {
123   BOOST_CHECK_EQUAL(1024, TMemoryBuffer().getBufferSize());
124 }
125 
BOOST_AUTO_TEST_CASE(test_error_set_max_buffer_size_too_small)126 BOOST_AUTO_TEST_CASE(test_error_set_max_buffer_size_too_small)
127 {
128   TMemoryBuffer buf;
129   BOOST_CHECK_THROW(buf.setMaxBufferSize(buf.getBufferSize() - 1), TTransportException);
130 }
131 
BOOST_AUTO_TEST_CASE(test_observe)132 BOOST_AUTO_TEST_CASE(test_observe) {
133 #ifdef _MSC_VER
134   #define N 73
135 #else
136   constexpr size_t N = 73;
137 #endif
138   constexpr size_t M = 42;
139   uint8_t one_byte = 42;
140   std::vector<uint8_t> scratch;
141   auto filler = [=]() {
142     std::array<uint8_t, N> x;
143     // Fill buf_mem with a sequence from 0 to N - 1
144     std::iota(x.begin(), x.end(), 0);
145     return x;
146   };
147   static const std::array<uint8_t, N> buf_mem = filler();
148 
149   BOOST_STATIC_ASSERT(M < N);
150 
151   TMemoryBuffer buf((uint8_t*)&buf_mem.front(), N, TMemoryBuffer::MemoryPolicy::OBSERVE);
152 
153   // Readable
154   BOOST_CHECK_EQUAL(N, buf.available_read());
155   // No available write space
156   BOOST_CHECK_EQUAL(0, buf.available_write());
157   // Not writeable
158   BOOST_CHECK_THROW(buf.write(&one_byte, 1), TTransportException);
159 
160   // Read some but not all
161   scratch.resize(M);
162   BOOST_CHECK_EQUAL(M, buf.read(&scratch[0], M));
163   // Check remaining
164   BOOST_CHECK_EQUAL(N - M, buf.available_read());
165   // No available write space
166   BOOST_CHECK_EQUAL(0, buf.available_write());
167   // Not writeable
168   BOOST_CHECK_THROW(buf.write(&one_byte, 1), TTransportException);
169   // Contents
170   BOOST_CHECK_EQUAL_COLLECTIONS(scratch.begin(), scratch.end(), buf_mem.begin(),
171                                 buf_mem.begin() + M);
172 
173   // Readable (drain remaining)
174   scratch.resize(N);
175   BOOST_CHECK_EQUAL(N - M, buf.read(&scratch[M], N - M));
176   // No available write space
177   BOOST_CHECK_EQUAL(0, buf.available_write());
178   // Not writeable
179   BOOST_CHECK_THROW(buf.write(&one_byte, 1), TTransportException);
180   // Contents
181   BOOST_CHECK_EQUAL_COLLECTIONS(scratch.begin(), scratch.end(), buf_mem.begin(), buf_mem.end());
182 
183   // Not readable
184   BOOST_CHECK_EQUAL(0, buf.read(&one_byte, 1));
185   BOOST_CHECK_EQUAL(0, buf.available_read());
186   // No available write space
187   BOOST_CHECK_EQUAL(0, buf.available_write());
188   // Not writeable
189   BOOST_CHECK_THROW(buf.write(&one_byte, 1), TTransportException);
190 
191   /* OBSERVE buffer cannot be reread with the default reset */
192 
193   buf.resetBuffer();
194   // Not Readable
195   BOOST_CHECK_EQUAL(0, buf.available_read());
196   // No available write space
197   BOOST_CHECK_EQUAL(0, buf.available_write());
198   // Not writeable
199   BOOST_CHECK_THROW(buf.write(&one_byte, 1), TTransportException);
200 
201   /* OBSERVE buffers do not auto-resize when written to (implicit) */
202   /* OBSERVE buffers can be appended-to (implicit) */
203 }
204 
BOOST_AUTO_TEST_CASE(test_copy)205 BOOST_AUTO_TEST_CASE(test_copy) {
206 #ifdef _MSC_VER
207   #define N 73
208 #else
209   constexpr size_t N = 73;
210 #endif
211   constexpr size_t M = 42;
212   uint8_t one_byte = 42;
213   std::vector<uint8_t> scratch;
214   auto filler = [&]() {
215     std::array<uint8_t, N> x;
216     // Fill buf_mem with a sequence from 0 to N - 1
217     std::iota(x.begin(), x.end(), 0);
218     return x;
219   };
220   static const std::array<uint8_t, N> buf_mem = filler();
221 
222   BOOST_STATIC_ASSERT(M < N);
223 
224   TMemoryBuffer buf((uint8_t*)&buf_mem.front(), N, TMemoryBuffer::MemoryPolicy::COPY);
225 
226   // Readable
227   BOOST_CHECK_EQUAL(N, buf.available_read());
228   // No available write space
229   BOOST_CHECK_EQUAL(0, buf.available_write());
230 
231   // Read some but not all
232   scratch.resize(M);
233   BOOST_CHECK_EQUAL(M, buf.read(&scratch[0], M));
234   // Check remaining
235   BOOST_CHECK_EQUAL(N - M, buf.available_read());
236   // No available write space
237   BOOST_CHECK_EQUAL(0, buf.available_write());
238   // Contents
239   BOOST_CHECK_EQUAL_COLLECTIONS(scratch.begin(), scratch.end(), buf_mem.begin(),
240                                 buf_mem.begin() + M);
241 
242   // Readable (drain remaining)
243   scratch.resize(N);
244   BOOST_CHECK_EQUAL(N - M, buf.read(&scratch[M], N - M));
245   // No available write space
246   BOOST_CHECK_EQUAL(0, buf.available_write());
247   // Contents
248   BOOST_CHECK_EQUAL_COLLECTIONS(scratch.begin(), scratch.end(), buf_mem.begin(), buf_mem.end());
249 
250   // Not readable
251   BOOST_CHECK_EQUAL(0, buf.read(&one_byte, 1));
252   BOOST_CHECK_EQUAL(0, buf.available_read());
253   // No available write space
254   BOOST_CHECK_EQUAL(0, buf.available_write());
255 
256   /* COPY buffer cannot be reread with the default reset */
257 
258   buf.resetBuffer();
259   // Not readable
260   BOOST_CHECK_EQUAL(0, buf.available_read());
261   // Has available write space
262   BOOST_CHECK_EQUAL(N, buf.available_write());
263 
264   /* COPY buffers auto-resize when written to */
265 
266   // Not readable
267   BOOST_CHECK_EQUAL(0, buf.read(&one_byte, 1));
268   BOOST_CHECK_EQUAL(0, buf.available_read());
269   // No available write space
270   BOOST_CHECK_GT(buf.available_write(), 0);
271   // Writeable
272   one_byte = M;
273   BOOST_CHECK_NO_THROW(buf.write(&one_byte, 1));
274   // Readable
275   one_byte = 0xff;
276   BOOST_CHECK_EQUAL(1, buf.read(&one_byte, 1));
277   BOOST_CHECK_EQUAL(one_byte, M);
278 
279   /* COPY buffers can be appended-to (and auto-resize) */
280 
281   buf.resetBuffer((uint8_t*)&buf_mem.front(), N, TMemoryBuffer::MemoryPolicy::COPY);
282   // Appendable
283   one_byte = N + 1;
284   BOOST_CHECK_NO_THROW(buf.write(&one_byte, 1));
285   BOOST_CHECK_EQUAL(N, buf.read(&scratch[0], N));
286   BOOST_CHECK_EQUAL_COLLECTIONS(scratch.begin(), scratch.end(), buf_mem.begin(),
287                                 buf_mem.begin() + N);
288   one_byte = 0xff;
289   BOOST_CHECK_EQUAL(1, buf.read(&one_byte, 1));
290   BOOST_CHECK_EQUAL(one_byte, N + 1);
291 }
292 
BOOST_AUTO_TEST_CASE(test_take_ownership)293 BOOST_AUTO_TEST_CASE(test_take_ownership)
294 {
295 #ifdef _MSC_VER
296   #define N 73
297 #else
298   constexpr size_t N = 73;
299 #endif
300   constexpr size_t M = 42;
301   uint8_t one_byte = 42;
302   std::vector<uint8_t> scratch;
303   auto filler = [&]() {
304     /* TAKE_OWNERSHIP buffers MUST be malloc'ed */
305     uint8_t* x = static_cast<uint8_t*>(malloc(N));
306     // Fill buf_mem with a sequence from 0 to N - 1
307     std::iota(&x[0], &x[N], 0);
308     return x;
309   };
310   uint8_t* buf_mem = filler();
311 
312   BOOST_STATIC_ASSERT(M < N);
313 
314   TMemoryBuffer buf(buf_mem, N, TMemoryBuffer::MemoryPolicy::TAKE_OWNERSHIP);
315 
316   // Readable
317   BOOST_CHECK_EQUAL(N, buf.available_read());
318   // No available write space
319   BOOST_CHECK_EQUAL(0, buf.available_write());
320 
321   // Read some but not all
322   scratch.resize(M);
323   BOOST_CHECK_EQUAL(M, buf.read(&scratch[0], M));
324   // Check remaining
325   BOOST_CHECK_EQUAL(N - M, buf.available_read());
326   // No available write space
327   BOOST_CHECK_EQUAL(0, buf.available_write());
328   // Contents
329   BOOST_CHECK_EQUAL_COLLECTIONS(scratch.begin(), scratch.end(), &buf_mem[0], &buf_mem[M]);
330 
331   // Readable (drain remaining)
332   scratch.resize(N);
333   BOOST_CHECK_EQUAL(N - M, buf.read(&scratch[M], N - M));
334   // No available write space
335   BOOST_CHECK_EQUAL(0, buf.available_write());
336   // Contents
337   BOOST_CHECK_EQUAL_COLLECTIONS(scratch.begin(), scratch.end(), &buf_mem[0], &buf_mem[N]);
338 
339   // Not readable
340   BOOST_CHECK_EQUAL(0, buf.read(&one_byte, 1));
341   BOOST_CHECK_EQUAL(0, buf.available_read());
342   // No available write space
343   BOOST_CHECK_EQUAL(0, buf.available_write());
344 
345   /* TAKE_OWNERSHIP buffers auto-resize when written to */
346 
347   // Not readable
348   BOOST_CHECK_EQUAL(0, buf.read(&one_byte, 1));
349   BOOST_CHECK_EQUAL(0, buf.available_read());
350   // No available write space
351   BOOST_CHECK_EQUAL(buf.available_write(), 0);
352   // Writeable
353   one_byte = M;
354   BOOST_CHECK_NO_THROW(buf.write(&one_byte, 1));
355   // Readable
356   one_byte = 0xff;
357   BOOST_CHECK_EQUAL(1, buf.read(&one_byte, 1));
358   BOOST_CHECK_EQUAL(one_byte, M);
359 
360   /* TAKE_OWNERSHIP buffers can be appended-to (and auto-resize) */
361 
362   buf_mem = filler();
363   buf.resetBuffer(buf_mem, N, TMemoryBuffer::MemoryPolicy::COPY);
364   // Appendable
365   one_byte = N + 1;
366   BOOST_CHECK_NO_THROW(buf.write(&one_byte, 1));
367   BOOST_CHECK_EQUAL(N, buf.read(&scratch[0], N));
368   BOOST_CHECK_EQUAL_COLLECTIONS(scratch.begin(), scratch.end(), &buf_mem[0], &buf_mem[N]);
369   one_byte = 0xff;
370   BOOST_CHECK_EQUAL(1, buf.read(&one_byte, 1));
371   BOOST_CHECK_EQUAL(one_byte, N + 1);
372 }
373 
BOOST_AUTO_TEST_CASE(test_maximum_buffer_size)374 BOOST_AUTO_TEST_CASE(test_maximum_buffer_size)
375 {
376   TMemoryBuffer buf;
377   buf.setMaxBufferSize(8192);
378   std::vector<uint8_t> small_buff(1);
379 
380   for (size_t i = 0; i < 8192; ++i)
381   {
382     buf.write(&small_buff[0], 1);
383   }
384 
385   BOOST_CHECK_THROW(buf.write(&small_buff[0], 1), TTransportException);
386 }
387 
BOOST_AUTO_TEST_CASE(test_memory_buffer_to_get_sizeof_objects)388 BOOST_AUTO_TEST_CASE(test_memory_buffer_to_get_sizeof_objects)
389 {
390   // This is a demonstration of how to use TMemoryBuffer to determine
391   // the serialized size of a thrift object in the Binary protocol.
392   // See THRIFT-3480
393 
394   shared_ptr<TMemoryBuffer> memBuffer(new TMemoryBuffer());
395   shared_ptr<TBinaryProtocol> binaryProtcol(new TBinaryProtocol(memBuffer));
396 
397   thrift::test::Xtruct object;
398   object.i32_thing = 10;
399   object.i64_thing = 30;
400   object.string_thing = "who's your daddy?";
401 
402   uint32_t size = object.write(binaryProtcol.get());
403   BOOST_CHECK_EQUAL(47, size);
404 }
405 
406 BOOST_AUTO_TEST_SUITE_END()
407