Prevent IOBufQueue::append() from packing into a shared tail

Summary:
If a shared `IOBuf` is appended to `IOBufQueue` its tail is not considered writable, but we may write into it in the packing logic.

Fix that, and share the implementation between the pointer and value versions.

Reviewed By: luciang

Differential Revision: D31890708

fbshipit-source-id: 2b590cbcab7028af9421d29babb760869833bc65

folly/io/IOBufQueue.cpp

@@ -31,28 +31,36 @@ const size_t MIN_ALLOC_SIZE = 2000;
 const size_t MAX_ALLOC_SIZE = 8000;
 
 /**
- * Convenience function to append chain src to chain dst.
+ * Convenience functions to append chain src to chain dst.
  */
-void appendToChain(unique_ptr<IOBuf>& dst, unique_ptr<IOBuf>&& src, bool pack) {
-  if (dst == nullptr) {
-    dst = std::move(src);
-  } else {
-    IOBuf* tail = dst->prev();
-    if (pack) {
-      // Copy up to kMaxPackCopy bytes if we can free buffers; this helps
-      // reduce wastage (the tail's tailroom and the head's headroom) when
-      // joining two IOBufQueues together.
+template <class Src, class Next>
+void packInto(IOBuf* tail, Src& src, Next next) {
+  if (tail->isSharedOne()) {
+    return;
+  }
+
+  // Copy up to kMaxPackCopy bytes if we can free buffers; this helps reduce
+  // waste (the tail's tailroom and the head's headroom) when joining two
+  // IOBufQueues together.
   size_t copyRemaining = folly::IOBufQueue::kMaxPackCopy;
   std::size_t n;
-      while (src && (n = src->length()) <= copyRemaining &&
-             n <= tail->tailroom()) {
+  while (src && (n = src->length()) <= copyRemaining && n <= tail->tailroom()) {
     if (n > 0) {
       memcpy(tail->writableTail(), src->data(), n);
       tail->append(n);
       copyRemaining -= n;
     }
-        src = src->pop();
+    src = next(std::move(src));
   }
+}
+
+void appendToChain(unique_ptr<IOBuf>& dst, unique_ptr<IOBuf>&& src, bool pack) {
+  if (dst == nullptr) {
+    dst = std::move(src);
+  } else {
+    IOBuf* tail = dst->prev();
+    if (pack) {
+      packInto(tail, src, [](auto&& cur) { return cur->pop(); });
     }
     if (src) {
       tail->appendChain(std::move(src));
@@ -164,22 +172,14 @@ void IOBufQueue::append(
     chainLength_ += buf.computeChainDataLength();
   }
 
-  size_t copyRemaining = kMaxPackCopy;
-  std::size_t n;
-  const folly::IOBuf* src = &buf;
   folly::IOBuf* tail = head_->prev();
-  while ((n = src->length()) <= copyRemaining && n <= tail->tailroom()) {
-    if (n > 0) {
-      memcpy(tail->writableTail(), src->data(), n);
-      tail->append(n);
-      copyRemaining -= n;
-    }
-    src = src->next();
-
-    // Consumed full input.
-    if (src == &buf) {
-      return;
-    }
+  const folly::IOBuf* src = &buf;
+  packInto(tail, src, [&](auto&& cur) {
+    auto next = cur->next();
+    return next != &buf ? next : nullptr;
+  });
+  if (!src) {
+    return; // Consumed full input.
   }
 
   // Clone the rest.

folly/io/test/IOBufQueueTest.cpp

@@ -42,8 +42,9 @@ struct Initializer {
 };
 Initializer initializer;
 
-unique_ptr<IOBuf> stringToIOBuf(const char* s, size_t len) {
-  unique_ptr<IOBuf> buf = IOBuf::create(len);
+unique_ptr<IOBuf> stringToIOBuf(
+    const char* s, size_t len, size_t tailroom = 0) {
+  unique_ptr<IOBuf> buf = IOBuf::create(len + tailroom);
   memcpy(buf->writableTail(), s, len);
   buf->append(len);
   return buf;
@@ -601,3 +602,14 @@ TEST(IOBufQueue, ReuseTail) {
     }
   }
 }
+
+TEST(IOBufQueue, PackWithSharedTail) {
+  auto buf = stringToIOBuf("Hello ", 6, 10);
+  IOBufQueue queue;
+  queue.append(buf->clone());
+  *buf->writableTail() = 'X';
+  buf->append(1);
+  queue.append(stringToIOBuf("world", 5), /* pack */ true);
+  // buf is shared, packing should not modify it.
+  EXPECT_EQ(buf->data()[6], 'X');
+}