make IOBuf::gather() safe

Summary:
Update IOBuf::gather() and RWCursor::gather() to check their arguments
more carefully, and throw std::overflow_errors if the caller attempts to
gather more data than is available.

The comments for IOBuf::gather() claimed that it ensured that maxLength
bytes would always be available when it returned.  However, if the total
chain length was less than maxLength, it would simply coalesce the full
chain and return successfully, even though fewer than maxLength bytes
were available.  This fixes the code to throw a std::overflow_error
rather than coalescing the chain.

Additional, this updates RWCursor::gather() to ensure that it does not
attempt to gather past the end of the IOBuf chain.  Previously it could
gather past the end of the chain, coalescing the head of the chain into
the tail.  This would free the IOBuf head, which was owned by an
external caller.

A new RWCursor::gatherAtMost() API is provided for callers that wish to
gather as much as requested if possible, but still succeed if less than
this is available.

Test Plan:
Updated the unit tests to test calling gather() with more bytes than
actually available.

Reviewed By: davejwatson@fb.com

FB internal diff: D1081995

folly/io/Cursor.h

@@ -54,13 +54,30 @@ class CursorBase {
     return crtBuf_->data() + offset_;
   }
 
-  // Space available in the current IOBuf.  May be 0; use peek() instead which
-  // will always point to a non-empty chunk of data or at the end of the
-  // chain.
+  /*
+   * Return the remaining space available in the current IOBuf.
+   *
+   * May return 0 if the cursor is at the end of an IOBuf.  Use peek() instead
+   * if you want to avoid this.  peek() will advance to the next non-empty
+   * IOBuf (up to the end of the chain) if the cursor is currently pointing at
+   * the end of a buffer.
+   */
   size_t length() const {
     return crtBuf_->length() - offset_;
   }
 
