Add DCHECKs for checking that underlying IOBuf wasn't modified

[folly.git] / folly / io / Cursor.h

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index bdaeabe3d292bf0a9d67f2b3e9be1a9fd34b3783..be231d9b6246e31974274fff3cad1b2ae1c709f2 100644 (file)
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -73,10 +73,11 @@ class CursorBase {
   template <class OtherDerived, class OtherBuf>
   explicit CursorBase(const CursorBase<OtherDerived, OtherBuf>& cursor)
       : crtBuf_(cursor.crtBuf_),
+        buffer_(cursor.buffer_),
         crtBegin_(cursor.crtBegin_),
         crtEnd_(cursor.crtEnd_),
         crtPos_(cursor.crtPos_),
-        buffer_(cursor.buffer_) {}
+        absolutePos_(cursor.absolutePos_) {}
 
   /**
    * Reset cursor to point to a new buffer.
@@ -84,13 +85,23 @@ class CursorBase {
   void reset(BufType* buf) {
     crtBuf_ = buf;
     buffer_ = buf;
+    absolutePos_ = 0;
     if (crtBuf_) {
       crtPos_ = crtBegin_ = crtBuf_->data();
       crtEnd_ = crtBuf_->tail();
     }
   }
 
+  /**
+   * Get the current Cursor position relative to the head of IOBuf chain.
+   */
+  size_t getCurrentPosition() const {
+    dcheckIntegrity();
+    return (crtPos_ - crtBegin_) + absolutePos_;
+  }
+
   const uint8_t* data() const {
+    dcheckIntegrity();
     return crtPos_;
   }
 
@@ -103,6 +114,7 @@ class CursorBase {
    * pointing at the end of a buffer.
    */
   size_t length() const {
+    dcheckIntegrity();
     return crtEnd_ - crtPos_;
   }
 
@@ -143,6 +155,7 @@ class CursorBase {
    * Return true if the cursor is at the end of the entire IOBuf chain.
    */
   bool isAtEnd() const {
+    dcheckIntegrity();
     // Check for the simple cases first.
     if (crtPos_ != crtEnd_) {
       return false;
@@ -167,10 +180,21 @@ class CursorBase {
    * Advances the cursor to the end of the entire IOBuf chain.
    */
   void advanceToEnd() {
-    crtBegin_ = buffer_->prev()->data();
-    crtPos_ = crtEnd_ = buffer_->prev()->tail();
-    if (crtBuf_ != buffer_->prev()) {
-      crtBuf_ = buffer_->prev();
+    // Simple case, we're already in the last IOBuf.
+    if (crtBuf_ == buffer_->prev()) {
+      crtPos_ = crtEnd_;
+      return;
+    }
+
+    auto* nextBuf = crtBuf_->next();
+    while (nextBuf != buffer_) {
+      absolutePos_ += crtEnd_ - crtBegin_;
+
+      crtBuf_ = nextBuf;
+      nextBuf = crtBuf_->next();
+      crtBegin_ = crtBuf_->data();
+      crtPos_ = crtEnd_ = crtBuf_->tail();
+
       static_cast<Derived*>(this)->advanceDone();
     }
   }
@@ -333,6 +357,7 @@ class CursorBase {
   void skipWhile(const Predicate& predicate);
 
   size_t skipAtMost(size_t len) {
+    dcheckIntegrity();
     if (LIKELY(crtPos_ + len < crtEnd_)) {
       crtPos_ += len;
       return len;
@@ -341,6 +366,7 @@ class CursorBase {
   }
 
   void skip(size_t len) {
+    dcheckIntegrity();
     if (LIKELY(crtPos_ + len < crtEnd_)) {
       crtPos_ += len;
     } else {
@@ -358,6 +384,7 @@ class CursorBase {
   }
 
   size_t retreatAtMost(size_t len) {
+    dcheckIntegrity();
     if (len <= static_cast<size_t>(crtPos_ - crtBegin_)) {
       crtPos_ -= len;
       return len;
@@ -366,6 +393,7 @@ class CursorBase {
   }
 
   void retreat(size_t len) {
+    dcheckIntegrity();
     if (len <= static_cast<size_t>(crtPos_ - crtBegin_)) {
       crtPos_ -= len;
     } else {
@@ -374,6 +402,7 @@ class CursorBase {
   }
 
   size_t pullAtMost(void* buf, size_t len) {
+    dcheckIntegrity();
     // Fast path: it all fits in one buffer.
     if (LIKELY(crtPos_ + len <= crtEnd_)) {
       memcpy(buf, data(), len);
@@ -384,6 +413,7 @@ class CursorBase {
   }
 
   void pull(void* buf, size_t len) {
+    dcheckIntegrity();
     if (LIKELY(crtPos_ + len <= crtEnd_)) {
       memcpy(buf, data(), len);
       crtPos_ += len;
@@ -531,6 +561,14 @@ class CursorBase {
   }
 
  protected:
+  void dcheckIntegrity() const {
+    DCHECK(crtBegin_ <= crtPos_ && crtPos_ <= crtEnd_);
+    DCHECK(crtBuf_ == nullptr || crtBegin_ == crtBuf_->data());
+    DCHECK(
+        crtBuf_ == nullptr ||
+        (uint64_t)(crtEnd_ - crtBegin_) == crtBuf_->length());
+  }
+
   ~CursorBase() { }
 
   BufType* head() {
@@ -544,6 +582,7 @@ class CursorBase {
       return false;
     }
 
+    absolutePos_ += crtEnd_ - crtBegin_;
     crtBuf_ = nextBuf;
     crtPos_ = crtBegin_ = crtBuf_->data();
     crtEnd_ = crtBuf_->tail();
@@ -559,20 +598,24 @@ class CursorBase {
     crtBuf_ = crtBuf_->prev();
     crtBegin_ = crtBuf_->data();
     crtPos_ = crtEnd_ = crtBuf_->tail();
+    absolutePos_ -= crtEnd_ - crtBegin_;
     static_cast<Derived*>(this)->advanceDone();
     return true;
   }
 
   void advanceBufferIfEmpty() {
+    dcheckIntegrity();
     if (crtPos_ == crtEnd_) {
       tryAdvanceBuffer();
     }
   }
 
   BufType* crtBuf_;
+  BufType* buffer_;
   const uint8_t* crtBegin_{nullptr};
   const uint8_t* crtEnd_{nullptr};
   const uint8_t* crtPos_{nullptr};
+  size_t absolutePos_{0};
 
  private:
   template <class T>
@@ -660,8 +703,6 @@ class CursorBase {
 
   void advanceDone() {
   }
-
-  BufType* buffer_;
 };
 
 } // namespace detail
@@ -800,6 +841,7 @@ class RWCursor
     this->crtPos_ = this->crtBegin_ + offset;
   }
   void gatherAtMost(size_t n) {
+    this->dcheckIntegrity();
     size_t size = std::min(n, this->totalLength());
     size_t offset = this->crtPos_ - this->crtBegin_;
     this->crtBuf_->gather(offset + size);
@@ -845,12 +887,13 @@ class RWCursor
   }
 
   void insert(std::unique_ptr<folly::IOBuf> buf) {
-    folly::IOBuf* nextBuf;
-    if (this->crtPos_ == this->crtBegin_) {
+    this->dcheckIntegrity();
+    this->absolutePos_ += buf->computeChainDataLength();
+    if (this->crtPos_ == this->crtBegin_ && this->crtBuf_ != this->buffer_) {
       // Can just prepend
-      nextBuf = this->crtBuf_;
       this->crtBuf_->prependChain(std::move(buf));
     } else {
+      IOBuf* nextBuf;
       std::unique_ptr<folly::IOBuf> remaining;
       if (this->crtPos_ != this->crtEnd_) {
         // Need to split current IOBuf in two.
@@ -863,16 +906,27 @@ class RWCursor
         nextBuf = this->crtBuf_->next();
       }
       this->crtBuf_->trimEnd(this->length());
+      this->absolutePos_ += this->crtPos_ - this->crtBegin_;
       this->crtBuf_->appendChain(std::move(buf));
 
-      // Jump past the new links
-      this->crtBuf_ = nextBuf;
-      this->crtPos_ = this->crtBegin_ = this->crtBuf_->data();
-      this->crtEnd_ = this->crtBuf_->tail();
+      if (nextBuf == this->buffer_) {
+        // We've just appended to the end of the buffer, so advance to the end.
+        this->crtBuf_ = this->buffer_->prev();
+        this->crtBegin_ = this->crtBuf_->data();
+        this->crtPos_ = this->crtEnd_ = this->crtBuf_->tail();
+        // This has already been accounted for, so remove it.
+        this->absolutePos_ -= this->crtEnd_ - this->crtBegin_;
+      } else {
+        // Jump past the new links
+        this->crtBuf_ = nextBuf;
+        this->crtPos_ = this->crtBegin_ = this->crtBuf_->data();
+        this->crtEnd_ = this->crtBuf_->tail();
+      }
     }
   }
 
   uint8_t* writableData() {
+    this->dcheckIntegrity();
     return this->crtBuf_->writableData() + (this->crtPos_ - this->crtBegin_);
   }