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 #ifndef _GNU_SOURCE
21 #define _GNU_SOURCE // needed for getopt_long
22 #endif
23 
24 #if defined(_MSC_VER) && (_MSC_VER <= 1700)
25 // polynomial and std::fill_t warning happens in MSVC 2010, 2013, maybe others
26 // https://svn.boost.org/trac/boost/ticket/11426
27 #pragma warning(disable:4996)
28 #endif
29 
30 #ifdef HAVE_STDINT_H
31 #include <stdint.h>
32 #endif
33 #ifdef HAVE_INTTYPES_H
34 #include <inttypes.h>
35 #endif
36 #include <cstddef>
37 #include <fstream>
38 #include <iostream>
39 #include <memory>
40 
41 #include <boost/random.hpp>
42 #include <boost/shared_array.hpp>
43 #include <boost/test/unit_test.hpp>
44 #include <boost/version.hpp>
45 
46 #include <thrift/transport/TBufferTransports.h>
47 #include <thrift/transport/TZlibTransport.h>
48 
49 using namespace apache::thrift::transport;
50 using std::shared_ptr;
51 using std::string;
52 
53 boost::mt19937 rng;
54 
55 /*
56  * Utility code
57  */
58 
59 class SizeGenerator {
60 public:
61   virtual ~SizeGenerator() = default;
62   virtual unsigned int getSize() = 0;
63 };
64 
65 class ConstantSizeGenerator : public SizeGenerator {
66 public:
ConstantSizeGenerator(unsigned int value)67   ConstantSizeGenerator(unsigned int value) : value_(value) {}
getSize()68   unsigned int getSize() override { return value_; }
69 
70 private:
71   unsigned int value_;
72 };
73 
74 class LogNormalSizeGenerator : public SizeGenerator {
75 public:
LogNormalSizeGenerator(double mean,double std_dev)76   LogNormalSizeGenerator(double mean, double std_dev)
77     : gen_(rng, boost::lognormal_distribution<double>(mean, std_dev)) {}
78 
getSize()79   unsigned int getSize() override {
80     // Loop until we get a size of 1 or more
81     while (true) {
82       auto value = static_cast<unsigned int>(gen_());
83       if (value >= 1) {
84         return value;
85       }
86     }
87   }
88 
89 private:
90   boost::variate_generator<boost::mt19937, boost::lognormal_distribution<double> > gen_;
91 };
92 
gen_uniform_buffer(uint32_t buf_len,uint8_t c)93 boost::shared_array<uint8_t> gen_uniform_buffer(uint32_t buf_len, uint8_t c) {
94   auto* buf = new uint8_t[buf_len];
95   memset(buf, c, buf_len);
96   return boost::shared_array<uint8_t>(buf);
97 }
98 
gen_compressible_buffer(uint32_t buf_len)99 boost::shared_array<uint8_t> gen_compressible_buffer(uint32_t buf_len) {
100   auto* buf = new uint8_t[buf_len];
101 
102   // Generate small runs of alternately increasing and decreasing bytes
103   boost::uniform_smallint<uint32_t> run_length_distribution(1, 64);
104   boost::uniform_smallint<uint8_t> byte_distribution(0, UINT8_MAX);
105   boost::variate_generator<boost::mt19937, boost::uniform_smallint<uint8_t> >
106       byte_generator(rng, byte_distribution);
107   boost::variate_generator<boost::mt19937, boost::uniform_smallint<uint32_t> >
108       run_len_generator(rng, run_length_distribution);
109 
110   uint32_t idx = 0;
111   int8_t step = 1;
112   while (idx < buf_len) {
113     uint32_t run_length = run_len_generator();
114     if (idx + run_length > buf_len) {
115       run_length = buf_len - idx;
116     }
117 
118     uint8_t byte = byte_generator();
119     for (uint32_t n = 0; n < run_length; ++n) {
120       buf[idx] = byte;
121       ++idx;
122       byte += step;
123     }
124 
125     step *= -1;
126   }
127 
128   return boost::shared_array<uint8_t>(buf);
129 }
130 
gen_random_buffer(uint32_t buf_len)131 boost::shared_array<uint8_t> gen_random_buffer(uint32_t buf_len) {
132   auto* buf = new uint8_t[buf_len];
133 
134   boost::uniform_smallint<uint8_t> distribution(0, UINT8_MAX);
135   boost::variate_generator<boost::mt19937, boost::uniform_smallint<uint8_t> >
136       generator(rng, distribution);
137 
138   for (uint32_t n = 0; n < buf_len; ++n) {
139     buf[n] = generator();
140   }
141 
142   return boost::shared_array<uint8_t>(buf);
143 }
144 
145 /*
146  * Test functions
147  */
148 
test_write_then_read(const boost::shared_array<uint8_t> buf,uint32_t buf_len)149 void test_write_then_read(const boost::shared_array<uint8_t> buf, uint32_t buf_len) {
150   shared_ptr<TMemoryBuffer> membuf(new TMemoryBuffer());
151   shared_ptr<TZlibTransport> zlib_trans(new TZlibTransport(membuf));
152   zlib_trans->write(buf.get(), buf_len);
153   zlib_trans->finish();
154 
155   boost::shared_array<uint8_t> mirror(new uint8_t[buf_len]);
156   uint32_t got = zlib_trans->readAll(mirror.get(), buf_len);
157   BOOST_REQUIRE_EQUAL(got, buf_len);
158   BOOST_CHECK_EQUAL(memcmp(mirror.get(), buf.get(), buf_len), 0);
159   zlib_trans->verifyChecksum();
160 }
161 
test_separate_checksum(const boost::shared_array<uint8_t> buf,uint32_t buf_len)162 void test_separate_checksum(const boost::shared_array<uint8_t> buf, uint32_t buf_len) {
163   // This one is tricky.  I separate the last byte of the stream out
164   // into a separate crbuf_.  The last byte is part of the checksum,
165   // so the entire read goes fine, but when I go to verify the checksum
166   // it isn't there.  The original implementation complained that
167   // the stream was not complete.  I'm about to go fix that.
168   // It worked.  Awesome.
169   shared_ptr<TMemoryBuffer> membuf(new TMemoryBuffer());
170   shared_ptr<TZlibTransport> zlib_trans(new TZlibTransport(membuf));
171   zlib_trans->write(buf.get(), buf_len);
172   zlib_trans->finish();
173   string tmp_buf;
174   membuf->appendBufferToString(tmp_buf);
175   zlib_trans.reset(new TZlibTransport(membuf,
176                                       TZlibTransport::DEFAULT_URBUF_SIZE,
177                                       static_cast<uint32_t>(tmp_buf.length() - 1)));
178 
179   boost::shared_array<uint8_t> mirror(new uint8_t[buf_len]);
180   uint32_t got = zlib_trans->readAll(mirror.get(), buf_len);
181   BOOST_REQUIRE_EQUAL(got, buf_len);
182   BOOST_CHECK_EQUAL(memcmp(mirror.get(), buf.get(), buf_len), 0);
183   zlib_trans->verifyChecksum();
184 }
185 
test_incomplete_checksum(const boost::shared_array<uint8_t> buf,uint32_t buf_len)186 void test_incomplete_checksum(const boost::shared_array<uint8_t> buf, uint32_t buf_len) {
187   // Make sure we still get that "not complete" error if
188   // it really isn't complete.
189   shared_ptr<TMemoryBuffer> membuf(new TMemoryBuffer());
190   shared_ptr<TZlibTransport> zlib_trans(new TZlibTransport(membuf));
191   zlib_trans->write(buf.get(), buf_len);
192   zlib_trans->finish();
193   string tmp_buf;
194   membuf->appendBufferToString(tmp_buf);
195   tmp_buf.erase(tmp_buf.length() - 1);
196   membuf->resetBuffer(const_cast<uint8_t*>(reinterpret_cast<const uint8_t*>(tmp_buf.data())),
197                       static_cast<uint32_t>(tmp_buf.length()));
198 
199   boost::shared_array<uint8_t> mirror(new uint8_t[buf_len]);
200   uint32_t got = zlib_trans->readAll(mirror.get(), buf_len);
201   BOOST_REQUIRE_EQUAL(got, buf_len);
202   BOOST_CHECK_EQUAL(memcmp(mirror.get(), buf.get(), buf_len), 0);
203   try {
204     zlib_trans->verifyChecksum();
205     BOOST_ERROR("verifyChecksum() did not report an error");
206   } catch (TTransportException& ex) {
207     BOOST_CHECK_EQUAL(ex.getType(), TTransportException::CORRUPTED_DATA);
208   }
209 }
210 
test_read_write_mix(const boost::shared_array<uint8_t> buf,uint32_t buf_len,const shared_ptr<SizeGenerator> & write_gen,const shared_ptr<SizeGenerator> & read_gen)211 void test_read_write_mix(const boost::shared_array<uint8_t> buf,
212                          uint32_t buf_len,
213                          const shared_ptr<SizeGenerator>& write_gen,
214                          const shared_ptr<SizeGenerator>& read_gen) {
215   // Try it with a mix of read/write sizes.
216   shared_ptr<TMemoryBuffer> membuf(new TMemoryBuffer());
217   shared_ptr<TZlibTransport> zlib_trans(new TZlibTransport(membuf));
218   unsigned int tot;
219 
220   tot = 0;
221   while (tot < buf_len) {
222     uint32_t write_len = write_gen->getSize();
223     if (tot + write_len > buf_len) {
224       write_len = buf_len - tot;
225     }
226     zlib_trans->write(buf.get() + tot, write_len);
227     tot += write_len;
228   }
229 
230   zlib_trans->finish();
231 
232   tot = 0;
233   boost::shared_array<uint8_t> mirror(new uint8_t[buf_len]);
234   while (tot < buf_len) {
235     uint32_t read_len = read_gen->getSize();
236     uint32_t expected_read_len = read_len;
237     if (tot + read_len > buf_len) {
238       expected_read_len = buf_len - tot;
239     }
240     uint32_t got = zlib_trans->read(mirror.get() + tot, read_len);
241     BOOST_REQUIRE_LE(got, expected_read_len);
242     BOOST_REQUIRE_NE(got, (uint32_t)0);
243     tot += got;
244   }
245 
246   BOOST_CHECK_EQUAL(memcmp(mirror.get(), buf.get(), buf_len), 0);
247   zlib_trans->verifyChecksum();
248 }
249 
test_invalid_checksum(const boost::shared_array<uint8_t> buf,uint32_t buf_len)250 void test_invalid_checksum(const boost::shared_array<uint8_t> buf, uint32_t buf_len) {
251   // Verify checksum checking.
252   shared_ptr<TMemoryBuffer> membuf(new TMemoryBuffer());
253   shared_ptr<TZlibTransport> zlib_trans(new TZlibTransport(membuf));
254   zlib_trans->write(buf.get(), buf_len);
255   zlib_trans->finish();
256   string tmp_buf;
257   membuf->appendBufferToString(tmp_buf);
258   // Modify a byte at the end of the buffer (part of the checksum).
259   // On rare occasions, modifying a byte in the middle of the buffer
260   // isn't caught by the checksum.
261   //
262   // (This happens especially often for the uniform buffer.  The
263   // re-inflated data is correct, however.  I suspect in this case that
264   // we're more likely to modify bytes that are part of zlib metadata
265   // instead of the actual compressed data.)
266   //
267   // I've also seen some failure scenarios where a checksum failure isn't
268   // reported, but zlib keeps trying to decode past the end of the data.
269   // (When this occurs, verifyChecksum() throws an exception indicating
270   // that the end of the data hasn't been reached.)  I haven't seen this
271   // error when only modifying checksum bytes.
272   int index = static_cast<int>(tmp_buf.size() - 1);
273   tmp_buf[index]++;
274   membuf->resetBuffer(const_cast<uint8_t*>(reinterpret_cast<const uint8_t*>(tmp_buf.data())),
275                       static_cast<uint32_t>(tmp_buf.length()));
276 
277   boost::shared_array<uint8_t> mirror(new uint8_t[buf_len]);
278   try {
279     zlib_trans->readAll(mirror.get(), buf_len);
280     zlib_trans->verifyChecksum();
281     BOOST_ERROR("verifyChecksum() did not report an error");
282   } catch (TZlibTransportException& ex) {
283     BOOST_CHECK_EQUAL(ex.getType(), TTransportException::INTERNAL_ERROR);
284   }
285 }
286 
test_write_after_flush(const boost::shared_array<uint8_t> buf,uint32_t buf_len)287 void test_write_after_flush(const boost::shared_array<uint8_t> buf, uint32_t buf_len) {
288   // write some data
289   shared_ptr<TMemoryBuffer> membuf(new TMemoryBuffer());
290   shared_ptr<TZlibTransport> zlib_trans(new TZlibTransport(membuf));
291   zlib_trans->write(buf.get(), buf_len);
292 
293   // call finish()
294   zlib_trans->finish();
295 
296   // make sure write() throws an error
297   try {
298     uint8_t write_buf[] = "a";
299     zlib_trans->write(write_buf, 1);
300     BOOST_ERROR("write() after finish() did not raise an exception");
301   } catch (TTransportException& ex) {
302     BOOST_CHECK_EQUAL(ex.getType(), TTransportException::BAD_ARGS);
303   }
304 
305   // make sure flush() throws an error
306   try {
307     zlib_trans->flush();
308     BOOST_ERROR("flush() after finish() did not raise an exception");
309   } catch (TTransportException& ex) {
310     BOOST_CHECK_EQUAL(ex.getType(), TTransportException::BAD_ARGS);
311   }
312 
313   // make sure finish() throws an error
314   try {
315     zlib_trans->finish();
316     BOOST_ERROR("finish() after finish() did not raise an exception");
317   } catch (TTransportException& ex) {
318     BOOST_CHECK_EQUAL(ex.getType(), TTransportException::BAD_ARGS);
319   }
320 }
321 
test_no_write()322 void test_no_write() {
323   // Verify that no data is written to the underlying transport if we
324   // never write data to the TZlibTransport.
325   shared_ptr<TMemoryBuffer> membuf(new TMemoryBuffer());
326   {
327     // Create a TZlibTransport object, and immediately destroy it
328     // when it goes out of scope.
329     TZlibTransport w_zlib_trans(membuf);
330   }
331 
332   BOOST_CHECK_EQUAL(membuf->available_read(), (uint32_t)0);
333 }
334 
test_get_underlying_transport()335 void test_get_underlying_transport() {
336   shared_ptr<TMemoryBuffer> membuf(new TMemoryBuffer());
337   shared_ptr<TZlibTransport> zlib_trans(new TZlibTransport(membuf));
338   BOOST_CHECK_EQUAL(membuf.get(), zlib_trans->getUnderlyingTransport().get());
339 }
340 
341 /*
342  * Initialization
343  */
344 
345 #if (BOOST_VERSION >= 105900)
346 #define ADD_TEST_CASE(suite, name, _FUNC, ...)                                                     \
347   do {                                                                                             \
348     ::std::ostringstream name_ss;                                                                  \
349     name_ss << name << "-" << BOOST_STRINGIZE(_FUNC);                                              \
350     ::std::function<void ()> test_func =                                        \
351         ::std::bind(_FUNC, ##__VA_ARGS__);                                      \
352     ::boost::unit_test::test_case* tc                                                              \
353         = ::boost::unit_test::make_test_case(test_func, name_ss.str(), __FILE__, __LINE__);        \
354     (suite)->add(tc);                                                                              \
355   } while (0)
356 #else
357 #define ADD_TEST_CASE(suite, name, _FUNC, ...)                                                     \
358   do {                                                                                             \
359     ::std::ostringstream name_ss;                                                                  \
360     name_ss << name << "-" << BOOST_STRINGIZE(_FUNC);                                              \
361     ::boost::unit_test::test_case* tc                                                              \
362         = ::boost::unit_test::make_test_case(::std::bind(_FUNC,                 \
363                                                                             ##__VA_ARGS__),        \
364                                              name_ss.str());                                       \
365     (suite)->add(tc);                                                                              \
366   } while (0)
367 #endif
368 
add_tests(boost::unit_test::test_suite * suite,const boost::shared_array<uint8_t> & buf,uint32_t buf_len,const char * name)369 void add_tests(boost::unit_test::test_suite* suite,
370                const boost::shared_array<uint8_t>& buf,
371                uint32_t buf_len,
372                const char* name) {
373   ADD_TEST_CASE(suite, name, test_write_then_read, buf, buf_len);
374   ADD_TEST_CASE(suite, name, test_separate_checksum, buf, buf_len);
375   ADD_TEST_CASE(suite, name, test_incomplete_checksum, buf, buf_len);
376   ADD_TEST_CASE(suite, name, test_invalid_checksum, buf, buf_len);
377   ADD_TEST_CASE(suite, name, test_write_after_flush, buf, buf_len);
378 
379   shared_ptr<SizeGenerator> size_32k(new ConstantSizeGenerator(1 << 15));
380   shared_ptr<SizeGenerator> size_lognormal(new LogNormalSizeGenerator(20, 30));
381   ADD_TEST_CASE(suite, name << "-constant", test_read_write_mix, buf, buf_len, size_32k, size_32k);
382   ADD_TEST_CASE(suite,
383                 name << "-lognormal-write",
384                 test_read_write_mix,
385                 buf,
386                 buf_len,
387                 size_lognormal,
388                 size_32k);
389   ADD_TEST_CASE(suite,
390                 name << "-lognormal-read",
391                 test_read_write_mix,
392                 buf,
393                 buf_len,
394                 size_32k,
395                 size_lognormal);
396   ADD_TEST_CASE(suite,
397                 name << "-lognormal-both",
398                 test_read_write_mix,
399                 buf,
400                 buf_len,
401                 size_lognormal,
402                 size_lognormal);
403 
404   // Test with a random size distribution,
405   // but use the exact same distribution for reading as for writing.
406   //
407   // Because the SizeGenerator makes a copy of the random number generator,
408   // both SizeGenerators should return the exact same set of values, since they
409   // both start with random number generators in the same state.
410   shared_ptr<SizeGenerator> write_size_gen(new LogNormalSizeGenerator(20, 30));
411   shared_ptr<SizeGenerator> read_size_gen(new LogNormalSizeGenerator(20, 30));
412   ADD_TEST_CASE(suite,
413                 name << "-lognormal-same-distribution",
414                 test_read_write_mix,
415                 buf,
416                 buf_len,
417                 write_size_gen,
418                 read_size_gen);
419 }
420 
print_usage(FILE * f,const char * argv0)421 void print_usage(FILE* f, const char* argv0) {
422   fprintf(f, "Usage: %s [boost_options] [options]\n", argv0);
423   fprintf(f, "Options:\n");
424   fprintf(f, "  --seed=<N>, -s <N>\n");
425   fprintf(f, "  --help\n");
426 }
427 
428 #ifdef BOOST_TEST_DYN_LINK
init_unit_test_suite()429 bool init_unit_test_suite() {
430   auto seed = static_cast<uint32_t>(time(nullptr));
431 #ifdef HAVE_INTTYPES_H
432   printf("seed: %" PRIu32 "\n", seed);
433 #endif
434   rng.seed(seed);
435 
436   boost::unit_test::test_suite* suite = &boost::unit_test::framework::master_test_suite();
437   suite->p_name.value = "ZlibTest";
438 
439   uint32_t buf_len = 1024 * 32;
440   add_tests(suite, gen_uniform_buffer(buf_len, 'a'), buf_len, "uniform");
441   add_tests(suite, gen_compressible_buffer(buf_len), buf_len, "compressible");
442   add_tests(suite, gen_random_buffer(buf_len), buf_len, "random");
443 
444   suite->add(BOOST_TEST_CASE(test_no_write));
445   suite->add(BOOST_TEST_CASE(test_get_underlying_transport));
446 
447   return true;
448 }
449 
main(int argc,char * argv[])450 int main( int argc, char* argv[] ) {
451   return ::boost::unit_test::unit_test_main(&init_unit_test_suite,argc,argv);
452 }
453 #else
init_unit_test_suite(int argc,char * argv[])454 boost::unit_test::test_suite* init_unit_test_suite(int argc, char* argv[]) {
455   THRIFT_UNUSED_VARIABLE(argc);
456   THRIFT_UNUSED_VARIABLE(argv);
457   uint32_t seed = static_cast<uint32_t>(time(nullptr));
458 #ifdef HAVE_INTTYPES_H
459   printf("seed: %" PRIu32 "\n", seed);
460 #endif
461   rng.seed(seed);
462 
463   boost::unit_test::test_suite* suite = &boost::unit_test::framework::master_test_suite();
464   suite->p_name.value = "ZlibTest";
465 
466   uint32_t buf_len = 1024 * 32;
467   add_tests(suite, gen_uniform_buffer(buf_len, 'a'), buf_len, "uniform");
468   add_tests(suite, gen_compressible_buffer(buf_len), buf_len, "compressible");
469   add_tests(suite, gen_random_buffer(buf_len), buf_len, "random");
470 
471   suite->add(BOOST_TEST_CASE(test_no_write));
472 
473   return nullptr;
474 }
475 #endif
476