+  /*
+   * Return the space available until the end of the entire IOBuf chain.
+   */
+  size_t totalLength() const {
+    if (crtBuf_ == buffer_) {
+      return crtBuf_->computeChainDataLength() - offset_;
+    }
+    CursorBase end(buffer_->prev());
+    end.offset_ = end.buffer_->length();
+    return end - *this;
+  }
+
   Derived& operator+=(size_t offset) {
     Derived* p = static_cast<Derived*>(this);
     p->skip(offset);
@@ -344,6 +361,10 @@ class CursorBase {
 
   ~CursorBase(){}
 
+  BufType* head() {
+    return buffer_;
+  }
+
   bool tryAdvanceBuffer() {
     BufType* nextBuf = crtBuf_->next();
     if (UNLIKELY(nextBuf == buffer_)) {
@@ -431,8 +452,23 @@ class RWCursor
    * by coalescing subsequent buffers from the chain as necessary.
    */
   void gather(size_t n) {
+    // Forbid attempts to gather beyond the end of this IOBuf chain.
+    // Otherwise we could try to coalesce the head of the chain and end up
+    // accidentally freeing it, invalidating the pointer owned by external
+    // code.
+    //
+    // If crtBuf_ == head() then IOBuf::gather() will perform all necessary
+    // checking.  We only have to perform an explicit check here when calling
+    // gather() on a non-head element.
+    if (this->crtBuf_ != this->head() && this->totalLength() < n) {
+      throw std::overflow_error("cannot gather() past the end of the chain");
+    }
     this->crtBuf_->gather(this->offset_ + n);
   }
+  void gatherAtMost(size_t n) {
+    size_t size = std::min(n, this->totalLength());
+    return this->crtBuf_->gather(this->offset_ + size);
+  }
 
   size_t pushAtMost(const uint8_t* buf, size_t len) {
     size_t copied = 0;

folly/io/IOBuf.cpp

@@ -431,11 +431,10 @@ void IOBuf::unshareChained() {
   coalesceSlow();
 }
 
-void IOBuf::coalesceSlow(size_t maxLength) {
+void IOBuf::coalesceSlow() {
   // coalesceSlow() should only be called if we are part of a chain of multiple
   // IOBufs.  The caller should have already verified this.
-  assert(isChained());
-  assert(length_ < maxLength);
+  DCHECK(isChained());
 
   // Compute the length of the entire chain
   uint64_t newLength = 0;
@@ -443,13 +442,37 @@ void IOBuf::coalesceSlow(size_t maxLength) {
   do {
     newLength += end->length_;
     end = end->next_;
-  } while (newLength < maxLength && end != this);
+  } while (end != this);
 
-  uint64_t newHeadroom = headroom();
-  uint64_t newTailroom = end->prev_->tailroom();
-  coalesceAndReallocate(newHeadroom, newLength, end, newTailroom);
+  coalesceAndReallocate(newLength, end);
   // We should be only element left in the chain now
-  assert(length_ >= maxLength || !isChained());
+  DCHECK(!isChained());
+}
+
+void IOBuf::coalesceSlow(size_t maxLength) {
+  // coalesceSlow() should only be called if we are part of a chain of multiple
+  // IOBufs.  The caller should have already verified this.
+  DCHECK(isChained());
+  DCHECK_LT(length_, maxLength);
+
+  // Compute the length of the entire chain
+  uint64_t newLength = 0;
+  IOBuf* end = this;
+  while (true) {
+    newLength += end->length_;
+    end = end->next_;
+    if (newLength >= maxLength) {
+      break;
+    }
+    if (end == this) {
+      throw std::overflow_error("attempted to coalesce more data than "
+                                "available");
+    }
+  }
+
+  coalesceAndReallocate(newLength, end);
+  // We should have the requested length now
+  DCHECK_GE(length_, maxLength);
 }
 
 void IOBuf::coalesceAndReallocate(size_t newHeadroom,

folly/io/IOBuf.h

@@ -941,8 +941,9 @@ class IOBuf {
    * first IOBuf in the chain, and at least as much tailroom as the last IOBuf
    * that was coalesced.
    *
-   * Throws std::bad_alloc on error.  On error the IOBuf chain will be
-   * unmodified.  Throws std::overflow_error if the length of the coalesced
+   * Throws std::bad_alloc or std::overflow_error on error.  On error the IOBuf
+   * chain will be unmodified.  Throws std::overflow_error if maxLength is
+   * longer than the total chain length, or if the length of the coalesced
    * portion of the chain is larger than can be described by a uint32_t
    * capacity.  (Although maxLength is uint32_t, gather() doesn't split
    * buffers, so coalescing whole buffers may result in a capacity that can't
@@ -1066,7 +1067,8 @@ class IOBuf {
 
   void unshareOneSlow();
   void unshareChained();
-  void coalesceSlow(size_t maxLength=std::numeric_limits<size_t>::max());
+  void coalesceSlow();
+  void coalesceSlow(size_t maxLength);
   // newLength must be the entire length of the buffers between this and
   // end (no truncation)
   void coalesceAndReallocate(
@@ -1074,6 +1076,9 @@ class IOBuf {
       size_t newLength,
       IOBuf* end,
       size_t newTailroom);
+  void coalesceAndReallocate(size_t newLength, IOBuf* end) {
+    coalesceAndReallocate(headroom(), newLength, end, end->prev_->tailroom());
+  }
   void decrementRefcount();
   void reserveSlow(uint32_t minHeadroom, uint32_t minTailroom);
   void freeExtBuffer();

folly/io/test/IOBufCursorTest.cpp

@@ -236,6 +236,48 @@ TEST(IOBuf, PullAndPeek) {
   }
 }
 
+TEST(IOBuf, Gather) {
+  std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
+  append(iobuf1, "he");
+  std::unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
+  append(iobuf2, "llo ");
+  std::unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
+  append(iobuf3, "world");
+  iobuf1->prependChain(std::move(iobuf2));
+  iobuf1->prependChain(std::move(iobuf3));
+  EXPECT_EQ(3, iobuf1->countChainElements());
+  EXPECT_EQ(11, iobuf1->computeChainDataLength());
+
+  // Attempting to gather() more data than available in the chain should fail.
+  // Try from the very beginning of the chain.
+  RWPrivateCursor cursor(iobuf1.get());
+  EXPECT_THROW(cursor.gather(15), std::overflow_error);
+  // Now try from the middle of the chain
+  cursor += 3;
+  EXPECT_THROW(cursor.gather(10), std::overflow_error);
+
+  // Calling gatherAtMost() should succeed, however, and just gather
+  // as much as it can
+  cursor.gatherAtMost(10);
+  EXPECT_EQ(8, cursor.length());
+  EXPECT_EQ(8, cursor.totalLength());
+  EXPECT_EQ("lo world",
+            folly::StringPiece(reinterpret_cast<const char*>(cursor.data()),
+                               cursor.length()));
+  EXPECT_EQ(2, iobuf1->countChainElements());
+  EXPECT_EQ(11, iobuf1->computeChainDataLength());
+
+  // Now try gather again on the chain head
+  cursor = RWPrivateCursor(iobuf1.get());
+  cursor.gather(5);
+  // Since gather() doesn't split buffers, everything should be collapsed into
+  // a single buffer now.
+  EXPECT_EQ(1, iobuf1->countChainElements());
+  EXPECT_EQ(11, iobuf1->computeChainDataLength());
+  EXPECT_EQ(11, cursor.length());
+  EXPECT_EQ(11, cursor.totalLength());
+}
+
 TEST(IOBuf, cloneAndInsert) {
   std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
   append(iobuf1, "he");

folly/io/test/IOBufTest.cpp

@@ -486,6 +486,8 @@ TEST(IOBuf, Chaining) {
   EXPECT_EQ(0, arrayBufFreeCount);
 
   // Buffer lengths: 1500 20 1234 900 321
+  // Attempting to gather more data than available should fail
+  EXPECT_THROW(chainClone->gather(4000), std::overflow_error);
   // Coalesce the first 3 buffers
   chainClone->gather(1521);
   EXPECT_EQ(3, chainClone->countChainElements());