From b26808414c9cdc5d8d9243771fd8fca4c76128e4 Mon Sep 17 00:00:00 2001
From: "Alexander A. Klimov" <alexander.klimov@icinga.com>
Date: Tue, 19 Feb 2019 11:20:39 +0100
Subject: [PATCH] NetString::ReadStringFromStream(): add Boost ASIO overload

---
 lib/base/netstring.cpp | 82 ++++++++++++++++++++++++++++++++++++++++++
 lib/base/netstring.hpp |  2 ++
 2 files changed, 84 insertions(+)

diff --git a/lib/base/netstring.cpp b/lib/base/netstring.cpp
index debed2bda..489a8b40d 100644
--- 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 f54d70c17..f84eac7a3 100644
--- 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);
-- 
2.40.0