1 /*
2  * Copyright (c) Facebook, Inc. and its affiliates.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <cassert>
20 #include <cstdarg>
21 #include <cstring>
22 #include <memory>
23 #include <stdexcept>
24 #include <type_traits>
25 
26 #include <folly/Likely.h>
27 #include <folly/Memory.h>
28 #include <folly/Portability.h>
29 #include <folly/Range.h>
30 #include <folly/io/IOBuf.h>
31 #include <folly/io/IOBufQueue.h>
32 #include <folly/lang/Bits.h>
33 #include <folly/lang/Exception.h>
34 
35 /**
36  * Cursor class for fast iteration over IOBuf chains.
37  *
38  * Cursor - Read-only access
39  *
40  * RWPrivateCursor - Read-write access, assumes private access to IOBuf chain
41  * RWUnshareCursor - Read-write access, calls unshare on write (COW)
42  * Appender        - Write access, assumes private access to IOBuf chain
43  *
44  * Note that RW cursors write in the preallocated part of buffers (that is,
45  * between the buffer's data() and tail()), while Appenders append to the end
46  * of the buffer (between the buffer's tail() and bufferEnd()).  Appenders
47  * automatically adjust the buffer pointers, so you may only use one
48  * Appender with a buffer chain; for this reason, Appenders assume private
49  * access to the buffer (you need to call unshare() yourself if necessary).
50  **/
51 namespace folly {
52 namespace io {
53 
54 namespace detail {
55 
56 template <class Derived, class BufType>
57 class CursorBase {
58   // Make all the templated classes friends for copy constructor.
59   template <class D, typename B>
60   friend class CursorBase;
61 
62  public:
CursorBase(BufType * buf)63   explicit CursorBase(BufType* buf) : crtBuf_(buf), buffer_(buf) {
64     if (crtBuf_) {
65       crtPos_ = crtBegin_ = crtBuf_->data();
66       crtEnd_ = crtBuf_->tail();
67     }
68   }
69 
CursorBase(BufType * buf,size_t len)70   CursorBase(BufType* buf, size_t len) : crtBuf_(buf), buffer_(buf) {
71     if (crtBuf_) {
72       crtPos_ = crtBegin_ = crtBuf_->data();
73       crtEnd_ = crtBuf_->tail();
74       if (crtPos_ + len < crtEnd_) {
75         crtEnd_ = crtPos_ + len;
76       }
77       remainingLen_ = len - (crtEnd_ - crtPos_);
78     }
79   }
80 
81   /**
82    * Copy constructor.
83    *
84    * This also allows constructing a CursorBase from other derived types.
85    * For instance, this allows constructing a Cursor from an RWPrivateCursor.
86    */
87   template <class OtherDerived, class OtherBuf>
CursorBase(const CursorBase<OtherDerived,OtherBuf> & cursor)88   explicit CursorBase(const CursorBase<OtherDerived, OtherBuf>& cursor)
89       : crtBuf_(cursor.crtBuf_),
90         buffer_(cursor.buffer_),
91         crtBegin_(cursor.crtBegin_),
92         crtEnd_(cursor.crtEnd_),
93         crtPos_(cursor.crtPos_),
94         absolutePos_(cursor.absolutePos_),
95         remainingLen_(cursor.remainingLen_) {}
96 
97   template <class OtherDerived, class OtherBuf>
CursorBase(const CursorBase<OtherDerived,OtherBuf> & cursor,size_t len)98   explicit CursorBase(
99       const CursorBase<OtherDerived, OtherBuf>& cursor, size_t len)
100       : crtBuf_(cursor.crtBuf_),
101         buffer_(cursor.buffer_),
102         crtBegin_(cursor.crtBegin_),
103         crtEnd_(cursor.crtEnd_),
104         crtPos_(cursor.crtPos_),
105         absolutePos_(cursor.absolutePos_) {
106     if (cursor.isBounded() && len > cursor.remainingLen_ + cursor.length()) {
107       throw_exception<std::out_of_range>("underflow");
108     }
109     if (crtPos_ + len < crtEnd_) {
110       crtEnd_ = crtPos_ + len;
111     }
112     remainingLen_ = len - (crtEnd_ - crtPos_);
113   }
114 
115   /**
116    * Reset cursor to point to a new buffer.
117    */
reset(BufType * buf)118   void reset(BufType* buf) {
119     crtBuf_ = buf;
120     buffer_ = buf;
121     absolutePos_ = 0;
122     remainingLen_ = std::numeric_limits<size_t>::max();
123     if (crtBuf_) {
124       crtPos_ = crtBegin_ = crtBuf_->data();
125       crtEnd_ = crtBuf_->tail();
126     }
127   }
128 
129   /**
130    * Get the current Cursor position relative to the head of IOBuf chain.
131    */
getCurrentPosition()132   size_t getCurrentPosition() const {
133     dcheckIntegrity();
134     return (crtPos_ - crtBegin_) + absolutePos_;
135   }
136 
data()137   const uint8_t* data() const {
138     dcheckIntegrity();
139     return crtPos_;
140   }
141 
142   /**
143    * Return the remaining space available in the current IOBuf.
144    *
145    * May return 0 if the cursor is at the end of an IOBuf.  Use peekBytes()
146    * instead if you want to avoid this.  peekBytes() will advance to the next
147    * non-empty IOBuf (up to the end of the chain) if the cursor is currently
148    * pointing at the end of a buffer.
149    */
length()150   size_t length() const {
151     dcheckIntegrity();
152     return crtEnd_ - crtPos_;
153   }
154 
155   /**
156    * Return the space available until the end of the entire IOBuf chain.
157    * For bounded Cursors, return the available space until the boundary.
158    */
totalLength()159   size_t totalLength() const {
160     size_t len = 0;
161     const IOBuf* buf = crtBuf_->next();
162     while (buf != buffer_ && len < remainingLen_) {
163       len += buf->length();
164       buf = buf->next();
165     }
166     return std::min(len, remainingLen_) + length();
167   }
168 
169   /**
170    * Return true if the cursor could advance the specified number of bytes
171    * from its current position.
172    * This is useful for applications that want to do checked reads instead of
173    * catching exceptions and is more efficient than using totalLength as it
174    * walks the minimal set of buffers in the chain to determine the result.
175    */
canAdvance(size_t amount)176   bool canAdvance(size_t amount) const {
177     if (isBounded() && amount > remainingLen_ + length()) {
178       return false;
179     }
180     const IOBuf* nextBuf = crtBuf_;
181     size_t available = length();
182     do {
183       if (available >= amount) {
184         return true;
185       }
186       amount -= available;
187       nextBuf = nextBuf->next();
188       available = nextBuf->length();
189     } while (nextBuf != buffer_);
190     return false;
191   }
192 
193   /*
194    * Return true if the cursor is at the end of the entire IOBuf chain.
195    */
isAtEnd()196   bool isAtEnd() const {
197     dcheckIntegrity();
198     // Check for the simple cases first.
199     if (crtPos_ != crtEnd_) {
200       return false;
201     }
202     if (crtBuf_ == buffer_->prev()) {
203       return true;
204     }
205     if (isBounded() && remainingLen_ == 0) {
206       return true;
207     }
208     // We are at the end of a buffer, but it isn't the last buffer.
209     // We might still be at the end if the remaining buffers in the chain are
210     // empty.
211     const IOBuf* buf = crtBuf_->next();
212     while (buf != buffer_) {
213       if (buf->length() > 0) {
214         return false;
215       }
216       buf = buf->next();
217     }
218     return true;
219   }
220 
221   /**
222    * Advances the cursor to the end of the entire IOBuf chain.
223    */
advanceToEnd()224   void advanceToEnd() {
225     // Simple case, we're already in the last IOBuf.
226     if (crtBuf_ == buffer_->prev()) {
227       crtPos_ = crtEnd_;
228       return;
229     }
230 
231     auto* nextBuf = crtBuf_->next();
232     while (nextBuf != buffer_) {
233       if (isBounded() && remainingLen_ == 0) {
234         crtPos_ = crtEnd_;
235         return;
236       }
237       absolutePos_ += crtEnd_ - crtBegin_;
238 
239       crtBuf_ = nextBuf;
240       nextBuf = crtBuf_->next();
241       crtBegin_ = crtBuf_->data();
242       crtEnd_ = crtBuf_->tail();
243       if (isBounded()) {
244         if (crtBegin_ + remainingLen_ < crtEnd_) {
245           crtEnd_ = crtBegin_ + remainingLen_;
246         }
247         remainingLen_ -= crtEnd_ - crtBegin_;
248       }
249       crtPos_ = crtEnd_;
250       derived().advanceDone();
251     }
252   }
253 
254   Derived& operator+=(size_t offset) {
255     Derived* p = static_cast<Derived*>(this);
256     p->skip(offset);
257     return *p;
258   }
259   Derived operator+(size_t offset) const {
260     Derived other(*this);
261     other.skip(offset);
262     return other;
263   }
264 
265   Derived& operator-=(size_t offset) {
266     Derived* p = static_cast<Derived*>(this);
267     p->retreat(offset);
268     return *p;
269   }
270   Derived operator-(size_t offset) const {
271     Derived other(*this);
272     other.retreat(offset);
273     return other;
274   }
275 
276   /**
277    * Compare cursors for equality/inequality.
278    *
279    * Two cursors are equal if they are pointing to the same location in the
280    * same IOBuf chain.
281    */
282   bool operator==(const Derived& other) const {
283     const IOBuf* crtBuf = crtBuf_;
284     auto crtPos = crtPos_;
285     // We can be pointing to the end of a buffer chunk, find first non-empty.
286     while (crtPos == crtBuf->tail() && crtBuf != buffer_->prev()) {
287       crtBuf = crtBuf->next();
288       crtPos = crtBuf->data();
289     }
290 
291     const IOBuf* crtBufOther = other.crtBuf_;
292     auto crtPosOther = other.crtPos_;
293     // We can be pointing to the end of a buffer chunk, find first non-empty.
294     while (crtPosOther == crtBufOther->tail() &&
295            crtBufOther != other.buffer_->prev()) {
296       crtBufOther = crtBufOther->next();
297       crtPosOther = crtBufOther->data();
298     }
299     return (crtPos == crtPosOther) && (crtBuf == crtBufOther);
300   }
301   bool operator!=(const Derived& other) const { return !operator==(other); }
302 
303   template <class T>
tryRead(T & val)304   typename std::enable_if<std::is_arithmetic<T>::value, bool>::type tryRead(
305       T& val) {
306     if (LIKELY(crtPos_ + sizeof(T) <= crtEnd_)) {
307       val = loadUnaligned<T>(data());
308       crtPos_ += sizeof(T);
309       return true;
310     }
311     return pullAtMostSlow(&val, sizeof(T)) == sizeof(T);
312   }
313 
314   template <class T>
tryReadBE(T & val)315   bool tryReadBE(T& val) {
316     const bool result = tryRead(val);
317     val = Endian::big(val);
318     return result;
319   }
320 
321   template <class T>
tryReadLE(T & val)322   bool tryReadLE(T& val) {
323     const bool result = tryRead(val);
324     val = Endian::little(val);
325     return result;
326   }
327 
328   template <class T>
read()329   T read() {
330     if (LIKELY(crtPos_ + sizeof(T) <= crtEnd_)) {
331       T val = loadUnaligned<T>(data());
332       crtPos_ += sizeof(T);
333       return val;
334     } else {
335       return readSlow<T>();
336     }
337   }
338 
339   template <class T>
readBE()340   T readBE() {
341     return Endian::big(read<T>());
342   }
343 
344   template <class T>
readLE()345   T readLE() {
346     return Endian::little(read<T>());
347   }
348 
349   /**
350    * Read a fixed-length string.
351    *
352    * The std::string-based APIs should probably be avoided unless you
353    * ultimately want the data to live in an std::string. You're better off
354    * using the pull() APIs to copy into a raw buffer otherwise.
355    */
readFixedString(size_t len)356   std::string readFixedString(size_t len) {
357     std::string str;
358     str.reserve(len);
359     if (LIKELY(length() >= len)) {
360       str.append(reinterpret_cast<const char*>(data()), len);
361       crtPos_ += len;
362     } else {
363       readFixedStringSlow(&str, len);
364     }
365     return str;
366   }
367 
368   /**
369    * Read a string consisting of bytes until the given terminator character is
370    * seen. Raises an std::length_error if maxLength bytes have been processed
371    * before the terminator is seen.
372    *
373    * See comments in readFixedString() about when it's appropriate to use this
374    * vs. using pull().
375    */
376   std::string readTerminatedString(
377       char termChar = '\0',
378       size_t maxLength = std::numeric_limits<size_t>::max());
379 
380   /*
381    * Read all bytes until the specified predicate returns true.
382    *
383    * The predicate will be called on each byte in turn, until it returns false
384    * or until the end of the IOBuf chain is reached.
385    *
386    * Returns the result as a string.
387    */
388   template <typename Predicate>
389   std::string readWhile(const Predicate& predicate);
390 
391   /*
392    * Read all bytes until the specified predicate returns true.
393    *
394    * This is a more generic version of readWhile() takes an arbitrary Output
395    * object, and calls Output::append() with each chunk of matching data.
396    */
397   template <typename Predicate, typename Output>
398   void readWhile(const Predicate& predicate, Output& out);
399 
400   /*
401    * Skip all bytes until the specified predicate returns true.
402    *
403    * The predicate will be called on each byte in turn, until it returns false
404    * or until the end of the IOBuf chain is reached.
405    */
406   template <typename Predicate>
407   void skipWhile(const Predicate& predicate);
408 
skipAtMost(size_t len)409   size_t skipAtMost(size_t len) {
410     dcheckIntegrity();
411     if (LIKELY(crtPos_ + len < crtEnd_)) {
412       crtPos_ += len;
413       return len;
414     }
415     return skipAtMostSlow(len);
416   }
417 
skip(size_t len)418   void skip(size_t len) {
419     dcheckIntegrity();
420     if (LIKELY(crtPos_ + len < crtEnd_)) {
421       crtPos_ += len;
422     } else {
423       skipSlow(len);
424     }
425   }
426 
427   /**
428    * Skip bytes in the current IOBuf without advancing to the next one.
429    * Precondition: length() >= len
430    */
skipNoAdvance(size_t len)431   void skipNoAdvance(size_t len) {
432     DCHECK_LE(len, length());
433     crtPos_ += len;
434   }
435 
retreatAtMost(size_t len)436   size_t retreatAtMost(size_t len) {
437     dcheckIntegrity();
438     if (len <= static_cast<size_t>(crtPos_ - crtBegin_)) {
439       crtPos_ -= len;
440       return len;
441     }
442     return retreatAtMostSlow(len);
443   }
444 
retreat(size_t len)445   void retreat(size_t len) {
446     dcheckIntegrity();
447     if (len <= static_cast<size_t>(crtPos_ - crtBegin_)) {
448       crtPos_ -= len;
449     } else {
450       retreatSlow(len);
451     }
452   }
453 
pullAtMost(void * buf,size_t len)454   size_t pullAtMost(void* buf, size_t len) {
455     dcheckIntegrity();
456     // Fast path: it all fits in one buffer.
457     if (LIKELY(crtPos_ + len <= crtEnd_)) {
458       memcpy(buf, data(), len);
459       crtPos_ += len;
460       return len;
461     }
462     return pullAtMostSlow(buf, len);
463   }
464 
pull(void * buf,size_t len)465   void pull(void* buf, size_t len) {
466     if (UNLIKELY(len == 0)) {
467       return;
468     }
469     dcheckIntegrity();
470     if (LIKELY(crtPos_ + len <= crtEnd_)) {
471       memcpy(buf, data(), len);
472       crtPos_ += len;
473     } else {
474       pullSlow(buf, len);
475     }
476   }
477 
478   /**
479    * Return the available data in the current buffer.
480    * If you want to gather more data from the chain into a contiguous region
481    * (for hopefully zero-copy access), use gather() before peekBytes().
482    */
peekBytes()483   ByteRange peekBytes() {
484     // Ensure that we're pointing to valid data
485     size_t available = length();
486     while (UNLIKELY(available == 0 && tryAdvanceBuffer())) {
487       available = length();
488     }
489     return ByteRange{data(), available};
490   }
491 
492   /**
493    * Alternate version of peekBytes() that returns a std::pair
494    * instead of a ByteRange.  (This method pre-dates ByteRange.)
495    *
496    * This function will eventually be deprecated.
497    */
peek()498   std::pair<const uint8_t*, size_t> peek() {
499     auto bytes = peekBytes();
500     return std::make_pair(bytes.data(), bytes.size());
501   }
502 
clone(std::unique_ptr<folly::IOBuf> & buf,size_t len)503   void clone(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
504     if (UNLIKELY(cloneAtMost(buf, len) != len)) {
505       throw_exception<std::out_of_range>("underflow");
506     }
507   }
508 
clone(folly::IOBuf & buf,size_t len)509   void clone(folly::IOBuf& buf, size_t len) {
510     if (UNLIKELY(cloneAtMost(buf, len) != len)) {
511       throw_exception<std::out_of_range>("underflow");
512     }
513   }
514 
cloneAtMost(folly::IOBuf & buf,size_t len)515   size_t cloneAtMost(folly::IOBuf& buf, size_t len) {
516     // We might be at the end of buffer.
517     advanceBufferIfEmpty();
518 
519     std::unique_ptr<folly::IOBuf> tmp;
520     size_t copied = 0;
521     for (int loopCount = 0; true; ++loopCount) {
522       // Fast path: it all fits in one buffer.
523       size_t available = length();
524       if (LIKELY(available >= len)) {
525         if (loopCount == 0) {
526           crtBuf_->cloneOneInto(buf);
527           buf.trimStart(crtPos_ - crtBegin_);
528           buf.trimEnd(buf.length() - len);
529         } else {
530           tmp = crtBuf_->cloneOne();
531           tmp->trimStart(crtPos_ - crtBegin_);
532           tmp->trimEnd(tmp->length() - len);
533           buf.prependChain(std::move(tmp));
534         }
535 
536         crtPos_ += len;
537         advanceBufferIfEmpty();
538         return copied + len;
539       }
540 
541       if (loopCount == 0) {
542         crtBuf_->cloneOneInto(buf);
543         buf.trimStart(crtPos_ - crtBegin_);
544       } else {
545         tmp = crtBuf_->cloneOne();
546         tmp->trimStart(crtPos_ - crtBegin_);
547         buf.prependChain(std::move(tmp));
548       }
549 
550       copied += available;
551       if (UNLIKELY(!tryAdvanceBuffer())) {
552         return copied;
553       }
554       len -= available;
555     }
556   }
557 
cloneAtMost(std::unique_ptr<folly::IOBuf> & buf,size_t len)558   size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
559     if (!buf) {
560       buf = std::make_unique<folly::IOBuf>();
561     }
562     return cloneAtMost(*buf, len);
563   }
564 
565   /**
566    * Return the distance between two cursors.
567    */
568   size_t operator-(const CursorBase& other) const {
569     BufType* otherBuf = other.crtBuf_;
570     size_t len = 0;
571 
572     if (otherBuf != crtBuf_) {
573       if (other.remainingLen_ == 0) {
574         len += otherBuf->tail() - other.crtPos_;
575       } else {
576         len += other.crtEnd_ - other.crtPos_;
577       }
578 
579       for (otherBuf = otherBuf->next();
580            otherBuf != crtBuf_ && otherBuf != other.buffer_;
581            otherBuf = otherBuf->next()) {
582         len += otherBuf->length();
583       }
584 
585       if (otherBuf == other.buffer_) {
586         throw_exception<std::out_of_range>("wrap-around");
587       }
588 
589       len += crtPos_ - crtBegin_;
590     } else {
591       if (crtPos_ < other.crtPos_) {
592         throw_exception<std::out_of_range>("underflow");
593       }
594 
595       len += crtPos_ - other.crtPos_;
596     }
597 
598     return len;
599   }
600 
601   /**
602    * Return the distance from the given IOBuf to the this cursor.
603    */
604   size_t operator-(const BufType* buf) const {
605     size_t len = 0;
606 
607     const BufType* curBuf = buf;
608     while (curBuf != crtBuf_) {
609       len += curBuf->length();
610       curBuf = curBuf->next();
611       if (curBuf == buf || curBuf == buffer_) {
612         throw_exception<std::out_of_range>("wrap-around");
613       }
614     }
615 
616     len += crtPos_ - crtBegin_;
617     return len;
618   }
619 
isBounded()620   bool isBounded() const {
621     return remainingLen_ != std::numeric_limits<size_t>::max();
622   }
623 
624  protected:
dcheckIntegrity()625   void dcheckIntegrity() const {
626     DCHECK(crtBegin_ <= crtPos_ && crtPos_ <= crtEnd_);
627     DCHECK(crtBuf_ == nullptr || crtBegin_ == crtBuf_->data());
628     DCHECK(
629         crtBuf_ == nullptr ||
630         (std::size_t)(crtEnd_ - crtBegin_) <= crtBuf_->length());
631   }
632 
~CursorBase()633   ~CursorBase() {}
634 
head()635   BufType* head() { return buffer_; }
636 
tryAdvanceBuffer()637   bool tryAdvanceBuffer() {
638     BufType* nextBuf = crtBuf_->next();
639     if (UNLIKELY(nextBuf == buffer_) || remainingLen_ == 0) {
640       crtPos_ = crtEnd_;
641       return false;
642     }
643 
644     absolutePos_ += crtEnd_ - crtBegin_;
645     crtBuf_ = nextBuf;
646     crtPos_ = crtBegin_ = crtBuf_->data();
647     crtEnd_ = crtBuf_->tail();
648     if (isBounded()) {
649       if (crtPos_ + remainingLen_ < crtEnd_) {
650         crtEnd_ = crtPos_ + remainingLen_;
651       }
652       remainingLen_ -= crtEnd_ - crtPos_;
653     }
654     derived().advanceDone();
655     return true;
656   }
657 
tryRetreatBuffer()658   bool tryRetreatBuffer() {
659     if (UNLIKELY(crtBuf_ == buffer_)) {
660       crtPos_ = crtBegin_;
661       return false;
662     }
663     if (isBounded()) {
664       remainingLen_ += crtEnd_ - crtBegin_;
665     }
666     crtBuf_ = crtBuf_->prev();
667     crtBegin_ = crtBuf_->data();
668     crtPos_ = crtEnd_ = crtBuf_->tail();
669     absolutePos_ -= crtEnd_ - crtBegin_;
670     derived().advanceDone();
671     return true;
672   }
673 
advanceBufferIfEmpty()674   void advanceBufferIfEmpty() {
675     dcheckIntegrity();
676     if (crtPos_ == crtEnd_) {
677       tryAdvanceBuffer();
678     }
679   }
680 
681   BufType* crtBuf_;
682   BufType* buffer_;
683   const uint8_t* crtBegin_{nullptr};
684   const uint8_t* crtEnd_{nullptr};
685   const uint8_t* crtPos_{nullptr};
686   size_t absolutePos_{0};
687   // For bounded Cursor, remainingLen_ is the remaining number of data bytes
688   // in subsequent IOBufs in the chain. For unbounded Cursor, remainingLen_
689   // is set to the max of size_t
690   size_t remainingLen_{std::numeric_limits<size_t>::max()};
691 
692  private:
derived()693   Derived& derived() { return static_cast<Derived&>(*this); }
694 
derived()695   Derived const& derived() const { return static_cast<const Derived&>(*this); }
696 
697   template <class T>
readSlow()698   FOLLY_NOINLINE T readSlow() {
699     T val;
700     pullSlow(&val, sizeof(T));
701     return val;
702   }
703 
readFixedStringSlow(std::string * str,size_t len)704   FOLLY_NOINLINE void readFixedStringSlow(std::string* str, size_t len) {
705     for (size_t available; (available = length()) < len;) {
706       str->append(reinterpret_cast<const char*>(data()), available);
707       if (UNLIKELY(!tryAdvanceBuffer())) {
708         throw_exception<std::out_of_range>("string underflow");
709       }
710       len -= available;
711     }
712     str->append(reinterpret_cast<const char*>(data()), len);
713     crtPos_ += len;
714     advanceBufferIfEmpty();
715   }
716 
pullAtMostSlow(void * buf,size_t len)717   FOLLY_NOINLINE size_t pullAtMostSlow(void* buf, size_t len) {
718     // If the length of this buffer is 0 try advancing it.
719     // Otherwise on the first iteration of the following loop memcpy is called
720     // with a null source pointer.
721     if (UNLIKELY(length() == 0 && !tryAdvanceBuffer())) {
722       return 0;
723     }
724     uint8_t* p = reinterpret_cast<uint8_t*>(buf);
725     size_t copied = 0;
726     for (size_t available; (available = length()) < len;) {
727       memcpy(p, data(), available);
728       copied += available;
729       if (UNLIKELY(!tryAdvanceBuffer())) {
730         return copied;
731       }
732       p += available;
733       len -= available;
734     }
735     memcpy(p, data(), len);
736     crtPos_ += len;
737     advanceBufferIfEmpty();
738     return copied + len;
739   }
740 
pullSlow(void * buf,size_t len)741   FOLLY_NOINLINE void pullSlow(void* buf, size_t len) {
742     if (UNLIKELY(pullAtMostSlow(buf, len) != len)) {
743       throw_exception<std::out_of_range>("underflow");
744     }
745   }
746 
skipAtMostSlow(size_t len)747   FOLLY_NOINLINE size_t skipAtMostSlow(size_t len) {
748     size_t skipped = 0;
749     for (size_t available; (available = length()) < len;) {
750       skipped += available;
751       if (UNLIKELY(!tryAdvanceBuffer())) {
752         return skipped;
753       }
754       len -= available;
755     }
756     crtPos_ += len;
757     advanceBufferIfEmpty();
758     return skipped + len;
759   }
760 
skipSlow(size_t len)761   FOLLY_NOINLINE void skipSlow(size_t len) {
762     if (UNLIKELY(skipAtMostSlow(len) != len)) {
763       throw_exception<std::out_of_range>("underflow");
764     }
765   }
766 
retreatAtMostSlow(size_t len)767   FOLLY_NOINLINE size_t retreatAtMostSlow(size_t len) {
768     size_t retreated = 0;
769     for (size_t available; (available = crtPos_ - crtBegin_) < len;) {
770       retreated += available;
771       if (UNLIKELY(!tryRetreatBuffer())) {
772         return retreated;
773       }
774       len -= available;
775     }
776     crtPos_ -= len;
777     return retreated + len;
778   }
779 
retreatSlow(size_t len)780   FOLLY_NOINLINE void retreatSlow(size_t len) {
781     if (UNLIKELY(retreatAtMostSlow(len) != len)) {
782       throw_exception<std::out_of_range>("underflow");
783     }
784   }
785 
advanceDone()786   void advanceDone() {}
787 };
788 
789 } // namespace detail
790 
791 class Cursor : public detail::CursorBase<Cursor, const IOBuf> {
792  public:
Cursor(const IOBuf * buf)793   explicit Cursor(const IOBuf* buf)
794       : detail::CursorBase<Cursor, const IOBuf>(buf) {}
795 
Cursor(const IOBuf * buf,size_t len)796   explicit Cursor(const IOBuf* buf, size_t len)
797       : detail::CursorBase<Cursor, const IOBuf>(buf, len) {}
798 
799   template <class OtherDerived, class OtherBuf>
Cursor(const detail::CursorBase<OtherDerived,OtherBuf> & cursor)800   explicit Cursor(const detail::CursorBase<OtherDerived, OtherBuf>& cursor)
801       : detail::CursorBase<Cursor, const IOBuf>(cursor) {}
802 
803   template <class OtherDerived, class OtherBuf>
Cursor(const detail::CursorBase<OtherDerived,OtherBuf> & cursor,size_t len)804   Cursor(const detail::CursorBase<OtherDerived, OtherBuf>& cursor, size_t len)
805       : detail::CursorBase<Cursor, const IOBuf>(cursor, len) {}
806 };
807 
808 namespace detail {
809 
810 template <class Derived>
811 class Writable {
812  public:
813   template <class T>
814   typename std::enable_if<std::is_arithmetic<T>::value>::type write(
815       T value, size_t n = sizeof(T)) {
816     assert(n <= sizeof(T));
817     const uint8_t* u8 = reinterpret_cast<const uint8_t*>(&value);
818     Derived* d = static_cast<Derived*>(this);
819     d->push(u8, n);
820   }
821 
822   template <class T>
writeBE(T value)823   void writeBE(T value) {
824     Derived* d = static_cast<Derived*>(this);
825     d->write(Endian::big(value));
826   }
827 
828   template <class T>
writeLE(T value)829   void writeLE(T value) {
830     Derived* d = static_cast<Derived*>(this);
831     d->write(Endian::little(value));
832   }
833 
push(const uint8_t * buf,size_t len)834   void push(const uint8_t* buf, size_t len) {
835     Derived* d = static_cast<Derived*>(this);
836     if (d->pushAtMost(buf, len) != len) {
837       throw_exception<std::out_of_range>("overflow");
838     }
839   }
840 
push(ByteRange buf)841   void push(ByteRange buf) {
842     if (this->pushAtMost(buf) != buf.size()) {
843       throw_exception<std::out_of_range>("overflow");
844     }
845   }
846 
pushAtMost(ByteRange buf)847   size_t pushAtMost(ByteRange buf) {
848     Derived* d = static_cast<Derived*>(this);
849     return d->pushAtMost(buf.data(), buf.size());
850   }
851 
852   /**
853    * push len bytes of data from input cursor, data could be in an IOBuf chain.
854    * If input cursor contains less than len bytes, or this cursor has less than
855    * len bytes writable space, an out_of_range exception will be thrown.
856    */
push(Cursor cursor,size_t len)857   void push(Cursor cursor, size_t len) {
858     if (this->pushAtMost(cursor, len) != len) {
859       throw_exception<std::out_of_range>("overflow");
860     }
861   }
862 
pushAtMost(Cursor cursor,size_t len)863   size_t pushAtMost(Cursor cursor, size_t len) {
864     size_t written = 0;
865     for (;;) {
866       auto currentBuffer = cursor.peekBytes();
867       const uint8_t* crtData = currentBuffer.data();
868       size_t available = currentBuffer.size();
869       if (available == 0) {
870         // end of buffer chain
871         return written;
872       }
873       // all data is in current buffer
874       if (available >= len) {
875         this->push(crtData, len);
876         cursor.skip(len);
877         return written + len;
878       }
879 
880       // write the whole current IOBuf
881       this->push(crtData, available);
882       cursor.skip(available);
883       written += available;
884       len -= available;
885     }
886   }
887 };
888 
889 } // namespace detail
890 
891 enum class CursorAccess { PRIVATE, UNSHARE };
892 
893 template <CursorAccess access>
894 class RWCursor : public detail::CursorBase<RWCursor<access>, IOBuf>,
895                  public detail::Writable<RWCursor<access>> {
896   friend class detail::CursorBase<RWCursor<access>, IOBuf>;
897 
898  public:
RWCursor(IOBuf * buf)899   explicit RWCursor(IOBuf* buf)
900       : detail::CursorBase<RWCursor<access>, IOBuf>(buf), maybeShared_(true) {}
901 
RWCursor(IOBufQueue & queue)902   explicit RWCursor(IOBufQueue& queue)
903       : RWCursor((queue.flushCache(), queue.head_.get())) {}
904 
905   template <class OtherDerived, class OtherBuf>
RWCursor(const detail::CursorBase<OtherDerived,OtherBuf> & cursor)906   explicit RWCursor(const detail::CursorBase<OtherDerived, OtherBuf>& cursor)
907       : detail::CursorBase<RWCursor<access>, IOBuf>(cursor),
908         maybeShared_(true) {
909     CHECK(!cursor.isBounded())
910         << "Creating RWCursor from bounded Cursor is not allowed";
911   }
912   /**
913    * Gather at least n bytes contiguously into the current buffer,
914    * by coalescing subsequent buffers from the chain as necessary.
915    */
gather(size_t n)916   void gather(size_t n) {
917     // Forbid attempts to gather beyond the end of this IOBuf chain.
918     // Otherwise we could try to coalesce the head of the chain and end up
919     // accidentally freeing it, invalidating the pointer owned by external
920     // code.
921     //
922     // If crtBuf_ == head() then IOBuf::gather() will perform all necessary
923     // checking.  We only have to perform an explicit check here when calling
924     // gather() on a non-head element.
925     if (this->crtBuf_ != this->head() && this->totalLength() < n) {
926       throw std::overflow_error("cannot gather() past the end of the chain");
927     }
928     size_t offset = this->crtPos_ - this->crtBegin_;
929     this->crtBuf_->gather(offset + n);
930     this->crtBegin_ = this->crtBuf_->data();
931     this->crtEnd_ = this->crtBuf_->tail();
932     this->crtPos_ = this->crtBegin_ + offset;
933   }
gatherAtMost(size_t n)934   void gatherAtMost(size_t n) {
935     this->dcheckIntegrity();
936     size_t size = std::min(n, this->totalLength());
937     size_t offset = this->crtPos_ - this->crtBegin_;
938     this->crtBuf_->gather(offset + size);
939     this->crtBegin_ = this->crtBuf_->data();
940     this->crtEnd_ = this->crtBuf_->tail();
941     this->crtPos_ = this->crtBegin_ + offset;
942   }
943 
944   using detail::Writable<RWCursor<access>>::pushAtMost;
pushAtMost(const uint8_t * buf,size_t len)945   size_t pushAtMost(const uint8_t* buf, size_t len) {
946     // We have to explicitly check for an input length of 0.
947     // We support buf being nullptr in this case, but we need to avoid calling
948     // memcpy() with a null source pointer, since that is undefined behavior
949     // even if the length is 0.
950     if (len == 0) {
951       return 0;
952     }
953 
954     size_t copied = 0;
955     for (;;) {
956       // Fast path: the current buffer is big enough.
957       size_t available = this->length();
958       if (LIKELY(available >= len)) {
959         if (access == CursorAccess::UNSHARE) {
960           maybeUnshare();
961         }
962         memcpy(writableData(), buf, len);
963         this->crtPos_ += len;
964         return copied + len;
965       }
966 
967       if (access == CursorAccess::UNSHARE) {
968         maybeUnshare();
969       }
970       memcpy(writableData(), buf, available);
971       copied += available;
972       if (UNLIKELY(!this->tryAdvanceBuffer())) {
973         return copied;
974       }
975       buf += available;
976       len -= available;
977     }
978   }
979 
insert(std::unique_ptr<folly::IOBuf> buf)980   void insert(std::unique_ptr<folly::IOBuf> buf) {
981     this->dcheckIntegrity();
982     this->absolutePos_ += buf->computeChainDataLength();
983     if (this->crtPos_ == this->crtBegin_ && this->crtBuf_ != this->buffer_) {
984       // Can just prepend
985       this->crtBuf_->prependChain(std::move(buf));
986     } else {
987       IOBuf* nextBuf;
988       std::unique_ptr<folly::IOBuf> remaining;
989       if (this->crtPos_ != this->crtEnd_) {
990         // Need to split current IOBuf in two.
991         remaining = this->crtBuf_->cloneOne();
992         remaining->trimStart(this->crtPos_ - this->crtBegin_);
993         nextBuf = remaining.get();
994         buf->prependChain(std::move(remaining));
995       } else {
996         // Can just append
997         nextBuf = this->crtBuf_->next();
998       }
999       this->crtBuf_->trimEnd(this->length());
1000       this->absolutePos_ += this->crtPos_ - this->crtBegin_;
1001       this->crtBuf_->insertAfterThisOne(std::move(buf));
1002 
1003       if (nextBuf == this->buffer_) {
1004         // We've just appended to the end of the buffer, so advance to the end.
1005         this->crtBuf_ = this->buffer_->prev();
1006         this->crtBegin_ = this->crtBuf_->data();
1007         this->crtPos_ = this->crtEnd_ = this->crtBuf_->tail();
1008         // This has already been accounted for, so remove it.
1009         this->absolutePos_ -= this->crtEnd_ - this->crtBegin_;
1010       } else {
1011         // Jump past the new links
1012         this->crtBuf_ = nextBuf;
1013         this->crtPos_ = this->crtBegin_ = this->crtBuf_->data();
1014         this->crtEnd_ = this->crtBuf_->tail();
1015       }
1016     }
1017   }
1018 
writableData()1019   uint8_t* writableData() {
1020     this->dcheckIntegrity();
1021     return this->crtBuf_->writableData() + (this->crtPos_ - this->crtBegin_);
1022   }
1023 
1024  private:
maybeUnshare()1025   void maybeUnshare() {
1026     if (UNLIKELY(maybeShared_)) {
1027       size_t offset = this->crtPos_ - this->crtBegin_;
1028       this->crtBuf_->unshareOne();
1029       this->crtBegin_ = this->crtBuf_->data();
1030       this->crtEnd_ = this->crtBuf_->tail();
1031       this->crtPos_ = this->crtBegin_ + offset;
1032       maybeShared_ = false;
1033     }
1034   }
1035 
advanceDone()1036   void advanceDone() { maybeShared_ = true; }
1037 
1038   bool maybeShared_;
1039 };
1040 
1041 typedef RWCursor<CursorAccess::PRIVATE> RWPrivateCursor;
1042 typedef RWCursor<CursorAccess::UNSHARE> RWUnshareCursor;
1043 
1044 /**
1045  * Append to the end of a buffer chain, growing the chain (by allocating new
1046  * buffers) in increments of at least growth bytes every time.  Won't grow
1047  * (and push() and ensure() will throw) if growth == 0.
1048  *
1049  * TODO(tudorb): add a flavor of Appender that reallocates one IOBuf instead
1050  * of chaining.
1051  */
1052 class Appender : public detail::Writable<Appender> {
1053  public:
Appender(IOBuf * buf,std::size_t growth)1054   Appender(IOBuf* buf, std::size_t growth)
1055       : buffer_(buf), crtBuf_(buf->prev()), growth_(growth) {}
1056 
writableData()1057   uint8_t* writableData() { return crtBuf_->writableTail(); }
1058 
length()1059   size_t length() const { return crtBuf_->tailroom(); }
1060 
1061   /**
1062    * Mark n bytes (must be <= length()) as appended, as per the
1063    * IOBuf::append() method.
1064    */
append(size_t n)1065   void append(size_t n) { crtBuf_->append(n); }
1066 
1067   /**
1068    * Ensure at least n contiguous bytes available to write.
1069    * Postcondition: length() >= n.
1070    */
ensure(std::size_t n)1071   void ensure(std::size_t n) {
1072     if (LIKELY(length() >= n)) {
1073       return;
1074     }
1075 
1076     // Waste the rest of the current buffer and allocate a new one.
1077     // Don't make it too small, either.
1078     if (growth_ == 0) {
1079       throw_exception<std::out_of_range>("can't grow buffer chain");
1080     }
1081 
1082     n = std::max(n, growth_);
1083     buffer_->prependChain(IOBuf::create(n));
1084     crtBuf_ = buffer_->prev();
1085   }
1086 
1087   using detail::Writable<Appender>::pushAtMost;
pushAtMost(const uint8_t * buf,size_t len)1088   size_t pushAtMost(const uint8_t* buf, size_t len) {
1089     // We have to explicitly check for an input length of 0.
1090     // We support buf being nullptr in this case, but we need to avoid calling
1091     // memcpy() with a null source pointer, since that is undefined behavior
1092     // even if the length is 0.
1093     if (len == 0) {
1094       return 0;
1095     }
1096 
1097     // If the length of this buffer is 0 try growing it.
1098     // Otherwise on the first iteration of the following loop memcpy is called
1099     // with a null source pointer.
1100     if (UNLIKELY(length() == 0 && !tryGrowChain())) {
1101       return 0;
1102     }
1103 
1104     size_t copied = 0;
1105     for (;;) {
1106       // Fast path: it all fits in one buffer.
1107       size_t available = length();
1108       if (LIKELY(available >= len)) {
1109         memcpy(writableData(), buf, len);
1110         append(len);
1111         return copied + len;
1112       }
1113 
1114       memcpy(writableData(), buf, available);
1115       append(available);
1116       copied += available;
1117       if (UNLIKELY(!tryGrowChain())) {
1118         return copied;
1119       }
1120       buf += available;
1121       len -= available;
1122     }
1123   }
1124 
1125   /*
1126    * Append to the end of this buffer, using a printf() style
1127    * format specifier.
1128    *
1129    * Note that folly/Format.h provides nicer and more type-safe mechanisms
1130    * for formatting strings, which should generally be preferred over
1131    * printf-style formatting.  Appender objects can be used directly as an
1132    * output argument for Formatter objects.  For example:
1133    *
1134    *   Appender app(&iobuf);
1135    *   format("{} {}", "hello", "world")(app);
1136    *
1137    * However, printf-style strings are still needed when dealing with existing
1138    * third-party code in some cases.
1139    *
1140    * This will always add a nul-terminating character after the end
1141    * of the output.  However, the buffer data length will only be updated to
1142    * include the data itself.  The nul terminator will be the first byte in the
1143    * buffer tailroom.
1144    *
1145    * This method may throw exceptions on error.
1146    */
1147   void printf(FOLLY_PRINTF_FORMAT const char* fmt, ...)
1148       FOLLY_PRINTF_FORMAT_ATTR(2, 3);
1149 
1150   void vprintf(const char* fmt, va_list ap);
1151 
1152   /*
1153    * Calling an Appender object with a StringPiece will append the string
1154    * piece.  This allows Appender objects to be used directly with
1155    * Formatter.
1156    */
operator()1157   void operator()(StringPiece sp) { push(ByteRange(sp)); }
1158 
1159  private:
tryGrowChain()1160   bool tryGrowChain() {
1161     assert(crtBuf_->next() == buffer_);
1162     if (growth_ == 0) {
1163       return false;
1164     }
1165 
1166     buffer_->prependChain(IOBuf::create(growth_));
1167     crtBuf_ = buffer_->prev();
1168     return true;
1169   }
1170 
1171   IOBuf* buffer_;
1172   IOBuf* crtBuf_;
1173   std::size_t growth_;
1174 };
1175 
1176 class QueueAppender : public detail::Writable<QueueAppender> {
1177  public:
1178   /**
1179    * Create an Appender that writes to a IOBufQueue.  When we allocate
1180    * space in the queue, we grow no more than growth bytes at once
1181    * (unless you call ensure() with a bigger value yourself).
1182    */
QueueAppender(IOBufQueue * queue,std::size_t growth)1183   QueueAppender(IOBufQueue* queue, std::size_t growth)
1184       : queueCache_(queue), growth_(growth) {}
1185 
reset(IOBufQueue * queue,std::size_t growth)1186   void reset(IOBufQueue* queue, std::size_t growth) {
1187     queueCache_.reset(queue);
1188     growth_ = growth;
1189   }
1190 
writableData()1191   uint8_t* writableData() { return queueCache_.writableData(); }
1192 
length()1193   size_t length() { return queueCache_.length(); }
1194 
append(size_t n)1195   void append(size_t n) { queueCache_.append(n); }
1196 
1197   // Ensure at least n contiguous; can go above growth_, throws if
1198   // not enough room.
ensure(size_t n)1199   void ensure(size_t n) {
1200     if (length() < n) {
1201       ensureSlow(n);
1202     }
1203   }
1204 
1205   template <class T>
1206   typename std::enable_if<std::is_arithmetic<T>::value>::type write(
1207       T value, size_t n = sizeof(T)) {
1208     // We can't fail.
1209     assert(n <= sizeof(T));
1210     if (length() >= sizeof(T)) {
1211       storeUnaligned(queueCache_.writableData(), value);
1212       queueCache_.appendUnsafe(n);
1213     } else {
1214       writeSlow<T>(value, n);
1215     }
1216   }
1217 
1218   using detail::Writable<QueueAppender>::pushAtMost;
pushAtMost(const uint8_t * buf,size_t len)1219   size_t pushAtMost(const uint8_t* buf, size_t len) {
1220     // Fill the current buffer
1221     const size_t copyLength = std::min(len, length());
1222     if (copyLength != 0) {
1223       memcpy(writableData(), buf, copyLength);
1224       queueCache_.appendUnsafe(copyLength);
1225       buf += copyLength;
1226     }
1227     size_t remaining = len - copyLength;
1228     // Allocate more buffers as necessary
1229     while (remaining != 0) {
1230       auto p = queueCache_.queue()->preallocate(
1231           std::min(remaining, growth_), growth_, remaining);
1232       memcpy(p.first, buf, p.second);
1233       queueCache_.queue()->postallocate(p.second);
1234       buf += p.second;
1235       remaining -= p.second;
1236     }
1237     return len;
1238   }
1239 
insert(std::unique_ptr<folly::IOBuf> buf)1240   void insert(std::unique_ptr<folly::IOBuf> buf) {
1241     if (buf) {
1242       queueCache_.queue()->append(
1243           std::move(buf), /* pack */ true, /* allowTailReuse */ true);
1244     }
1245   }
1246 
insert(const folly::IOBuf & buf)1247   void insert(const folly::IOBuf& buf) {
1248     queueCache_.queue()->append(
1249         buf, /* pack */ true, /* allowTailReuse */ true);
1250   }
1251 
1252   template <CursorAccess access>
1253   explicit operator RWCursor<access>() {
1254     return RWCursor<access>(*queueCache_.queue());
1255   }
1256 
1257  private:
1258   folly::IOBufQueue::WritableRangeCache queueCache_{nullptr};
1259   size_t growth_{0};
1260 
ensureSlow(size_t n)1261   FOLLY_NOINLINE void ensureSlow(size_t n) {
1262     queueCache_.queue()->preallocate(n, growth_);
1263     queueCache_.fillCache();
1264   }
1265 
1266   template <class T>
1267   typename std::enable_if<std::is_arithmetic<T>::value>::type FOLLY_NOINLINE
1268   writeSlow(T value, size_t n = sizeof(T)) {
1269     assert(n <= sizeof(T));
1270     queueCache_.queue()->preallocate(sizeof(T), growth_);
1271     queueCache_.fillCache();
1272 
1273     storeUnaligned(queueCache_.writableData(), value);
1274     queueCache_.appendUnsafe(n);
1275   }
1276 };
1277 
1278 } // namespace io
1279 } // namespace folly
1280 
1281 #include <folly/io/Cursor-inl.h>
1282