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

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index 73e37965713924dc905f23c67ef8bd0c8c1c8fa6..0b79c23f26b501f9b35583163a52ac32239e7bab 100644 (file)
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -54,13 +54,30 @@ class CursorBase {
     return crtBuf_->data() + offset_;
   }
 
     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_;
   }
 
   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);
   Derived& operator+=(size_t offset) {
     Derived* p = static_cast<Derived*>(this);
     p->skip(offset);


@@ -344,6 +361,10 @@ class CursorBase {
 
   ~CursorBase(){}
 
 
   ~CursorBase(){}
 

+  BufType* head() {
+    return buffer_;
+  }
+

   bool tryAdvanceBuffer() {
     BufType* nextBuf = crtBuf_->next();
     if (UNLIKELY(nextBuf == 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) {
    * 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);
   }
     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;
 
   size_t pushAtMost(const uint8_t* buf, size_t len) {
     size_t copied = 0;

diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index 65f7624deba5eaf5e1fc9bde3da8289e03595ff1..923cbb2228ebffa122b73ace1cc7141ed512ac95 100644 (file)
--- a/folly/io/IOBuf.cpp
+++ b/folly/io/IOBuf.cpp
@@ -431,11 +431,10 @@ void IOBuf::unshareChained() {
   coalesceSlow();
 }
 
   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.
   // 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;
 
   // 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_;
   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
   // 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,
 }
 
 void IOBuf::coalesceAndReallocate(size_t newHeadroom,

diff --git a/folly/io/IOBuf.h b/folly/io/IOBuf.h
index ebba90a8b9884a32876f08bcb5bb5dc9f268859d..72ef2f7f093eacf522210b0cf50c825d0eae693a 100644 (file)
--- a/folly/io/IOBuf.h
+++ b/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.
    *
    * 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
    * 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 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(
   // 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);
       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();
   void decrementRefcount();
   void reserveSlow(uint32_t minHeadroom, uint32_t minTailroom);
   void freeExtBuffer();

diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index d349241b93fda7e2ebef0e20727181775436d43b..5cc6896792c71f2d21ba4b0c1718588678ef24e8 100644 (file)
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/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");
 TEST(IOBuf, cloneAndInsert) {
   std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
   append(iobuf1, "he");

diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index 5fc86644b7a8f84a7105374074bceec6e85ff8ba..baab1c854e0b6e02177f38f391819aff3ee8ca8b 100644 (file)
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -486,6 +486,8 @@ TEST(IOBuf, Chaining) {
   EXPECT_EQ(0, arrayBufFreeCount);
 
   // Buffer lengths: 1500 20 1234 900 321
   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());
   // Coalesce the first 3 buffers
   chainClone->gather(1521);
   EXPECT_EQ(3, chainClone->countChainElements());