folly/io/test/IOBufCursorTest.cpp

   1 /*
   2  * Copyright 2013-present Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <folly/io/IOBuf.h>
  18
  19 #include <folly/Format.h>
  20 #include <folly/Range.h>
  21 #include <folly/io/Cursor.h>
  22 #include <folly/portability/GTest.h>
  23 #include <numeric>
  24 #include <vector>
  25
  26 using folly::ByteRange;
  27 using folly::format;
  28 using folly::IOBuf;
  29 using folly::StringPiece;
  30 using std::unique_ptr;
  31 using namespace folly::io;
  32
  33 TEST(IOBuf, RWCursor) {
  34   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
  35   iobuf1->append(20);
  36   unique_ptr<IOBuf> iobuf2(IOBuf::create(20));
  37   iobuf2->append(20);
  38
  39   iobuf2.get();
  40   iobuf1->prependChain(std::move(iobuf2));
  41
  42   EXPECT_TRUE(iobuf1->isChained());
  43
  44   RWPrivateCursor wcursor(iobuf1.get());
  45   Cursor rcursor(iobuf1.get());
  46   wcursor.writeLE((uint64_t)1);
  47   wcursor.writeLE((uint64_t)1);
  48   wcursor.writeLE((uint64_t)1);
  49   wcursor.write((uint8_t)1);
  50
  51   EXPECT_EQ(1u, rcursor.readLE<uint64_t>());
  52   rcursor.skip(8);
  53   EXPECT_EQ(1u, rcursor.readLE<uint32_t>());
  54   rcursor.skip(0);
  55   EXPECT_EQ(0u, rcursor.read<uint8_t>());
  56   EXPECT_EQ(0u, rcursor.read<uint8_t>());
  57   EXPECT_EQ(0u, rcursor.read<uint8_t>());
  58   EXPECT_EQ(0u, rcursor.read<uint8_t>());
  59   EXPECT_EQ(1u, rcursor.read<uint8_t>());
  60 }
  61
  62 TEST(IOBuf, skip) {
  63   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
  64   iobuf1->append(20);
  65   RWPrivateCursor wcursor(iobuf1.get());
  66   wcursor.write((uint8_t)1);
  67   wcursor.write((uint8_t)2);
  68   Cursor cursor(iobuf1.get());
  69   cursor.skip(1);
  70   EXPECT_EQ(2, cursor.read<uint8_t>());
  71 }
  72
  73 TEST(IOBuf, reset) {
  74   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
  75   iobuf1->append(20);
  76   RWPrivateCursor wcursor(iobuf1.get());
  77   wcursor.write((uint8_t)1);
  78   wcursor.write((uint8_t)2);
  79   wcursor.reset(iobuf1.get());
  80   EXPECT_EQ(1, wcursor.read<uint8_t>());
  81 }
  82
  83 TEST(IOBuf, copy_assign_convert) {
  84   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
  85   iobuf1->append(20);
  86   RWPrivateCursor wcursor(iobuf1.get());
  87   RWPrivateCursor cursor2(wcursor);
  88   RWPrivateCursor cursor3(iobuf1.get());
  89
  90   wcursor.write((uint8_t)1);
  91   cursor3 = wcursor;
  92   wcursor.write((uint8_t)2);
  93   Cursor cursor4(wcursor);
  94   RWPrivateCursor cursor5(wcursor);
  95   wcursor.write((uint8_t)3);
  96
  97   EXPECT_EQ(1, cursor2.read<uint8_t>());
  98   EXPECT_EQ(2, cursor3.read<uint8_t>());
  99   EXPECT_EQ(3, cursor4.read<uint8_t>());
 100 }
 101
 102 TEST(IOBuf, arithmetic) {
 103   IOBuf iobuf1(IOBuf::CREATE, 20);
 104   iobuf1.append(20);
 105   RWPrivateCursor wcursor(&iobuf1);
 106   wcursor += 1;
 107   wcursor.write((uint8_t)1);
 108   Cursor cursor(&iobuf1);
 109   cursor += 1;
 110   EXPECT_EQ(1, cursor.read<uint8_t>());
 111
 112   Cursor start(&iobuf1);
 113   Cursor cursor2 = start + 9;
 114   EXPECT_EQ(7, cursor2 - cursor);
 115   EXPECT_NE(cursor, cursor2);
 116   cursor += 8;
 117   cursor2 = cursor2 + 1;
 118   EXPECT_EQ(cursor, cursor2);
 119 }
 120
 121 TEST(IOBuf, endian) {
 122   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
 123   iobuf1->append(20);
 124   RWPrivateCursor wcursor(iobuf1.get());
 125   Cursor rcursor(iobuf1.get());
 126   uint16_t v = 1;
 127   int16_t vu = -1;
 128   wcursor.writeBE(v);
 129   wcursor.writeBE(vu);
 130   // Try a couple combinations to ensure they were generated correctly
 131   wcursor.writeBE(vu);
 132   wcursor.writeLE(vu);
 133   wcursor.writeLE(vu);
 134   wcursor.writeLE(v);
 135   EXPECT_EQ(v, rcursor.readBE<uint16_t>());
 136 }
 137
 138 TEST(IOBuf, Cursor) {
 139   unique_ptr<IOBuf> iobuf1(IOBuf::create(1));
 140   iobuf1->append(1);
 141   RWPrivateCursor c(iobuf1.get());
 142   c.write((uint8_t)40); // OK
 143   try {
 144     c.write((uint8_t)10); // Bad write, checked should except.
 145     ADD_FAILURE();
 146   } catch (...) {
 147   }
 148 }
 149
 150 TEST(IOBuf, UnshareCursor) {
 151   uint8_t buf = 0;
 152   unique_ptr<IOBuf> iobuf1(IOBuf::wrapBuffer(&buf, 1));
 153   unique_ptr<IOBuf> iobuf2(IOBuf::wrapBuffer(&buf, 1));
 154   RWUnshareCursor c1(iobuf1.get());
 155   RWUnshareCursor c2(iobuf2.get());
 156
 157   c1.write((uint8_t)10); // This should duplicate the two buffers.
 158   uint8_t t = c2.read<uint8_t>();
 159   EXPECT_EQ(0, t);
 160
 161   iobuf1 = IOBuf::wrapBuffer(&buf, 1);
 162   iobuf2 = IOBuf::wrapBuffer(&buf, 1);
 163   RWPrivateCursor c3(iobuf1.get());
 164   RWPrivateCursor c4(iobuf2.get());
 165
 166   c3.write((uint8_t)10); // This should _not_ duplicate the two buffers.
 167   t = c4.read<uint8_t>();
 168   EXPECT_EQ(10, t);
 169 }
 170
 171 namespace {
 172 void append(std::unique_ptr<IOBuf>& buf, folly::StringPiece data) {
 173   EXPECT_LE(data.size(), buf->tailroom());
 174   memcpy(buf->writableData(), data.data(), data.size());
 175   buf->append(data.size());
 176 }
 177
 178 void append(Appender& appender, StringPiece data) {
 179   appender.push(ByteRange(data));
 180 }
 181
 182 std::string toString(const IOBuf& buf) {
 183   std::string str;
 184   Cursor cursor(&buf);
 185   ByteRange b;
 186   while (!(b = cursor.peekBytes()).empty()) {
 187     str.append(reinterpret_cast<const char*>(b.data()), b.size());
 188     cursor.skip(b.size());
 189   }
 190   return str;
 191 }
 192
 193 } // namespace
 194
 195 TEST(IOBuf, PullAndPeek) {
 196   std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
 197   append(iobuf1, "he");
 198   std::unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
 199   append(iobuf2, "llo ");
 200   std::unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
 201   append(iobuf3, "world");
 202   iobuf1->prependChain(std::move(iobuf2));
 203   iobuf1->prependChain(std::move(iobuf3));
 204   EXPECT_EQ(3, iobuf1->countChainElements());
 205   EXPECT_EQ(11, iobuf1->computeChainDataLength());
 206
 207   char buf[12];
 208   memset(buf, 0, sizeof(buf));
 209   Cursor(iobuf1.get()).pull(buf, 11);
 210   EXPECT_EQ("hello world", std::string(buf));
 211
 212   memset(buf, 0, sizeof(buf));
 213   EXPECT_EQ(11, Cursor(iobuf1.get()).pullAtMost(buf, 20));
 214   EXPECT_EQ("hello world", std::string(buf));
 215
 216   EXPECT_THROW({Cursor(iobuf1.get()).pull(buf, 20);},
 217                std::out_of_range);
 218
 219   {
 220     RWPrivateCursor cursor(iobuf1.get());
 221     auto b = cursor.peekBytes();
 222     EXPECT_EQ("he", StringPiece(b));
 223     cursor.skip(b.size());
 224     b = cursor.peekBytes();
 225     EXPECT_EQ("llo ", StringPiece(b));
 226     cursor.skip(b.size());
 227     b = cursor.peekBytes();
 228     EXPECT_EQ("world", StringPiece(b));
 229     cursor.skip(b.size());
 230     EXPECT_EQ(3, iobuf1->countChainElements());
 231     EXPECT_EQ(11, iobuf1->computeChainDataLength());
 232   }
 233
 234   {
 235     RWPrivateCursor cursor(iobuf1.get());
 236     cursor.gather(11);
 237     auto b = cursor.peekBytes();
 238     EXPECT_EQ("hello world", StringPiece(b));
 239     EXPECT_EQ(1, iobuf1->countChainElements());
 240     EXPECT_EQ(11, iobuf1->computeChainDataLength());
 241   }
 242 }
 243
 244 TEST(IOBuf, pushCursorData) {
 245   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
 246   iobuf1->append(15);
 247   iobuf1->trimStart(5);
 248   unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
 249   unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
 250   iobuf3->append(10);
 251
 252   iobuf1->prependChain(std::move(iobuf2));
 253   iobuf1->prependChain(std::move(iobuf3));
 254   EXPECT_TRUE(iobuf1->isChained());
 255
 256   //write 20 bytes to the buffer chain
 257   RWPrivateCursor wcursor(iobuf1.get());
 258   EXPECT_FALSE(wcursor.isAtEnd());
 259   wcursor.writeBE<uint64_t>(1);
 260   wcursor.writeBE<uint64_t>(10);
 261   wcursor.writeBE<uint32_t>(20);
 262   EXPECT_TRUE(wcursor.isAtEnd());
 263
 264   // create a read buffer for the buffer chain
 265   Cursor rcursor(iobuf1.get());
 266   EXPECT_EQ(1, rcursor.readBE<uint64_t>());
 267   EXPECT_EQ(10, rcursor.readBE<uint64_t>());
 268   EXPECT_EQ(20, rcursor.readBE<uint32_t>());
 269   EXPECT_EQ(0, rcursor.totalLength());
 270   rcursor.reset(iobuf1.get());
 271   EXPECT_EQ(20, rcursor.totalLength());
 272
 273   // create another write buffer
 274   unique_ptr<IOBuf> iobuf4(IOBuf::create(30));
 275   iobuf4->append(30);
 276   RWPrivateCursor wcursor2(iobuf4.get());
 277   // write buffer chain data into it, now wcursor2 should only
 278   // have 10 bytes writable space
 279   wcursor2.push(rcursor, 20);
 280   EXPECT_EQ(wcursor2.totalLength(), 10);
 281   // write again with not enough space in rcursor
 282   EXPECT_THROW(wcursor2.push(rcursor, 20), std::out_of_range);
 283
 284   // create a read cursor to check iobuf3 data back
 285   Cursor rcursor2(iobuf4.get());
 286   EXPECT_EQ(1, rcursor2.readBE<uint64_t>());
 287   EXPECT_EQ(10, rcursor2.readBE<uint64_t>());
 288   EXPECT_EQ(20, rcursor2.readBE<uint32_t>());
 289 }
 290
 291 TEST(IOBuf, Gather) {
 292   std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
 293   append(iobuf1, "he");
 294   std::unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
 295   append(iobuf2, "llo ");
 296   std::unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
 297   append(iobuf3, "world");
 298   iobuf1->prependChain(std::move(iobuf2));
 299   iobuf1->prependChain(std::move(iobuf3));
 300   EXPECT_EQ(3, iobuf1->countChainElements());
 301   EXPECT_EQ(11, iobuf1->computeChainDataLength());
 302
 303   // Attempting to gather() more data than available in the chain should fail.
 304   // Try from the very beginning of the chain.
 305   RWPrivateCursor cursor(iobuf1.get());
 306   EXPECT_THROW(cursor.gather(15), std::overflow_error);
 307   // Now try from the middle of the chain
 308   cursor += 3;
 309   EXPECT_THROW(cursor.gather(10), std::overflow_error);
 310
 311   // Calling gatherAtMost() should succeed, however, and just gather
 312   // as much as it can
 313   cursor.gatherAtMost(10);
 314   EXPECT_EQ(8, cursor.length());
 315   EXPECT_EQ(8, cursor.totalLength());
 316   EXPECT_FALSE(cursor.isAtEnd());
 317   EXPECT_EQ("lo world",
 318             folly::StringPiece(reinterpret_cast<const char*>(cursor.data()),
 319                                cursor.length()));
 320   EXPECT_EQ(2, iobuf1->countChainElements());
 321   EXPECT_EQ(11, iobuf1->computeChainDataLength());
 322
 323   // Now try gather again on the chain head
 324   cursor = RWPrivateCursor(iobuf1.get());
 325   cursor.gather(5);
 326   // Since gather() doesn't split buffers, everything should be collapsed into
 327   // a single buffer now.
 328   EXPECT_EQ(1, iobuf1->countChainElements());
 329   EXPECT_EQ(11, iobuf1->computeChainDataLength());
 330   EXPECT_EQ(11, cursor.length());
 331   EXPECT_EQ(11, cursor.totalLength());
 332 }
 333
 334 TEST(IOBuf, cloneAndInsert) {
 335   std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
 336   append(iobuf1, "he");
 337   std::unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
 338   append(iobuf2, "llo ");
 339   std::unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
 340   append(iobuf3, "world");
 341   iobuf1->prependChain(std::move(iobuf2));
 342   iobuf1->prependChain(std::move(iobuf3));
 343   EXPECT_EQ(3, iobuf1->countChainElements());
 344   EXPECT_EQ(11, iobuf1->computeChainDataLength());
 345
 346   std::unique_ptr<IOBuf> cloned;
 347
 348   Cursor(iobuf1.get()).clone(cloned, 3);
 349   EXPECT_EQ(2, cloned->countChainElements());
 350   EXPECT_EQ(3, cloned->computeChainDataLength());
 351
 352
 353   EXPECT_EQ(11, Cursor(iobuf1.get()).cloneAtMost(cloned, 20));
 354   EXPECT_EQ(3, cloned->countChainElements());
 355   EXPECT_EQ(11, cloned->computeChainDataLength());
 356
 357
 358   EXPECT_THROW({Cursor(iobuf1.get()).clone(cloned, 20);},
 359                std::out_of_range);
 360
 361   {
 362     // Check that inserting in the middle of an iobuf splits
 363     RWPrivateCursor cursor(iobuf1.get());
 364     Cursor(iobuf1.get()).clone(cloned, 3);
 365     EXPECT_EQ(2, cloned->countChainElements());
 366     EXPECT_EQ(3, cloned->computeChainDataLength());
 367
 368     cursor.skip(1);
 369
 370     cursor.insert(std::move(cloned));
 371     cursor.insert(folly::IOBuf::create(0));
 372     EXPECT_EQ(7, iobuf1->countChainElements());
 373     EXPECT_EQ(14, iobuf1->computeChainDataLength());
 374     // Check that nextBuf got set correctly to the buffer with 1 byte left
 375     EXPECT_EQ(1, cursor.peekBytes().size());
 376     cursor.read<uint8_t>();
 377   }
 378
 379   {
 380     // Check that inserting at the end doesn't create empty buf
 381     RWPrivateCursor cursor(iobuf1.get());
 382     Cursor(iobuf1.get()).clone(cloned, 1);
 383     EXPECT_EQ(1, cloned->countChainElements());
 384     EXPECT_EQ(1, cloned->computeChainDataLength());
 385
 386     cursor.skip(1);
 387
 388     cursor.insert(std::move(cloned));
 389     EXPECT_EQ(8, iobuf1->countChainElements());
 390     EXPECT_EQ(15, iobuf1->computeChainDataLength());
 391     // Check that nextBuf got set correctly
 392     cursor.read<uint8_t>();
 393   }
 394   {
 395     // Check that inserting at the beginning doesn't create empty buf
 396     RWPrivateCursor cursor(iobuf1.get());
 397     Cursor(iobuf1.get()).clone(cloned, 1);
 398     EXPECT_EQ(1, cloned->countChainElements());
 399     EXPECT_EQ(1, cloned->computeChainDataLength());
 400
 401     cursor.insert(std::move(cloned));
 402     EXPECT_EQ(9, iobuf1->countChainElements());
 403     EXPECT_EQ(16, iobuf1->computeChainDataLength());
 404     // Check that nextBuf got set correctly
 405     cursor.read<uint8_t>();
 406   }
 407 }
 408
 409 TEST(IOBuf, cloneWithEmptyBufAtStart) {
 410   folly::IOBufEqual eq;
 411   auto empty = IOBuf::create(0);
 412   auto hel = IOBuf::create(3);
 413   append(hel, "hel");
 414   auto lo = IOBuf::create(2);
 415   append(lo, "lo");
 416
 417   auto iobuf = empty->clone();
 418   iobuf->prependChain(hel->clone());
 419   iobuf->prependChain(lo->clone());
 420   iobuf->prependChain(empty->clone());
 421   iobuf->prependChain(hel->clone());
 422   iobuf->prependChain(lo->clone());
 423   iobuf->prependChain(empty->clone());
 424   iobuf->prependChain(lo->clone());
 425   iobuf->prependChain(hel->clone());
 426   iobuf->prependChain(lo->clone());
 427   iobuf->prependChain(lo->clone());
 428
 429   Cursor cursor(iobuf.get());
 430   std::unique_ptr<IOBuf> cloned;
 431   char data[3];
 432   cursor.pull(&data, 3);
 433   cursor.clone(cloned, 2);
 434   EXPECT_EQ(1, cloned->countChainElements());
 435   EXPECT_EQ(2, cloned->length());
 436   EXPECT_TRUE(eq(lo, cloned));
 437
 438   cursor.pull(&data, 3);
 439   EXPECT_EQ("hel", std::string(data, sizeof(data)));
 440
 441   cursor.skip(2);
 442   cursor.clone(cloned, 2);
 443   EXPECT_TRUE(eq(lo, cloned));
 444
 445   std::string hello = cursor.readFixedString(5);
 446   cursor.clone(cloned, 2);
 447   EXPECT_TRUE(eq(lo, cloned));
 448 }
 449
 450 TEST(IOBuf, Appender) {
 451   std::unique_ptr<IOBuf> head(IOBuf::create(10));
 452   append(head, "hello");
 453
 454   Appender app(head.get(), 10);
 455   auto cap = head->capacity();
 456   auto len1 = app.length();
 457   EXPECT_EQ(cap - 5, len1);
 458   app.ensure(len1);  // won't grow
 459   EXPECT_EQ(len1, app.length());
 460   app.ensure(len1 + 1);  // will grow
 461   EXPECT_LE(len1 + 1, app.length());
 462
 463   append(app, " world");
 464   EXPECT_EQ("hello world", toString(*head));
 465 }
 466
 467 TEST(IOBuf, Printf) {
 468   IOBuf head(IOBuf::CREATE, 24);
 469   Appender app(&head, 32);
 470
 471   app.printf("%s", "test");
 472   EXPECT_EQ(head.length(), 4);
 473   EXPECT_EQ(0, memcmp(head.data(), "test\0", 5));
 474
 475   app.printf("%d%s %s%s %#x", 32, "this string is",
 476              "longer than our original allocation size,",
 477              "and will therefore require a new allocation", 0x12345678);
 478   // The tailroom should start with a nul byte now.
 479   EXPECT_GE(head.prev()->tailroom(), 1u);
 480   EXPECT_EQ(0, *head.prev()->tail());
 481
 482   EXPECT_EQ("test32this string is longer than our original "
 483             "allocation size,and will therefore require a "
 484             "new allocation 0x12345678",
 485             head.moveToFbString().toStdString());
 486 }
 487
 488 TEST(IOBuf, Format) {
 489   IOBuf head(IOBuf::CREATE, 24);
 490   Appender app(&head, 32);
 491
 492   format("{}", "test")(app);
 493   EXPECT_EQ(head.length(), 4);
 494   EXPECT_EQ(0, memcmp(head.data(), "test", 4));
 495
 496   auto fmt = format("{}{} {}{} {:#x}",
 497                     32, "this string is",
 498                     "longer than our original allocation size,",
 499                     "and will therefore require a new allocation",
 500                     0x12345678);
 501   fmt(app);
 502   EXPECT_EQ("test32this string is longer than our original "
 503             "allocation size,and will therefore require a "
 504             "new allocation 0x12345678",
 505             head.moveToFbString().toStdString());
 506 }
 507
 508 TEST(IOBuf, QueueAppender) {
 509   folly::IOBufQueue queue;
 510
 511   // Allocate 100 bytes at once, but don't grow past 1024
 512   QueueAppender app(&queue, 100);
 513   size_t n = 1024 / sizeof(uint32_t);
 514   for (uint32_t i = 0; i < n; ++i) {
 515     app.writeBE(i);
 516   }
 517
 518   // There must be a goodMallocSize between 100 and 1024...
 519   EXPECT_LT(1u, queue.front()->countChainElements());
 520   const IOBuf* buf = queue.front();
 521   do {
 522     EXPECT_LE(100u, buf->capacity());
 523     buf = buf->next();
 524   } while (buf != queue.front());
 525
 526   Cursor cursor(queue.front());
 527   for (uint32_t i = 0; i < n; ++i) {
 528     EXPECT_EQ(i, cursor.readBE<uint32_t>());
 529   }
 530
 531   EXPECT_THROW({cursor.readBE<uint32_t>();}, std::out_of_range);
 532 }
 533
 534 TEST(IOBuf, QueueAppenderPushAtMostFillBuffer) {
 535   folly::IOBufQueue queue;
 536   // There should be a goodMallocSize between 125 and 1000
 537   QueueAppender appender{&queue, 125};
 538   std::vector<uint8_t> data;
 539   data.resize(1000);
 540   std::iota(data.begin(), data.end(), uint8_t(0));
 541   // Add 100 byte
 542   appender.pushAtMost(data.data(), 100);
 543   // Add 900 bytes
 544   appender.pushAtMost(data.data() + 100, data.size() - 100);
 545   const auto buf = queue.front();
 546   // Should fill the current buffer before adding another
 547   EXPECT_LE(2, buf->countChainElements());
 548   EXPECT_EQ(0, buf->tailroom());
 549   EXPECT_LE(125, buf->length());
 550   EXPECT_EQ(1000, buf->computeChainDataLength());
 551   const StringPiece sp{(const char*)data.data(), data.size()};
 552   EXPECT_EQ(sp, toString(*buf));
 553 }
 554
 555 TEST(IOBuf, QueueAppenderInsertOwn) {
 556   auto buf = IOBuf::create(10);
 557   folly::IOBufQueue queue;
 558   QueueAppender appender{&queue, 128};
 559   appender.insert(std::move(buf));
 560
 561   std::vector<uint8_t> data;
 562   data.resize(256);
 563   std::iota(data.begin(), data.end(), 0);
 564   appender.pushAtMost(folly::range(data));
 565   // Buffer is owned, so we should write to it
 566   EXPECT_LE(2, queue.front()->countChainElements());
 567   EXPECT_EQ(0, queue.front()->tailroom());
 568   const StringPiece sp{(const char*)data.data(), data.size()};
 569   EXPECT_EQ(sp, toString(*queue.front()));
 570 }
 571
 572 TEST(IOBuf, QueueAppenderInsertClone) {
 573   IOBuf buf{IOBuf::CREATE, 100};
 574   folly::IOBufQueue queue;
 575   QueueAppender appender{&queue, 100};
 576   // Buffer is shared, so we create a new buffer to write to
 577   appender.insert(buf);
 578   uint8_t x = 42;
 579   appender.pushAtMost(&x, 1);
 580   EXPECT_EQ(2, queue.front()->countChainElements());
 581   EXPECT_EQ(0, queue.front()->length());
 582   EXPECT_LT(0, queue.front()->tailroom());
 583   EXPECT_EQ(1, queue.front()->next()->length());
 584   EXPECT_EQ(x, queue.front()->next()->data()[0]);
 585 }
 586
 587 TEST(IOBuf, CursorOperators) {
 588   // Test operators on a single-item chain
 589   {
 590     std::unique_ptr<IOBuf> chain1(IOBuf::create(20));
 591     chain1->append(10);
 592
 593     Cursor curs1(chain1.get());
 594     EXPECT_EQ(0, curs1 - chain1.get());
 595     EXPECT_FALSE(curs1.isAtEnd());
 596     curs1.skip(3);
 597     EXPECT_EQ(3, curs1 - chain1.get());
 598     EXPECT_FALSE(curs1.isAtEnd());
 599     curs1.skip(7);
 600     EXPECT_EQ(10, curs1 - chain1.get());
 601     EXPECT_TRUE(curs1.isAtEnd());
 602
 603     Cursor curs2(chain1.get());
 604     EXPECT_EQ(0, curs2 - chain1.get());
 605     EXPECT_EQ(10, curs1 - curs2);
 606     EXPECT_THROW(curs2 - curs1, std::out_of_range);
 607   }
 608
 609   // Test cross-chain operations
 610   {
 611     std::unique_ptr<IOBuf> chain1(IOBuf::create(20));
 612     chain1->append(10);
 613     std::unique_ptr<IOBuf> chain2 = chain1->clone();
 614
 615     Cursor curs1(chain1.get());
 616     Cursor curs2(chain2.get());
 617     EXPECT_THROW(curs1 - curs2, std::out_of_range);
 618     EXPECT_THROW(curs1 - chain2.get(), std::out_of_range);
 619   }
 620
 621   // Test operations on multi-item chains
 622   {
 623     std::unique_ptr<IOBuf> chain(IOBuf::create(20));
 624     chain->append(10);
 625     chain->appendChain(chain->clone());
 626     EXPECT_EQ(20, chain->computeChainDataLength());
 627
 628     Cursor curs1(chain.get());
 629     curs1.skip(5);
 630     Cursor curs2(chain.get());
 631     curs2.skip(3);
 632     EXPECT_EQ(2, curs1 - curs2);
 633     EXPECT_EQ(5, curs1 - chain.get());
 634     EXPECT_THROW(curs2 - curs1, std::out_of_range);
 635
 636     curs1.skip(7);
 637     EXPECT_EQ(9, curs1 - curs2);
 638     EXPECT_EQ(12, curs1 - chain.get());
 639     EXPECT_THROW(curs2 - curs1, std::out_of_range);
 640
 641     curs2.skip(7);
 642     EXPECT_EQ(2, curs1 - curs2);
 643     EXPECT_THROW(curs2 - curs1, std::out_of_range);
 644   }
 645
 646   // Test isAtEnd() with empty buffers at the end of a chain
 647   {
 648     auto iobuf1 = IOBuf::create(20);
 649     iobuf1->append(15);
 650     iobuf1->trimStart(5);
 651
 652     Cursor c(iobuf1.get());
 653     EXPECT_FALSE(c.isAtEnd());
 654     c.skip(10);
 655     EXPECT_TRUE(c.isAtEnd());
 656
 657     iobuf1->prependChain(IOBuf::create(10));
 658     iobuf1->prependChain(IOBuf::create(10));
 659     EXPECT_TRUE(c.isAtEnd());
 660     iobuf1->prev()->append(5);
 661     EXPECT_FALSE(c.isAtEnd());
 662     c.skip(5);
 663     EXPECT_TRUE(c.isAtEnd());
 664   }
 665
 666   // Test canAdvance with a chain of items
 667   {
 668     auto chain = IOBuf::create(10);
 669     chain->append(10);
 670     chain->appendChain(chain->clone());
 671     EXPECT_EQ(2, chain->countChainElements());
 672     EXPECT_EQ(20, chain->computeChainDataLength());
 673
 674     Cursor c(chain.get());
 675     for (size_t i = 0; i <= 20; ++i) {
 676       EXPECT_TRUE(c.canAdvance(i));
 677     }
 678     EXPECT_FALSE(c.canAdvance(21));
 679     c.skip(10);
 680     EXPECT_TRUE(c.canAdvance(10));
 681     EXPECT_FALSE(c.canAdvance(11));
 682   }
 683 }
 684
 685 TEST(IOBuf, StringOperations) {
 686   // Test a single buffer with two null-terminated strings and an extra uint8_t
 687   // at the end
 688   {
 689     std::unique_ptr<IOBuf> chain(IOBuf::create(16));
 690     Appender app(chain.get(), 0);
 691     app.push(reinterpret_cast<const uint8_t*>("hello\0world\0\x01"), 13);
 692
 693     Cursor curs(chain.get());
 694     EXPECT_STREQ("hello", curs.readTerminatedString().c_str());
 695     EXPECT_STREQ("world", curs.readTerminatedString().c_str());
 696     EXPECT_EQ(1, curs.read<uint8_t>());
 697   }
 698
 699   // Test multiple buffers where the first is empty and the string starts in
 700   // the second buffer.
 701   {
 702     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
 703     chain->prependChain(IOBuf::create(12));
 704     Appender app(chain.get(), 0);
 705     app.push(reinterpret_cast<const uint8_t*>("hello world\0"), 12);
 706
 707     Cursor curs(chain.get());
 708     EXPECT_STREQ("hello world", curs.readTerminatedString().c_str());
 709   }
 710
 711   // Test multiple buffers with a single null-terminated string spanning them
 712   {
 713     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
 714     chain->prependChain(IOBuf::create(8));
 715     chain->append(8);
 716     chain->next()->append(4);
 717     RWPrivateCursor rwc(chain.get());
 718     rwc.push(reinterpret_cast<const uint8_t*>("hello world\0"), 12);
 719
 720     Cursor curs(chain.get());
 721     EXPECT_STREQ("hello world", curs.readTerminatedString().c_str());
 722   }
 723
 724   // Test a reading a null-terminated string that's longer than the maximum
 725   // allowable length
 726   {
 727     std::unique_ptr<IOBuf> chain(IOBuf::create(16));
 728     Appender app(chain.get(), 0);
 729     app.push(reinterpret_cast<const uint8_t*>("hello world\0"), 12);
 730
 731     Cursor curs(chain.get());
 732     EXPECT_THROW(curs.readTerminatedString('\0', 5), std::length_error);
 733   }
 734
 735   // Test reading a null-terminated string from a chain with an empty buffer at
 736   // the front
 737   {
 738     std::unique_ptr<IOBuf> buf(IOBuf::create(8));
 739     Appender app(buf.get(), 0);
 740     app.push(reinterpret_cast<const uint8_t*>("hello\0"), 6);
 741     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
 742     chain->prependChain(std::move(buf));
 743
 744     Cursor curs(chain.get());
 745     EXPECT_STREQ("hello", curs.readTerminatedString().c_str());
 746   }
 747
 748   // Test reading a null-terminated string from a chain that doesn't contain the
 749   // terminator
 750   {
 751     std::unique_ptr<IOBuf> buf(IOBuf::create(8));
 752     Appender app(buf.get(), 0);
 753     app.push(reinterpret_cast<const uint8_t*>("hello"), 5);
 754     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
 755     chain->prependChain(std::move(buf));
 756
 757     Cursor curs(chain.get());
 758     EXPECT_THROW(curs.readTerminatedString(),
 759                  std::out_of_range);
 760   }
 761
 762   // Test reading a null-terminated string past the maximum length
 763   {
 764     std::unique_ptr<IOBuf> buf(IOBuf::create(8));
 765     Appender app(buf.get(), 0);
 766     app.push(reinterpret_cast<const uint8_t*>("hello\0"), 6);
 767     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
 768     chain->prependChain(std::move(buf));
 769
 770     Cursor curs(chain.get());
 771     EXPECT_THROW(curs.readTerminatedString('\0', 3),
 772                  std::length_error);
 773   }
 774
 775   // Test reading a two fixed-length strings from a single buffer with an extra
 776   // uint8_t at the end
 777   {
 778     std::unique_ptr<IOBuf> chain(IOBuf::create(16));
 779     Appender app(chain.get(), 0);
 780     app.push(reinterpret_cast<const uint8_t*>("helloworld\x01"), 11);
 781
 782     Cursor curs(chain.get());
 783     EXPECT_STREQ("hello", curs.readFixedString(5).c_str());
 784     EXPECT_STREQ("world", curs.readFixedString(5).c_str());
 785     EXPECT_EQ(1, curs.read<uint8_t>());
 786   }
 787
 788   // Test multiple buffers where the first is empty and a fixed-length string
 789   // starts in the second buffer.
 790   {
 791     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
 792     chain->prependChain(IOBuf::create(16));
 793     Appender app(chain.get(), 0);
 794     app.push(reinterpret_cast<const uint8_t*>("hello world"), 11);
 795
 796     Cursor curs(chain.get());
 797     EXPECT_STREQ("hello world", curs.readFixedString(11).c_str());
 798   }
 799
 800   // Test multiple buffers with a single fixed-length string spanning them
 801   {
 802     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
 803     chain->prependChain(IOBuf::create(8));
 804     chain->append(7);
 805     chain->next()->append(4);
 806     RWPrivateCursor rwc(chain.get());
 807     rwc.push(reinterpret_cast<const uint8_t*>("hello world"), 11);
 808
 809     Cursor curs(chain.get());
 810     EXPECT_STREQ("hello world", curs.readFixedString(11).c_str());
 811   }
 812
 813   // Test reading a fixed-length string from a chain with an empty buffer at
 814   // the front
 815   {
 816     std::unique_ptr<IOBuf> buf(IOBuf::create(8));
 817     Appender app(buf.get(), 0);
 818     app.push(reinterpret_cast<const uint8_t*>("hello"), 5);
 819     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
 820     chain->prependChain(std::move(buf));
 821
 822     Cursor curs(chain.get());
 823     EXPECT_STREQ("hello", curs.readFixedString(5).c_str());
 824   }
 825 }
 826
 827 TEST(IOBuf, ReadWhileTrue) {
 828   auto isAlpha = [](uint8_t ch) {
 829     return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z');
 830   };
 831   auto isDigit = [](uint8_t ch) { return (ch >= '0' && ch <= '9'); };
 832
 833   // Test reading alternating alphabetic and numeric strings
 834   {
 835     std::unique_ptr<IOBuf> chain(IOBuf::create(32));
 836     Appender app(chain.get(), 0);
 837     app.push(StringPiece("hello123world456"));
 838
 839     Cursor curs(chain.get());
 840     EXPECT_STREQ("hello", curs.readWhile(isAlpha).c_str());
 841     EXPECT_STREQ("123", curs.readWhile(isDigit).c_str());
 842     EXPECT_STREQ("world", curs.readWhile(isAlpha).c_str());
 843     EXPECT_STREQ("456", curs.readWhile(isDigit).c_str());
 844     EXPECT_TRUE(curs.isAtEnd());
 845   }
 846
 847   // The same, but also use skipWhile()
 848   {
 849     std::unique_ptr<IOBuf> chain(IOBuf::create(16));
 850     Appender app(chain.get(), 0);
 851     app.push(StringPiece("hello123world456"));
 852
 853     Cursor curs(chain.get());
 854     EXPECT_STREQ("hello", curs.readWhile(isAlpha).c_str());
 855     curs.skipWhile(isDigit);
 856     curs.skipWhile(isAlpha);
 857     EXPECT_STREQ("456", curs.readWhile(isDigit).c_str());
 858     EXPECT_TRUE(curs.isAtEnd());
 859   }
 860
 861   // Test readWhile() using data split across multiple buffers,
 862   // including some empty buffers in the middle of the chain.
 863   {
 864     std::unique_ptr<IOBuf> chain;
 865
 866     // First element in the chain has "he"
 867     auto buf = IOBuf::create(40);
 868     Appender app(buf.get(), 0);
 869     app.push(StringPiece("he"));
 870     chain = std::move(buf);
 871
 872     // The second element has "ll", after 10 bytes of headroom
 873     buf = IOBuf::create(40);
 874     buf->advance(10);
 875     app = Appender{buf.get(), 0};
 876     app.push(StringPiece("ll"));
 877     chain->prependChain(std::move(buf));
 878
 879     // The third element is empty
 880     buf = IOBuf::create(40);
 881     buf->advance(15);
 882     chain->prependChain(std::move(buf));
 883
 884     // The fourth element has "o12"
 885     buf = IOBuf::create(40);
 886     buf->advance(37);
 887     app = Appender{buf.get(), 0};
 888     app.push(StringPiece("o12"));
 889     chain->prependChain(std::move(buf));
 890
 891     // The fifth element has "3"
 892     buf = IOBuf::create(40);
 893     app = Appender{buf.get(), 0};
 894     app.push(StringPiece("3"));
 895     chain->prependChain(std::move(buf));
 896
 897     // The sixth element is empty
 898     buf = IOBuf::create(40);
 899     chain->prependChain(std::move(buf));
 900
 901     // The seventh element has "world456"
 902     buf = IOBuf::create(40);
 903     app = Appender{buf.get(), 0};
 904     app.push(StringPiece("world456"));
 905     chain->prependChain(std::move(buf));
 906
 907     // The eighth element is empty
 908     buf = IOBuf::create(40);
 909     chain->prependChain(std::move(buf));
 910
 911     Cursor curs(chain.get());
 912     EXPECT_STREQ("hello", curs.readWhile(isAlpha).c_str());
 913     EXPECT_STREQ("123", curs.readWhile(isDigit).c_str());
 914     EXPECT_STREQ("world", curs.readWhile(isAlpha).c_str());
 915     EXPECT_STREQ("456", curs.readWhile(isDigit).c_str());
 916     EXPECT_TRUE(curs.isAtEnd());
 917   }
 918 }
 919
 920 TEST(IOBuf, TestAdvanceToEndSingle) {
 921   std::unique_ptr<IOBuf> chain(IOBuf::create(10));
 922   chain->append(10);
 923
 924   Cursor curs(chain.get());
 925   curs.advanceToEnd();
 926   EXPECT_TRUE(curs.isAtEnd());
 927   EXPECT_EQ(curs - chain.get(), 10);
 928 }
 929
 930 TEST(IOBuf, TestAdvanceToEndMulti) {
 931   std::unique_ptr<IOBuf> chain(IOBuf::create(10));
 932   chain->append(10);
 933
 934   std::unique_ptr<IOBuf> buf(IOBuf::create(5));
 935   buf->append(5);
 936   chain->prependChain(std::move(buf));
 937
 938   buf = IOBuf::create(20);
 939   buf->append(20);
 940   chain->prependChain(std::move(buf));
 941
 942   Cursor curs(chain.get());
 943   curs.advanceToEnd();
 944   EXPECT_TRUE(curs.isAtEnd());
 945   EXPECT_EQ(curs - chain.get(), 35);
 946
 947   curs.reset(chain.get());
 948   curs.skip(12);
 949   curs.advanceToEnd();
 950   EXPECT_TRUE(curs.isAtEnd());
 951 }
 952
 953 TEST(IOBuf, TestRetreatSingle) {
 954   std::unique_ptr<IOBuf> chain(IOBuf::create(20));
 955   chain->append(20);
 956
 957   Cursor curs(chain.get());
 958   EXPECT_EQ(curs.retreatAtMost(0), 0);
 959   EXPECT_EQ(curs.totalLength(), 20);
 960   EXPECT_EQ(curs.retreatAtMost(5), 0);
 961   EXPECT_EQ(curs.totalLength(), 20);
 962   EXPECT_EQ(curs.retreatAtMost(25), 0);
 963   EXPECT_EQ(curs.totalLength(), 20);
 964
 965   curs.retreat(0);
 966   EXPECT_THROW(curs.retreat(5), std::out_of_range);
 967   curs.reset(chain.get());
 968   EXPECT_THROW(curs.retreat(25), std::out_of_range);
 969   curs.reset(chain.get());
 970
 971   curs.advanceToEnd();
 972   curs.retreat(5);
 973   EXPECT_EQ(curs.totalLength(), 5);
 974   curs.retreat(10);
 975   EXPECT_EQ(curs.totalLength(), 15);
 976   EXPECT_THROW(curs.retreat(10), std::out_of_range);
 977
 978   curs.reset(chain.get());
 979   curs.advanceToEnd();
 980   EXPECT_EQ(curs.retreatAtMost(5), 5);
 981   EXPECT_EQ(curs.totalLength(), 5);
 982   EXPECT_EQ(curs.retreatAtMost(10), 10);
 983   EXPECT_EQ(curs.totalLength(), 15);
 984   EXPECT_EQ(curs.retreatAtMost(10), 5);
 985   EXPECT_EQ(curs.totalLength(), 20);
 986 }
 987
 988 TEST(IOBuf, TestRetreatMulti) {
 989   std::unique_ptr<IOBuf> chain(IOBuf::create(10));
 990   chain->append(10);
 991
 992   std::unique_ptr<IOBuf> buf(IOBuf::create(5));
 993   buf->append(5);
 994   chain->prependChain(std::move(buf));
 995
 996   buf = IOBuf::create(20);
 997   buf->append(20);
 998   chain->prependChain(std::move(buf));
 999
1000   Cursor curs(chain.get());
1001   EXPECT_EQ(curs.retreatAtMost(10), 0);
1002   EXPECT_THROW(curs.retreat(10), std::out_of_range);
1003   curs.reset(chain.get());
1004
1005   curs.advanceToEnd();
1006   curs.retreat(20);
1007   EXPECT_EQ(curs.totalLength(), 20);
1008   EXPECT_EQ(curs.length(), 20);
1009   curs.retreat(1);
1010   EXPECT_EQ(curs.totalLength(), 21);
1011   EXPECT_EQ(curs.length(), 1);
1012   EXPECT_EQ(curs.retreatAtMost(50), 14);
1013   EXPECT_EQ(curs.totalLength(), 35);
1014
1015   curs.advanceToEnd();
1016   curs.retreat(30);
1017   EXPECT_EQ(curs.totalLength(), 30);
1018 }
1019
1020 TEST(IOBuf, TestRetreatOperators) {
1021   std::unique_ptr<IOBuf> chain(IOBuf::create(20));
1022   chain->append(20);
1023
1024   Cursor curs(chain.get());
1025   curs.advanceToEnd();
1026   curs -= 5;
1027   EXPECT_EQ(curs.totalLength(), 5);
1028
1029   curs.advanceToEnd();
1030   auto retreated = curs - 5;
1031   EXPECT_EQ(retreated.totalLength(), 5);
1032   EXPECT_EQ(curs.totalLength(), 0);
1033 }
1034
1035 TEST(IOBuf, tryRead) {
1036   unique_ptr<IOBuf> iobuf1(IOBuf::create(6));
1037   iobuf1->append(6);
1038   unique_ptr<IOBuf> iobuf2(IOBuf::create(24));
1039   iobuf2->append(24);
1040
1041   iobuf1->prependChain(std::move(iobuf2));
1042
1043   EXPECT_TRUE(iobuf1->isChained());
1044
1045   RWPrivateCursor wcursor(iobuf1.get());
1046   Cursor rcursor(iobuf1.get());
1047   wcursor.writeLE((uint32_t)1);
1048   wcursor.writeLE((uint64_t)1);
1049   wcursor.writeLE((uint64_t)1);
1050   wcursor.writeLE((uint64_t)1);
1051   wcursor.writeLE((uint16_t)1);
1052   EXPECT_EQ(0, wcursor.totalLength());
1053
1054   EXPECT_EQ(1u, rcursor.readLE<uint32_t>());
1055
1056   EXPECT_EQ(1u, rcursor.readLE<uint32_t>());
1057   EXPECT_EQ(0u, rcursor.readLE<uint32_t>());
1058
1059   EXPECT_EQ(1u, rcursor.readLE<uint32_t>());
1060   rcursor.skip(4);
1061
1062   uint32_t val;
1063   EXPECT_TRUE(rcursor.tryRead(val));
1064   EXPECT_EQ(1, val);
1065   EXPECT_TRUE(rcursor.tryRead(val));
1066
1067   EXPECT_EQ(0, val);
1068   EXPECT_FALSE(rcursor.tryRead(val));
1069 }
1070
1071 TEST(IOBuf, tryReadLE) {
1072   IOBuf buf{IOBuf::CREATE, 4};
1073   buf.append(4);
1074
1075   RWPrivateCursor wcursor(&buf);
1076   Cursor rcursor(&buf);
1077
1078   const uint32_t expected = 0x01020304;
1079   wcursor.writeLE(expected);
1080   uint32_t actual;
1081   EXPECT_TRUE(rcursor.tryReadLE(actual));
1082   EXPECT_EQ(expected, actual);
1083 }
1084
1085 TEST(IOBuf, tryReadBE) {
1086   IOBuf buf{IOBuf::CREATE, 4};
1087   buf.append(4);
1088
1089   RWPrivateCursor wcursor(&buf);
1090   Cursor rcursor(&buf);
1091
1092   const uint32_t expected = 0x01020304;
1093   wcursor.writeBE(expected);
1094   uint32_t actual;
1095   EXPECT_TRUE(rcursor.tryReadBE(actual));
1096   EXPECT_EQ(expected, actual);
1097 }
1098
1099 TEST(IOBuf, tryReadConsumesAllInputOnFailure) {
1100   IOBuf buf{IOBuf::CREATE, 2};
1101   buf.append(2);
1102
1103   Cursor rcursor(&buf);
1104   uint32_t val;
1105   EXPECT_FALSE(rcursor.tryRead(val));
1106   EXPECT_EQ(0, rcursor.totalLength());
1107 }
1108
1109 TEST(IOBuf, readConsumesAllInputOnFailure) {
1110   IOBuf buf{IOBuf::CREATE, 2};
1111   buf.append(2);
1112
1113   Cursor rcursor(&buf);
1114   EXPECT_THROW(rcursor.read<uint32_t>(), std::out_of_range);
1115   EXPECT_EQ(0, rcursor.totalLength());
1116 }
1117
1118 TEST(IOBuf, pushEmptyByteRange) {
1119   // Test pushing an empty ByteRange.  This mainly tests that we do not
1120   // trigger UBSAN warnings by calling memcpy() with an null source pointer,
1121   // which is undefined behavior even if the length is 0.
1122   IOBuf buf{IOBuf::CREATE, 2};
1123   ByteRange emptyBytes;
1124
1125   // Test calling Cursor::push()
1126   RWPrivateCursor wcursor(&buf);
1127   wcursor.push(emptyBytes);
1128   EXPECT_EQ(0, buf.computeChainDataLength());
1129
1130   // Test calling Appender::push()
1131   Appender app(&buf, 16);
1132   app.push(emptyBytes);
1133   EXPECT_EQ(0, buf.computeChainDataLength());
1134 }