NetString::ReadStringFromStream(): add Boost ASIO overload

diff --git a/lib/base/netstring.cpp b/lib/base/netstring.cpp
index debed2bda8dfb0e38619afd681d45c09b80a322a..489a8b40db727a00db9dee1652be5b2d6151d598 100644 (file)
--- a/lib/base/netstring.cpp
+++ b/lib/base/netstring.cpp
@@ -3,9 +3,12 @@
 #include "base/netstring.hpp"
 #include "base/debug.hpp"
 #include "base/tlsstream.hpp"
+#include <cstdint>
 #include <memory>
 #include <sstream>
+#include <utility>
 #include <boost/asio/buffer.hpp>
+#include <boost/asio/read.hpp>
 #include <boost/asio/spawn.hpp>
 #include <boost/asio/write.hpp>
 
@@ -115,6 +118,85 @@ size_t NetString::WriteStringToStream(const Stream::Ptr& stream, const String& s
        return msg.GetLength();
 }
 
+/**
+ * Reads data from a stream in netstring format.
+ *
+ * @param stream The stream to read from.
+ * @returns The String that has been read from the IOQueue.
+ * @exception invalid_argument The input stream is invalid.
+ * @see https://github.com/PeterScott/netstring-c/blob/master/netstring.c
+ */
+String NetString::ReadStringFromStream(const std::shared_ptr<AsioTlsStream>& stream,
+       boost::asio::yield_context yc, ssize_t maxMessageLength)
+{
+       namespace asio = boost::asio;
+
+       size_t len = 0;
+       bool leadingZero = false;
+
+       for (uint_fast8_t readBytes = 0;; ++readBytes) {
+               char byte = 0;
+
+               {
+                       asio::mutable_buffer byteBuf (&byte, 1);
+                       asio::async_read(*stream, byteBuf, yc);
+               }
+
+               if (isdigit(byte)) {
+                       if (readBytes == 9) {
+                               BOOST_THROW_EXCEPTION(std::invalid_argument("Length specifier must not exceed 9 characters"));
+                       }
+
+                       if (leadingZero) {
+                               BOOST_THROW_EXCEPTION(std::invalid_argument("Invalid NetString (leading zero)"));
+                       }
+
+                       len = len * 10u + size_t(byte - '0');
+
+                       if (!readBytes && byte == '0') {
+                               leadingZero = true;
+                       }
+               } else if (byte == ':') {
+                       if (!readBytes) {
+                               BOOST_THROW_EXCEPTION(std::invalid_argument("Invalid NetString (no length specifier)"));
+                       }
+
+                       break;
+               } else {
+                       BOOST_THROW_EXCEPTION(std::invalid_argument("Invalid NetString (missing :)"));
+               }
+       }
+
+       if (maxMessageLength >= 0 && len > maxMessageLength) {
+               std::stringstream errorMessage;
+               errorMessage << "Max data length exceeded: " << (maxMessageLength / 1024) << " KB";
+
+               BOOST_THROW_EXCEPTION(std::invalid_argument(errorMessage.str()));
+       }
+
+       String payload;
+
+       if (len) {
+               payload.Append(len, 0);
+
+               asio::mutable_buffer payloadBuf (&*payload.Begin(), payload.GetLength());
+               asio::async_read(*stream, payloadBuf, yc);
+       }
+
+       char trailer = 0;
+
+       {
+               asio::mutable_buffer trailerBuf (&trailer, 1);
+               asio::async_read(*stream, trailerBuf, yc);
+       }
+
+       if (trailer != ',') {
+               BOOST_THROW_EXCEPTION(std::invalid_argument("Invalid NetString (missing ,)"));
+       }
+
+       return std::move(payload);
+}
+
 /**
  * Writes data into a stream using the netstring format and returns bytes written.
  *

diff --git a/lib/base/netstring.hpp b/lib/base/netstring.hpp
index f54d70c1721ac0a08035c9438806d4b948fdba29..f84eac7a3138b6ac316f03a491258b17faf7d0f5 100644 (file)
--- a/lib/base/netstring.hpp
+++ b/lib/base/netstring.hpp
@@ -26,6 +26,8 @@ class NetString
 public:
        static StreamReadStatus ReadStringFromStream(const Stream::Ptr& stream, String *message, StreamReadContext& context,
                bool may_wait = false, ssize_t maxMessageLength = -1);
+       static String ReadStringFromStream(const std::shared_ptr<AsioTlsStream>& stream,
+               boost::asio::yield_context yc, ssize_t maxMessageLength = -1);
        static size_t WriteStringToStream(const Stream::Ptr& stream, const String& message);
        static size_t WriteStringToStream(const std::shared_ptr<AsioTlsStream>& stream, const String& message, boost::asio::yield_context yc);
        static void WriteStringToStream(std::ostream& stream, const String& message);