Add read*String() methods to Cursor

Summary:
- Add some convenience methods for reading std::string objects via
folly::Cursor

Test Plan: - Unit tests

Reviewed By: simpkins@fb.com

FB internal diff: D795428

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index 73a5493a7a612526b3e7b22d695f927d87e8a77a..790b9cca4a9e22994f38462dfd3af7a95cd7d10c 100644 (file)
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -84,6 +84,76 @@ class CursorBase {
     return Endian::little(read<T>());
   }
 
+  /**
+   * Read a fixed-length string.
+   *
+   * The std::string-based APIs should probably be avoided unless you
+   * ultimately want the data to live in an std::string. You're better off
+   * using the pull() APIs to copy into a raw buffer otherwise.
+   */
+  std::string readFixedString(size_t len) {
+    std::string str;
+
+    str.reserve(len);
+    for (;;) {
+      // Fast path: it all fits in one buffer.
+      size_t available = length();
+      if (LIKELY(available >= len)) {
+        str.append(reinterpret_cast<const char*>(data()), len);
+        offset_ += len;
+        return str;
+      }
+
+      str.append(reinterpret_cast<const char*>(data()), available);
+      if (UNLIKELY(!tryAdvanceBuffer())) {
+        throw std::out_of_range("string underflow");
+      }
+      len -= available;
+    }
+  }
+
+  /**
+   * Read a string consisting of bytes until the given terminator character is
+   * seen. Raises an std::length_error if maxLength bytes have been processed
+   * before the terminator is seen.
+   *
+   * See comments in readFixedString() about when it's appropriate to use this
+   * vs. using pull().
+   */
+  std::string readTerminatedString(
+    char termChar = '\0',
+    size_t maxLength = std::numeric_limits<size_t>::max()) {
+    std::string str;
+
+    for (;;) {
+      const uint8_t* buf = data();
+      size_t buflen = length();
+
+      size_t i = 0;
+      while (i < buflen && buf[i] != termChar) {
+        ++i;
+
+        // Do this check after incrementing 'i', as even though we start at the
+        // 0 byte, it still represents a single character
+        if (str.length() + i >= maxLength) {
+          throw std::length_error("string overflow");
+        }
+      }
+
+      str.append(reinterpret_cast<const char*>(buf), i);
+      if (i < buflen) {
+        skip(i + 1);
+        return str;
+      }
+
+      skip(i);
+
+      if (UNLIKELY(!tryAdvanceBuffer())) {
+        throw std::out_of_range("string underflow");
+      }
+    }
+  }
+
   explicit CursorBase(BufType* buf)
     : crtBuf_(buf)
     , offset_(0)

diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index 3ce32ad6869a3ad2976715143d3d2da1f5f0385f..4c7d622cdf36e6ada7c07e05c4dd95adeaf3b241 100644 (file)
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -383,6 +383,93 @@ TEST(IOBuf, CursorOperators) {
   }
 }
 
+TEST(IOBuf, StringOperations) {
+  // Test a single buffer with two null-terminated strings and an extra uint8_t
+  // at the end
+  {
+    std::unique_ptr<IOBuf> chain(IOBuf::create(16));
+    Appender app(chain.get(), 0);
+    app.pushAtMost(reinterpret_cast<const uint8_t*>("hello\0world\0\x01"), 13);
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello", curs.readTerminatedString().c_str());
+    EXPECT_STREQ("world", curs.readTerminatedString().c_str());
+    EXPECT_EQ(1, curs.read<uint8_t>());
+  }
+
+  // Test multiple buffers with a single null-terminated string spanning them
+  {
+    std::unique_ptr<IOBuf> chain(IOBuf::create(8));
+    chain->prependChain(IOBuf::create(8));
+    Appender app(chain.get(), 0);
+    app.pushAtMost(reinterpret_cast<const uint8_t*>("hello world\0"), 12);
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello world", curs.readTerminatedString().c_str());
+  }
+
+  // Test a reading a null-terminated string that's longer than the maximum
+  // allowable length
+  {
+    std::unique_ptr<IOBuf> chain(IOBuf::create(16));
+    Appender app(chain.get(), 0);
+    app.pushAtMost(reinterpret_cast<const uint8_t*>("hello world\0"), 12);
+
+    Cursor curs(chain.get());
+    EXPECT_THROW(curs.readTerminatedString('\0', 5), std::length_error);
+  }
+
+  // Test reading a null-termianted string from a chain with an empty buffer at
+  // the front
+  {
+    std::unique_ptr<IOBuf> buf(IOBuf::create(8));
+    Appender app(buf.get(), 0);
+    app.pushAtMost(reinterpret_cast<const uint8_t*>("hello\0"), 6);
+    std::unique_ptr<IOBuf> chain(IOBuf::create(8));
+    chain->prependChain(std::move(buf));
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello", curs.readTerminatedString().c_str());
+  }
+
+  // Test reading a two fixed-length strings from a single buffer with an extra
+  // uint8_t at the end
+  {
+    std::unique_ptr<IOBuf> chain(IOBuf::create(16));
+    Appender app(chain.get(), 0);
+    app.pushAtMost(reinterpret_cast<const uint8_t*>("helloworld\x01"), 11);
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello", curs.readFixedString(5).c_str());
+    EXPECT_STREQ("world", curs.readFixedString(5).c_str());
+    EXPECT_EQ(1, curs.read<uint8_t>());
+  }
+
+  // Test multiple buffers with a single fixed-length string spanning them
+  {
+    std::unique_ptr<IOBuf> chain(IOBuf::create(8));
+    chain->prependChain(IOBuf::create(8));
+    Appender app(chain.get(), 0);
+    app.pushAtMost(reinterpret_cast<const uint8_t*>("hello world"), 11);
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello world", curs.readFixedString(11).c_str());
+  }
+
+  // Test reading a fixed-length string from a chain with an empty buffer at
+  // the front
+  {
+    std::unique_ptr<IOBuf> buf(IOBuf::create(8));
+    Appender app(buf.get(), 0);
+    app.pushAtMost(reinterpret_cast<const uint8_t*>("hello"), 5);
+    std::unique_ptr<IOBuf> chain(IOBuf::create(8));
+    chain->prependChain(std::move(buf));
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello", curs.readFixedString(5).c_str());
+  }
+}
+
 int benchmark_size = 1000;
 unique_ptr<IOBuf> iobuf_benchmark;