Fix decompression of truncated data
[folly.git] / folly / io / test / IOBufTest.cpp
index 72a35dba33ec163186c2b7dc4b7effb0ade4034c..e33287bb5aa55d5c1d244f00d76a3c7e03ae4e26 100644 (file)
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
 #include <folly/io/IOBuf.h>
 #include <folly/io/TypedIOBuf.h>
 
-#include <gflags/gflags.h>
+#include <cstddef>
+
 #include <boost/random.hpp>
-#include <gtest/gtest.h>
 
 #include <folly/Malloc.h>
 #include <folly/Range.h>
+#include <folly/portability/GTest.h>
 
 using folly::fbstring;
 using folly::fbvector;
@@ -185,6 +186,14 @@ TEST(IOBuf, WrapBuffer) {
   EXPECT_EQ(size3, iobuf3.length());
   EXPECT_EQ(buf3.get(), iobuf3.buffer());
   EXPECT_EQ(size3, iobuf3.capacity());
+
+  const uint32_t size4 = 2345;
+  unique_ptr<uint8_t[]> buf4(new uint8_t[size4]);
+  IOBuf iobuf4 = IOBuf::wrapBufferAsValue(buf4.get(), size4);
+  EXPECT_EQ(buf4.get(), iobuf4.data());
+  EXPECT_EQ(size4, iobuf4.length());
+  EXPECT_EQ(buf4.get(), iobuf4.buffer());
+  EXPECT_EQ(size4, iobuf4.capacity());
 }
 
 TEST(IOBuf, CreateCombined) {
@@ -522,6 +531,15 @@ TEST(IOBuf, Chaining) {
   gen.seed(fillSeed);
   checkChain(chainClone.get(), gen);
 
+  // cloneCoalesced
+  {
+    auto chainCloneCoalesced = chainClone->cloneCoalesced();
+    EXPECT_EQ(1, chainCloneCoalesced->countChainElements());
+    EXPECT_EQ(fullLength, chainCloneCoalesced->computeChainDataLength());
+    gen.seed(fillSeed);
+    checkChain(chainCloneCoalesced.get(), gen);
+  }
+
   // Coalesce the entire chain
   chainClone->coalesce();
   EXPECT_EQ(1, chainClone->countChainElements());
@@ -626,7 +644,6 @@ TEST(IOBuf, Reserve) {
     EXPECT_EQ(0, iob->headroom());
     EXPECT_EQ(100, iob->length());
     const void* p1 = iob->buffer();
-    const uint8_t* d1 = iob->data();
     iob->reserve(100, 2512);  // allocation sizes are multiples of 256
     EXPECT_LE(100, iob->headroom());
     if (folly::usingJEMalloc()) {
@@ -709,6 +726,11 @@ TEST(IOBuf, maybeCopyBuffer) {
   EXPECT_EQ(nullptr, buf.get());
 }
 
+TEST(IOBuf, copyEmptyBuffer) {
+  auto buf = IOBuf::copyBuffer(nullptr, 0);
+  EXPECT_EQ(buf->length(), 0);
+}
+
 namespace {
 
 int customDeleterCount = 0;
@@ -792,11 +814,7 @@ TEST(IOBuf, takeOwnershipUniquePtr) {
 }
 
 TEST(IOBuf, Alignment) {
-  // max_align_t doesn't exist in gcc 4.6.2
-  struct MaxAlign {
-    char c;
-  } __attribute__((__aligned__));
-  size_t alignment = alignof(MaxAlign);
+  size_t alignment = alignof(std::max_align_t);
 
   std::vector<size_t> sizes {0, 1, 64, 256, 1024, 1 << 10};
   for (size_t size : sizes) {
@@ -832,7 +850,7 @@ enum BufType {
 class MoveToFbStringTest
   : public ::testing::TestWithParam<std::tr1::tuple<int, int, bool, BufType>> {
  protected:
-  void SetUp() {
+  void SetUp() override {
     elementSize_ = std::tr1::get<0>(GetParam());
     elementCount_ = std::tr1::get<1>(GetParam());
     shared_ = std::tr1::get<2>(GetParam());
@@ -879,7 +897,6 @@ class MoveToFbStringTest
       }
       default:
         throw std::invalid_argument("unexpected buffer type parameter");
-        break;
     }
     memset(buf->writableData(), 'x', elementSize_);
     return buf;
@@ -1081,9 +1098,299 @@ TEST(IOBuf, HashAndEqual) {
   EXPECT_EQ(hash(e), hash(f));
 }
 
-int main(int argc, char** argv) {
-  testing::InitGoogleTest(&argc, argv);
-  gflags::ParseCommandLineFlags(&argc, &argv, true);
+// reserveSlow() had a bug when reallocating the buffer in place. It would
+// preserve old headroom if it's not too much (heuristically) but wouldn't
+// adjust the requested amount of memory to account for that; the end result
+// would be that reserve() would return with less tailroom than requested.
+TEST(IOBuf, ReserveWithHeadroom) {
+  // This is assuming jemalloc, where we know that 4096 and 8192 bytes are
+  // valid (and consecutive) allocation sizes. We're hoping that our
+  // 4096-byte buffer can be expanded in place to 8192 (in practice, this
+  // usually happens).
+  const char data[] = "Lorem ipsum dolor sit amet, consectetur adipiscing elit";
+  constexpr size_t reservedSize = 24;  // sizeof(SharedInfo)
+  // chosen carefully so that the buffer is exactly 4096 bytes
+  IOBuf buf(IOBuf::CREATE, 4096 - reservedSize);
+  buf.advance(10);
+  memcpy(buf.writableData(), data, sizeof(data));
+  buf.append(sizeof(data));
+  EXPECT_EQ(sizeof(data), buf.length());
+
+  // Grow the buffer (hopefully in place); this would incorrectly reserve
+  // the 10 bytes of headroom, giving us 10 bytes less than requested.
+  size_t tailroom = 8192 - reservedSize - sizeof(data);
+  buf.reserve(0, tailroom);
+  EXPECT_LE(tailroom, buf.tailroom());
+  EXPECT_EQ(sizeof(data), buf.length());
+  EXPECT_EQ(0, memcmp(data, buf.data(), sizeof(data)));
+}
+
+TEST(IOBuf, CopyConstructorAndAssignmentOperator) {
+  auto buf = IOBuf::create(4096);
+  append(buf, "hello world");
+  auto buf2 = IOBuf::create(4096);
+  append(buf2, " goodbye");
+  buf->prependChain(std::move(buf2));
+  EXPECT_FALSE(buf->isShared());
+
+  {
+    auto copy = *buf;
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_TRUE(copy.isShared());
+    EXPECT_EQ((void*)buf->data(), (void*)copy.data());
+    EXPECT_NE(buf->next(), copy.next());  // actually different buffers
+
+    auto copy2 = *buf;
+    copy2.coalesce();
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_TRUE(copy.isShared());
+    EXPECT_FALSE(copy2.isShared());
+
+    auto p = reinterpret_cast<const char*>(copy2.data());
+    EXPECT_EQ("hello world goodbye", std::string(p, copy2.length()));
+  }
+
+  EXPECT_FALSE(buf->isShared());
+
+  {
+    folly::IOBuf newBuf(folly::IOBuf::CREATE, 4096);
+    EXPECT_FALSE(newBuf.isShared());
+
+    auto newBufCopy = newBuf;
+    EXPECT_TRUE(newBuf.isShared());
+    EXPECT_TRUE(newBufCopy.isShared());
+
+    newBufCopy = *buf;
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_FALSE(newBuf.isShared());
+    EXPECT_TRUE(newBufCopy.isShared());
+  }
+
+  EXPECT_FALSE(buf->isShared());
+}
+
+TEST(IOBuf, CloneAsValue) {
+  auto buf = IOBuf::create(4096);
+  append(buf, "hello world");
+  {
+    auto buf2 = IOBuf::create(4096);
+    append(buf2, " goodbye");
+    buf->prependChain(std::move(buf2));
+    EXPECT_FALSE(buf->isShared());
+  }
+
+  {
+    auto copy = buf->cloneOneAsValue();
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_TRUE(copy.isShared());
+    EXPECT_EQ((void*)buf->data(), (void*)copy.data());
+    EXPECT_TRUE(buf->isChained());
+    EXPECT_FALSE(copy.isChained());
+
+    auto copy2 = buf->cloneAsValue();
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_TRUE(copy.isShared());
+    EXPECT_TRUE(copy2.isShared());
+    EXPECT_TRUE(buf->isChained());
+    EXPECT_TRUE(copy2.isChained());
+
+    copy.unshareOne();
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_FALSE(copy.isShared());
+    EXPECT_NE((void*)buf->data(), (void*)copy.data());
+    EXPECT_TRUE(copy2.isShared());
+
+    auto p = reinterpret_cast<const char*>(copy.data());
+    EXPECT_EQ("hello world", std::string(p, copy.length()));
+
+    copy2.coalesce();
+    EXPECT_FALSE(buf->isShared());
+    EXPECT_FALSE(copy.isShared());
+    EXPECT_FALSE(copy2.isShared());
+    EXPECT_FALSE(copy2.isChained());
+
+    auto p2 = reinterpret_cast<const char*>(copy2.data());
+    EXPECT_EQ("hello world goodbye", std::string(p2, copy2.length()));
+  }
+
+  EXPECT_FALSE(buf->isShared());
+}
+
+namespace {
+// Use with string literals only
+std::unique_ptr<IOBuf> wrap(const char* str) {
+  return IOBuf::wrapBuffer(str, strlen(str));
+}
+
+std::unique_ptr<IOBuf> copy(const char* str) {
+  // At least 1KiB of tailroom, to ensure an external buffer
+  return IOBuf::copyBuffer(str, strlen(str), 0, 1024);
+}
+
+std::string toString(const folly::IOBuf& buf) {
+  std::string result;
+  result.reserve(buf.computeChainDataLength());
+  for (auto& b : buf) {
+    result.append(reinterpret_cast<const char*>(b.data()), b.size());
+  }
+  return result;
+}
+
+char* writableStr(folly::IOBuf& buf) {
+  return reinterpret_cast<char*>(buf.writableData());
+}
+
+}  // namespace
 
-  return RUN_ALL_TESTS();
+TEST(IOBuf, ExternallyShared) {
+  struct Item {
+    Item(const char* src, size_t len) : size(len) {
+      CHECK_LE(len, sizeof(buffer));
+      memcpy(buffer, src, len);
+    }
+    uint32_t refcount{0};
+    uint8_t size;
+    char buffer[256];
+  };
+
+  auto hello = "hello";
+  struct Item it(hello, strlen(hello));
+
+  {
+    auto freeFn = [](void* /* unused */, void* userData) {
+      auto it = static_cast<struct Item*>(userData);
+      it->refcount--;
+    };
+    it.refcount++;
+    auto buf1 = IOBuf::takeOwnership(it.buffer, it.size, freeFn, &it);
+    EXPECT_TRUE(buf1->isManagedOne());
+    EXPECT_FALSE(buf1->isSharedOne());
+
+    buf1->markExternallyShared();
+    EXPECT_TRUE(buf1->isSharedOne());
+
+    {
+      auto buf2 = buf1->clone();
+      EXPECT_TRUE(buf2->isManagedOne());
+      EXPECT_TRUE(buf2->isSharedOne());
+      EXPECT_EQ(buf1->data(), buf2->data());
+      EXPECT_EQ(it.refcount, 1);
+    }
+    EXPECT_EQ(it.refcount, 1);
+  }
+  EXPECT_EQ(it.refcount, 0);
+}
+
+TEST(IOBuf, Managed) {
+  auto hello = "hello";
+  auto buf1UP = wrap(hello);
+  auto buf1 = buf1UP.get();
+  EXPECT_FALSE(buf1->isManagedOne());
+  auto buf2UP = copy("world");
+  auto buf2 = buf2UP.get();
+  EXPECT_TRUE(buf2->isManagedOne());
+  auto buf3UP = wrap(hello);
+  auto buf3 = buf3UP.get();
+  auto buf4UP = buf2->clone();
+  auto buf4 = buf4UP.get();
+
+  // buf1 and buf3 share the same memory (but are unmanaged)
+  EXPECT_FALSE(buf1->isManagedOne());
+  EXPECT_FALSE(buf3->isManagedOne());
+  EXPECT_TRUE(buf1->isSharedOne());
+  EXPECT_TRUE(buf3->isSharedOne());
+  EXPECT_EQ(buf1->data(), buf3->data());
+
+  // buf2 and buf4 share the same memory (but are managed)
+  EXPECT_TRUE(buf2->isManagedOne());
+  EXPECT_TRUE(buf4->isManagedOne());
+  EXPECT_TRUE(buf2->isSharedOne());
+  EXPECT_TRUE(buf4->isSharedOne());
+  EXPECT_EQ(buf2->data(), buf4->data());
+
+  buf1->prependChain(std::move(buf2UP));
+  buf1->prependChain(std::move(buf3UP));
+  buf1->prependChain(std::move(buf4UP));
+
+  EXPECT_EQ("helloworldhelloworld", toString(*buf1));
+  EXPECT_FALSE(buf1->isManaged());
+
+  buf1->makeManaged();
+  EXPECT_TRUE(buf1->isManaged());
+
+  // buf1 and buf3 are now unshared (because they were unmanaged)
+  EXPECT_TRUE(buf1->isManagedOne());
+  EXPECT_TRUE(buf3->isManagedOne());
+  EXPECT_FALSE(buf1->isSharedOne());
+  EXPECT_FALSE(buf3->isSharedOne());
+  EXPECT_NE(buf1->data(), buf3->data());
+
+  // buf2 and buf4 are still shared
+  EXPECT_TRUE(buf2->isManagedOne());
+  EXPECT_TRUE(buf4->isManagedOne());
+  EXPECT_TRUE(buf2->isSharedOne());
+  EXPECT_TRUE(buf4->isSharedOne());
+  EXPECT_EQ(buf2->data(), buf4->data());
+
+  // And verify that the truth is what we expect: modify a byte in buf1 and
+  // buf2, see that the change from buf1 is *not* reflected in buf3, but the
+  // change from buf2 is reflected in buf4.
+  writableStr(*buf1)[0] = 'j';
+  writableStr(*buf2)[0] = 'x';
+  EXPECT_EQ("jelloxorldhelloxorld", toString(*buf1));
+}
+
+TEST(IOBuf, CoalesceEmptyBuffers) {
+  auto b1 = IOBuf::takeOwnership(nullptr, 0);
+  auto b2 = fromStr("hello");
+  auto b3 = IOBuf::takeOwnership(nullptr, 0);
+
+  b2->appendChain(std::move(b3));
+  b1->appendChain(std::move(b2));
+
+  auto br = b1->coalesce();
+
+  EXPECT_TRUE(ByteRange(StringPiece("hello")) == br);
+}
+
+TEST(IOBuf, CloneCoalescedChain) {
+  auto b = IOBuf::createChain(1000, 100);
+  b->advance(10);
+  const uint32_t fillSeed = 0x12345678;
+  boost::mt19937 gen(fillSeed);
+  {
+    auto c = b.get();
+    uint64_t length = c->tailroom();
+    do {
+      length = std::min(length, c->tailroom());
+      c->append(length--);
+      fillBuf(c, gen);
+      c = c->next();
+    } while (c != b.get());
+  }
+  auto c = b->cloneCoalescedAsValue();
+  EXPECT_FALSE(c.isChained()); // Not chained
+  EXPECT_FALSE(c.isSharedOne()); // Not shared
+  EXPECT_EQ(b->headroom(), c.headroom()); // Preserves headroom
+  EXPECT_LE(b->prev()->tailroom(), c.tailroom()); // Preserves minimum tailroom
+  EXPECT_EQ(b->computeChainDataLength(), c.length()); // Same length
+  gen.seed(fillSeed);
+  checkBuf(&c, gen); // Same contents
+}
+
+TEST(IOBuf, CloneCoalescedSingle) {
+  auto b = IOBuf::create(1000);
+  b->advance(10);
+  b->append(900);
+  const uint32_t fillSeed = 0x12345678;
+  boost::mt19937 gen(fillSeed);
+  fillBuf(b.get(), gen);
+
+  auto c = b->cloneCoalesced();
+  EXPECT_FALSE(c->isChained()); // Not chained
+  EXPECT_TRUE(c->isSharedOne()); // Shared
+  EXPECT_EQ(b->buffer(), c->buffer());
+  EXPECT_EQ(b->capacity(), c->capacity());
+  EXPECT_EQ(b->data(), c->data());
+  EXPECT_EQ(b->length(), c->length());
 }