Add boundary support into Cursor

Summary:
Added boundary support in folly/io/Cursor.h.

The bounded Cursor can have a boundary that is shorter than the underlying IOBuf. This gives a custom view of the IOBuf.

Reviewed By: spalamarchuk

Differential Revision: D16347838

fbshipit-source-id: 16d91e42e2acee788b31bc2a9b6cc8716b1d22f2

folly/io/Cursor.h

@@ -66,6 +66,17 @@ class CursorBase {
     }
   }
 
+  CursorBase(BufType* buf, size_t len) : crtBuf_(buf), buffer_(buf) {
+    if (crtBuf_) {
+      crtPos_ = crtBegin_ = crtBuf_->data();
+      crtEnd_ = crtBuf_->tail();
+      if (crtPos_ + len < crtEnd_) {
+        crtEnd_ = crtPos_ + len;
+      }
+      remainingLen_ = len - (crtEnd_ - crtPos_);
+    }
+  }
+
   /**
    * Copy constructor.
    *
@@ -79,7 +90,27 @@ class CursorBase {
         crtBegin_(cursor.crtBegin_),
         crtEnd_(cursor.crtEnd_),
         crtPos_(cursor.crtPos_),
-        absolutePos_(cursor.absolutePos_) {}
+        absolutePos_(cursor.absolutePos_),
+        remainingLen_(cursor.remainingLen_) {}
+
+  template <class OtherDerived, class OtherBuf>
+  explicit CursorBase(
+      const CursorBase<OtherDerived, OtherBuf>& cursor,
+      size_t len)
+      : crtBuf_(cursor.crtBuf_),
+        buffer_(cursor.buffer_),
+        crtBegin_(cursor.crtBegin_),
+        crtEnd_(cursor.crtEnd_),
+        crtPos_(cursor.crtPos_),
+        absolutePos_(cursor.absolutePos_) {
+    if (cursor.isBounded() && len > cursor.remainingLen_ + cursor.length()) {
+      throw_exception<std::out_of_range>("underflow");
+    }
+    if (crtPos_ + len < crtEnd_) {
+      crtEnd_ = crtPos_ + len;
+    }
+    remainingLen_ = len - (crtEnd_ - crtPos_);
+  }
 
   /**
    * Reset cursor to point to a new buffer.
@@ -88,6 +119,7 @@ class CursorBase {
     crtBuf_ = buf;
     buffer_ = buf;
     absolutePos_ = 0;
+    remainingLen_ = std::numeric_limits<size_t>::max();
     if (crtBuf_) {
       crtPos_ = crtBegin_ = crtBuf_->data();
       crtEnd_ = crtBuf_->tail();
@@ -122,14 +154,16 @@ class CursorBase {
 
   /**
    * Return the space available until the end of the entire IOBuf chain.
+   * For bounded Cursors, return the available space until the boundary.
    */
   size_t totalLength() const {
-    if (crtBuf_ == buffer_) {
-      return crtBuf_->computeChainDataLength() - (crtPos_ - crtBegin_);
+    size_t len = 0;
+    const IOBuf* buf = crtBuf_->next();
+    while (buf != buffer_ && len < remainingLen_) {
+      len += buf->length();
+      buf = buf->next();
     }
-    CursorBase end(buffer_->prev());
-    end.crtPos_ = end.crtEnd_;
-    return end - *this;
+    return std::min(len, remainingLen_) + length();
   }
 
   /**
@@ -140,6 +174,9 @@ class CursorBase {
    * walks the minimal set of buffers in the chain to determine the result.
    */
   bool canAdvance(size_t amount) const {
+    if (isBounded() && amount > remainingLen_ + length()) {
+      return false;
+    }
     const IOBuf* nextBuf = crtBuf_;
     size_t available = length();
     do {
@@ -165,11 +202,13 @@ class CursorBase {
     if (crtBuf_ == buffer_->prev()) {
       return true;
     }
+    if (isBounded() && remainingLen_ == 0) {
+      return true;
+    }
     // We are at the end of a buffer, but it isn't the last buffer.
     // We might still be at the end if the remaining buffers in the chain are
     // empty.
     const IOBuf* buf = crtBuf_->next();
-    ;
     while (buf != buffer_) {
       if (buf->length() > 0) {
         return false;
@@ -191,14 +230,24 @@ class CursorBase {
 
     auto* nextBuf = crtBuf_->next();
     while (nextBuf != buffer_) {
+      if (isBounded() && remainingLen_ == 0) {
+        crtPos_ = crtEnd_;
+        return;
+      }
       absolutePos_ += crtEnd_ - crtBegin_;
 
       crtBuf_ = nextBuf;
       nextBuf = crtBuf_->next();
       crtBegin_ = crtBuf_->data();
-      crtPos_ = crtEnd_ = crtBuf_->tail();
-
-      static_cast<Derived*>(this)->advanceDone();
+      crtEnd_ = crtBuf_->tail();
+      if (isBounded()) {
+        if (crtBegin_ + remainingLen_ < crtEnd_) {
+          crtEnd_ = crtBegin_ + remainingLen_;
+        }
+        remainingLen_ -= crtEnd_ - crtBegin_;
+      }
+      crtPos_ = crtEnd_;
+      derived().advanceDone();
     }
   }
 
@@ -523,7 +572,11 @@ class CursorBase {
     size_t len = 0;
 
     if (otherBuf != crtBuf_) {
-      len += other.crtEnd_ - other.crtPos_;
+      if (other.remainingLen_ == 0) {
+        len += otherBuf->tail() - other.crtPos_;
+      } else {
+        len += other.crtEnd_ - other.crtPos_;
+      }
 
       for (otherBuf = otherBuf->next();
            otherBuf != crtBuf_ && otherBuf != other.buffer_;
@@ -566,13 +619,17 @@ class CursorBase {
     return len;
   }
 
+  bool isBounded() const {
+    return remainingLen_ != std::numeric_limits<size_t>::max();
+  }
+
  protected:
   void dcheckIntegrity() const {
     DCHECK(crtBegin_ <= crtPos_ && crtPos_ <= crtEnd_);
     DCHECK(crtBuf_ == nullptr || crtBegin_ == crtBuf_->data());
     DCHECK(
         crtBuf_ == nullptr ||
-        (std::size_t)(crtEnd_ - crtBegin_) == crtBuf_->length());
+        (std::size_t)(crtEnd_ - crtBegin_) <= crtBuf_->length());
   }
 
   ~CursorBase() {}
@@ -583,7 +640,7 @@ class CursorBase {
 
   bool tryAdvanceBuffer() {
     BufType* nextBuf = crtBuf_->next();
-    if (UNLIKELY(nextBuf == buffer_)) {
+    if (UNLIKELY(nextBuf == buffer_) || remainingLen_ == 0) {
       crtPos_ = crtEnd_;
       return false;
     }
@@ -592,7 +649,13 @@ class CursorBase {
     crtBuf_ = nextBuf;
     crtPos_ = crtBegin_ = crtBuf_->data();
     crtEnd_ = crtBuf_->tail();
-    static_cast<Derived*>(this)->advanceDone();
+    if (isBounded()) {
+      if (crtPos_ + remainingLen_ < crtEnd_) {
+        crtEnd_ = crtPos_ + remainingLen_;
+      }
+      remainingLen_ -= crtEnd_ - crtPos_;
+    }
+    derived().advanceDone();
     return true;
   }
 
@@ -601,11 +664,14 @@ class CursorBase {
       crtPos_ = crtBegin_;
       return false;
     }
+    if (isBounded()) {
+      remainingLen_ += crtEnd_ - crtBegin_;
+    }
     crtBuf_ = crtBuf_->prev();
     crtBegin_ = crtBuf_->data();
     crtPos_ = crtEnd_ = crtBuf_->tail();
     absolutePos_ -= crtEnd_ - crtBegin_;
-    static_cast<Derived*>(this)->advanceDone();
+    derived().advanceDone();
     return true;
   }
 
@@ -622,8 +688,20 @@ class CursorBase {
   const uint8_t* crtEnd_{nullptr};
   const uint8_t* crtPos_{nullptr};
   size_t absolutePos_{0};
+  // For bounded Cursor, remainingLen_ is the remaining number of data bytes
+  // in subsequent IOBufs in the chain. For unbounded Cursor, remainingLen_
+  // is set to the max of size_t
+  size_t remainingLen_{std::numeric_limits<size_t>::max()};
 
  private:
+  Derived& derived() {
+    return static_cast<Derived&>(*this);
+  }
+
+  Derived const& derived() const {
+    return static_cast<const Derived&>(*this);
+  }
+
   template <class T>
   FOLLY_NOINLINE T readSlow() {
     T val;
@@ -723,9 +801,16 @@ class Cursor : public detail::CursorBase<Cursor, const IOBuf> {
   explicit Cursor(const IOBuf* buf)
       : detail::CursorBase<Cursor, const IOBuf>(buf) {}
 
+  explicit Cursor(const IOBuf* buf, size_t len)
+      : detail::CursorBase<Cursor, const IOBuf>(buf, len) {}
+
   template <class OtherDerived, class OtherBuf>
   explicit Cursor(const detail::CursorBase<OtherDerived, OtherBuf>& cursor)
       : detail::CursorBase<Cursor, const IOBuf>(cursor) {}
+
+  template <class OtherDerived, class OtherBuf>
+  Cursor(const detail::CursorBase<OtherDerived, OtherBuf>& cursor, size_t len)
+      : detail::CursorBase<Cursor, const IOBuf>(cursor, len) {}
 };
 
 namespace detail {
@@ -823,7 +908,10 @@ class RWCursor : public detail::CursorBase<RWCursor<access>, IOBuf>,
   template <class OtherDerived, class OtherBuf>
   explicit RWCursor(const detail::CursorBase<OtherDerived, OtherBuf>& cursor)
       : detail::CursorBase<RWCursor<access>, IOBuf>(cursor),
-        maybeShared_(true) {}
+        maybeShared_(true) {
+    CHECK(!cursor.isBounded())
+        << "Creating RWCursor from bounded Cursor is not allowed";
+  }
   /**
    * Gather at least n bytes contiguously into the current buffer,
    * by coalescing subsequent buffers from the chain as necessary.

folly/io/test/IOBufCursorTest.cpp

@@ -1210,3 +1210,288 @@ TEST(IOBuf, positionTracking) {
   EXPECT_EQ(0, cursor.getCurrentPosition());
   EXPECT_EQ(30, cursor.totalLength());
 }
+
+TEST(IOBuf, BoundedCursorSanity) {
+  // initialize bounded Cursor with length 0
+  std::unique_ptr<IOBuf> chain1(IOBuf::create(20));
+  chain1->append(10);
+  Cursor subC(chain1.get(), 0);
+  EXPECT_EQ(0, subC.totalLength());
+  EXPECT_EQ(0, subC.length());
+  EXPECT_EQ(0, subC - chain1.get());
+  EXPECT_TRUE(subC.isAtEnd());
+  EXPECT_FALSE(subC.canAdvance(1));
+  EXPECT_THROW(subC.skip(1), std::out_of_range);
+
+  // multi-item chain
+  chain1->appendChain(chain1->clone());
+  EXPECT_EQ(2, chain1->countChainElements());
+  EXPECT_EQ(20, chain1->computeChainDataLength());
+
+  Cursor subCurs1(chain1.get(), 0);
+  EXPECT_EQ(0, subCurs1.totalLength());
+  EXPECT_EQ(0, subCurs1.length());
+  EXPECT_TRUE(subCurs1.isAtEnd());
+  EXPECT_FALSE(subCurs1.canAdvance(1));
+  EXPECT_THROW(subCurs1.skip(1), std::out_of_range);
+
+  // initialize bounded Cursor with length greater than total IOBuf bytes
+  Cursor subCurs2(chain1.get(), 25);
+  EXPECT_EQ(20, subCurs2.totalLength());
+  EXPECT_FALSE(subCurs2.isAtEnd());
+  EXPECT_TRUE(subCurs2.canAdvance(1));
+  subCurs2.skip(10);
+  EXPECT_EQ(10, subCurs2.totalLength());
+  EXPECT_EQ(10, subCurs2 - chain1.get());
+  EXPECT_FALSE(subCurs2.isAtEnd());
+
+  subCurs2.skip(3);
+  EXPECT_EQ(7, subCurs2.totalLength());
+  EXPECT_EQ(13, subCurs2 - chain1.get());
+  EXPECT_EQ(7, subCurs2.skipAtMost(10));
+  EXPECT_TRUE(subCurs2.isAtEnd());
+
+  // initialize bounded Cursor with another bounded Cursor
+  Cursor subCurs3(chain1.get(), 16);
+  EXPECT_EQ(16, subCurs3.totalLength());
+  EXPECT_FALSE(subCurs3.isAtEnd());
+  subCurs3.skip(4);
+  EXPECT_EQ(12, subCurs3.totalLength());
+  EXPECT_FALSE(subCurs3.isAtEnd());
+  Cursor subCurs4(subCurs3, 9);
+  EXPECT_EQ(9, subCurs4.totalLength());
+  EXPECT_FALSE(subCurs4.isAtEnd());
+  EXPECT_EQ(0, subCurs3 - subCurs4);
+  subCurs4.skip(5);
+  EXPECT_EQ(4, subCurs4.totalLength());
+  EXPECT_FALSE(subCurs4.isAtEnd());
+  EXPECT_EQ(5, subCurs4 - subCurs3);
+  subCurs4.skip(4);
+  EXPECT_TRUE(subCurs4.isAtEnd());
+  EXPECT_EQ(0, subCurs4.totalLength());
+
+  // Initialize a bounded Cursor with length greater than the bytes left in the
+  // bounded Cursor in the ctor argument
+  EXPECT_THROW(Cursor(subCurs3, 20), std::out_of_range);
+  // Initialize a bounded Cursor with equal length as the bounded Cursor
+  // in the ctor arguments
+  Cursor subCurs5(subCurs3, 12);
+  EXPECT_EQ(12, subCurs5.totalLength());
+  EXPECT_FALSE(subCurs5.isAtEnd());
+
+  // Initialize a Cursor with bounded Cursor
+  Cursor subCurs6(chain1.get(), 17);
+  EXPECT_EQ(17, subCurs6.totalLength());
+  Cursor curs(subCurs6);
+  EXPECT_TRUE(curs.isBounded());
+  EXPECT_EQ(17, curs.totalLength());
+  EXPECT_EQ(0, curs - chain1.get());
+  EXPECT_FALSE(curs.isAtEnd());
+  curs.advanceToEnd();
+  EXPECT_EQ(0, curs.totalLength());
+  EXPECT_TRUE(curs.isAtEnd());
+  EXPECT_EQ(17, curs - chain1.get());
+
+  subCurs6.skip(7);
+  EXPECT_EQ(10, subCurs6.totalLength());
+  Cursor curs1(subCurs6);
+  EXPECT_TRUE(curs1.isBounded());
+  EXPECT_EQ(10, curs1.totalLength());
+  EXPECT_EQ(7, curs1 - chain1.get());
+  EXPECT_FALSE(curs1.isAtEnd());
+
+  // test reset
+  curs.reset(chain1.get());
+  EXPECT_FALSE(curs.isBounded());
+  EXPECT_EQ(20, curs.totalLength());
+  EXPECT_EQ(0, curs - chain1.get());
+  EXPECT_FALSE(curs.isAtEnd());
+}
+
+TEST(IOBuf, BoundedCursorOperators) {
+  // Test operators on a single-item chain
+  {
+    std::unique_ptr<IOBuf> chain1(IOBuf::create(20));
+    chain1->append(10);
+
+    Cursor subCurs1(chain1.get(), 8);
+    EXPECT_EQ(8, subCurs1.totalLength());
+    EXPECT_EQ(0, subCurs1 - chain1.get());
+    EXPECT_FALSE(subCurs1.isAtEnd());
+    subCurs1.skip(3);
+    EXPECT_EQ(5, subCurs1.totalLength());
+    EXPECT_EQ(3, subCurs1 - chain1.get());
+    EXPECT_FALSE(subCurs1.isAtEnd());
+    subCurs1.skip(5);
+    EXPECT_EQ(0, subCurs1.totalLength());
+    EXPECT_EQ(8, subCurs1 - chain1.get());
+    EXPECT_TRUE(subCurs1.isAtEnd());
+
+    Cursor curs1(chain1.get());
+    Cursor subCurs2(curs1, 6);
+    EXPECT_EQ(6, subCurs2.totalLength());
+    EXPECT_EQ(0, subCurs2 - chain1.get());
+    EXPECT_EQ(8, subCurs1 - subCurs2);
+    EXPECT_THROW(subCurs2 - subCurs1, std::out_of_range);
+    EXPECT_EQ(subCurs2.retreatAtMost(10), 0);
+    EXPECT_THROW(subCurs2.retreat(10), std::out_of_range);
+    subCurs2.advanceToEnd();
+    EXPECT_EQ(2, subCurs1 - subCurs2);
+    subCurs1.retreat(1);
+    EXPECT_EQ(1, subCurs1.totalLength());
+    EXPECT_EQ(subCurs1.retreatAtMost(10), 7);
+    EXPECT_EQ(0, subCurs1 - chain1.get());
+    EXPECT_EQ(8, subCurs1.totalLength());
+    EXPECT_EQ(6, subCurs2 - subCurs1);
+    subCurs1.advanceToEnd();
+    EXPECT_EQ(2, subCurs1 - subCurs2);
+  }
+
+  // Test cross-chain operations
+  {
+    std::unique_ptr<IOBuf> chain1(IOBuf::create(20));
+    chain1->append(10);
+    std::unique_ptr<IOBuf> chain2 = chain1->clone();
+
+    Cursor subCurs1(chain1.get(), 9);
+    Cursor subCurs2(chain2.get(), 9);
+    EXPECT_THROW(subCurs1 - subCurs2, std::out_of_range);
+    EXPECT_THROW(subCurs1 - chain2.get(), std::out_of_range);
+  }
+
+  // Test operations on multi-item chains
+  {
+    std::unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
+    iobuf1->append(10);
+    std::unique_ptr<IOBuf> iobuf2(IOBuf::create(20));
+    iobuf2->append(10);
+    std::unique_ptr<IOBuf> iobuf3(IOBuf::create(20));
+    iobuf3->append(10);
+    iobuf1->prependChain(std::move(iobuf2));
+    iobuf1->prependChain(std::move(iobuf3));
+    EXPECT_EQ(3, iobuf1->countChainElements());
+    EXPECT_EQ(30, iobuf1->computeChainDataLength());
+
+    Cursor curs1(iobuf1.get());
+    Cursor subCurs1(curs1, 28);
+    subCurs1.skip(5);
+    EXPECT_EQ(23, subCurs1.totalLength());
+    EXPECT_EQ(5, subCurs1 - iobuf1.get());
+    EXPECT_FALSE(subCurs1.isAtEnd());
+    Cursor subCurs2(iobuf1.get(), 15);
+    subCurs2.skip(3);
+    EXPECT_EQ(12, subCurs2.totalLength());
+    EXPECT_EQ(2, subCurs1 - subCurs2);
+    EXPECT_THROW(subCurs2 - subCurs1, std::out_of_range);
+
+    subCurs1.skip(7);
+    EXPECT_EQ(9, subCurs1 - subCurs2);
+    EXPECT_EQ(12, subCurs1 - iobuf1.get());
+    EXPECT_FALSE(subCurs1.isAtEnd());
+    EXPECT_EQ(subCurs1.length(), 8);
+    EXPECT_THROW(subCurs2 - subCurs1, std::out_of_range);
+
+    subCurs2.skip(7);
+    EXPECT_EQ(5, subCurs2.totalLength());
+    EXPECT_EQ(10, subCurs2 - iobuf1.get());
+    EXPECT_FALSE(subCurs2.isAtEnd());
+    EXPECT_EQ(2, subCurs1 - subCurs2);
+    EXPECT_THROW(subCurs2 - subCurs1, std::out_of_range);
+
+    subCurs1.advanceToEnd();
+    EXPECT_EQ(28, subCurs1 - iobuf1.get());
+    EXPECT_EQ(18, subCurs1 - subCurs2);
+    EXPECT_TRUE(subCurs1.isAtEnd());
+
+    subCurs1.retreat(8);
+    EXPECT_FALSE(subCurs1.isAtEnd());
+    EXPECT_EQ(subCurs1.totalLength(), 8);
+    EXPECT_EQ(subCurs1.length(), 8);
+    EXPECT_EQ(20, subCurs1 - iobuf1.get());
+    EXPECT_EQ(10, subCurs1 - subCurs2);
+    subCurs1.retreat(1);
+    EXPECT_EQ(subCurs1.totalLength(), 9);
+    EXPECT_EQ(subCurs1.length(), 1);
+    EXPECT_EQ(19, subCurs1 - iobuf1.get());
+    EXPECT_EQ(subCurs1.retreatAtMost(20), 19);
+    EXPECT_EQ(subCurs1.totalLength(), 28);
+    EXPECT_EQ(0, subCurs1 - iobuf1.get());
+    EXPECT_EQ(10, subCurs2 - subCurs1);
+    EXPECT_THROW(subCurs1.retreat(1), std::out_of_range);
+  }
+
+  // Test canAdvance with a chain of items
+  {
+    auto chain = IOBuf::create(10);
+    chain->append(10);
+    chain->appendChain(chain->clone());
+    EXPECT_EQ(2, chain->countChainElements());
+    EXPECT_EQ(20, chain->computeChainDataLength());
+
+    Cursor c(chain.get());
+    Cursor subC(c, 14);
+    for (size_t i = 0; i <= 14; ++i) {
+      EXPECT_TRUE(subC.canAdvance(i));
+    }
+    EXPECT_FALSE(subC.canAdvance(15));
+    subC.skip(10);
+    EXPECT_TRUE(subC.canAdvance(4));
+    EXPECT_FALSE(subC.canAdvance(5));
+  }
+}
+
+TEST(IOBuf, BoundedCursorPullAndPeek) {
+  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 abc");
+  iobuf1->prependChain(std::move(iobuf2));
+  iobuf1->prependChain(std::move(iobuf3));
+  EXPECT_EQ(3, iobuf1->countChainElements());
+  EXPECT_EQ(15, iobuf1->computeChainDataLength());
+
+  std::array<char, 20> buf;
+  buf.fill(0);
+  Cursor subCurs1(iobuf1.get(), 13);
+  subCurs1.pull(buf.data(), 13);
+  EXPECT_EQ("hello world a", std::string(buf.data()));
+  EXPECT_TRUE(subCurs1.isAtEnd());
+
+  buf.fill(0);
+  Cursor subCurs2(iobuf1.get(), 11);
+  EXPECT_EQ(11, subCurs2.pullAtMost(buf.data(), 20));
+  EXPECT_EQ("hello world", std::string(buf.data()));
+  EXPECT_TRUE(subCurs2.isAtEnd());
+
+  buf.fill(0);
+  Cursor subCurs3(iobuf1.get(), 2);
+  EXPECT_EQ(2, subCurs3.totalLength());
+  EXPECT_EQ(0, subCurs3 - iobuf1.get());
+  EXPECT_EQ(2, subCurs3.pullAtMost(buf.data(), 20));
+  EXPECT_EQ("he", std::string(buf.data()));
+  EXPECT_TRUE(subCurs3.isAtEnd());
+  EXPECT_EQ(0, subCurs3.totalLength());
+  // the following test makes sure that tryAdvanceBuffer inside pullAtMostSlow
+  // won't go over the entire IOBuf chain, but terminate at the right boundary
+  EXPECT_EQ(2, subCurs3 - iobuf1.get());
+
+  EXPECT_THROW(
+      { Cursor(iobuf1.get(), 11).pull(buf.data(), 20); }, std::out_of_range);
+
+  {
+    Cursor subCursor(iobuf1.get(), 13);
+    auto b = subCursor.peekBytes();
+    EXPECT_EQ("he", StringPiece(b));
+    subCursor.skip(b.size());
+    b = subCursor.peekBytes();
+    EXPECT_EQ("llo ", StringPiece(b));
+    subCursor.skip(b.size());
+    b = subCursor.peekBytes();
+    EXPECT_EQ("world a", StringPiece(b));
+    subCursor.skip(b.size());
+    EXPECT_EQ(3, iobuf1->countChainElements());
+    EXPECT_EQ(15, iobuf1->computeChainDataLength());
+  }
+}