From 54ff60cd8cbfc93d21622c611d32f4efdae0e155 Mon Sep 17 00:00:00 2001
From: Noah Hilverling <noah.hilverling@icinga.com>
Date: Thu, 1 Mar 2018 09:47:29 +0100
Subject: [PATCH] Limit JSON-RPC message size

---
 lib/base/netstring.cpp           | 10 +++++++++-
 lib/base/netstring.hpp           |  3 ++-
 lib/remote/jsonrpc.cpp           |  4 ++--
 lib/remote/jsonrpc.hpp           |  2 +-
 lib/remote/jsonrpcconnection.cpp |  7 ++++++-
 5 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/lib/base/netstring.cpp b/lib/base/netstring.cpp
index 0bd110bac..e970f3b61 100644
--- a/lib/base/netstring.cpp
+++ b/lib/base/netstring.cpp
@@ -32,7 +32,8 @@ using namespace icinga;
  * @exception invalid_argument The input stream is invalid.
  * @see https://github.com/PeterScott/netstring-c/blob/master/netstring.c
  */
-StreamReadStatus NetString::ReadStringFromStream(const Stream::Ptr& stream, String *str, StreamReadContext& context, bool may_wait)
+StreamReadStatus NetString::ReadStringFromStream(const Stream::Ptr& stream, String *str, StreamReadContext& context,
+	bool may_wait, ssize_t maxMessageLength)
 {
 	if (context.Eof)
 		return StatusEof;
@@ -84,6 +85,13 @@ StreamReadStatus NetString::ReadStringFromStream(const Stream::Ptr& stream, Stri
 	/* read the whole message */
 	size_t data_length = len + 1;
 
+	if (maxMessageLength >= 0 && data_length > maxMessageLength) {
+		std::stringstream errorMessage;
+		errorMessage << "Max data length exceeded: " << (maxMessageLength / 1024 / 1024) << " MB";
+
+		BOOST_THROW_EXCEPTION(std::invalid_argument(errorMessage.str()));
+	}
+
 	char *data = context.Buffer + header_length + 1;
 
 	if (context.Size < header_length + 1 + data_length) {
diff --git a/lib/base/netstring.hpp b/lib/base/netstring.hpp
index 4ba02ccd3..1daaa979a 100644
--- a/lib/base/netstring.hpp
+++ b/lib/base/netstring.hpp
@@ -38,7 +38,8 @@ class String;
 class NetString
 {
 public:
-	static StreamReadStatus ReadStringFromStream(const Stream::Ptr& stream, String *message, StreamReadContext& context, bool may_wait = false);
+	static StreamReadStatus ReadStringFromStream(const Stream::Ptr& stream, String *message, StreamReadContext& context,
+		bool may_wait = false, ssize_t maxMessageLength = -1);
 	static size_t WriteStringToStream(const Stream::Ptr& stream, const String& message);
 	static void WriteStringToStream(std::ostream& stream, const String& message);
 
diff --git a/lib/remote/jsonrpc.cpp b/lib/remote/jsonrpc.cpp
index 94bd9cb5b..3dbd13aa3 100644
--- a/lib/remote/jsonrpc.cpp
+++ b/lib/remote/jsonrpc.cpp
@@ -72,10 +72,10 @@ size_t JsonRpc::SendMessage(const Stream::Ptr& stream, const Dictionary::Ptr& me
 	return NetString::WriteStringToStream(stream, json);
 }
 
-StreamReadStatus JsonRpc::ReadMessage(const Stream::Ptr& stream, String *message, StreamReadContext& src, bool may_wait)
+StreamReadStatus JsonRpc::ReadMessage(const Stream::Ptr& stream, String *message, StreamReadContext& src, bool may_wait, size_t maxMessageLength)
 {
 	String jsonString;
-	StreamReadStatus srs = NetString::ReadStringFromStream(stream, &jsonString, src, may_wait);
+	StreamReadStatus srs = NetString::ReadStringFromStream(stream, &jsonString, src, may_wait, maxMessageLength);
 
 	if (srs != StatusNewItem)
 		return srs;
diff --git a/lib/remote/jsonrpc.hpp b/lib/remote/jsonrpc.hpp
index 1d3232f95..3b0eaabf8 100644
--- a/lib/remote/jsonrpc.hpp
+++ b/lib/remote/jsonrpc.hpp
@@ -36,7 +36,7 @@ class JsonRpc
 {
 public:
 	static size_t SendMessage(const Stream::Ptr& stream, const Dictionary::Ptr& message);
-	static StreamReadStatus ReadMessage(const Stream::Ptr& stream, String *message, StreamReadContext& src, bool may_wait = false);
+	static StreamReadStatus ReadMessage(const Stream::Ptr& stream, String *message, StreamReadContext& src, bool may_wait = false, size_t maxMessageLength = -1);
 	static Dictionary::Ptr DecodeMessage(const String& message);
 
 private:
diff --git a/lib/remote/jsonrpcconnection.cpp b/lib/remote/jsonrpcconnection.cpp
index 978cd9658..bbb6fe180 100644
--- a/lib/remote/jsonrpcconnection.cpp
+++ b/lib/remote/jsonrpcconnection.cpp
@@ -250,9 +250,14 @@ void JsonRpcConnection::MessageHandler(const String& jsonString)
 
 bool JsonRpcConnection::ProcessMessage()
 {
+	ssize_t maxMessageLength = 64 * 1024;
+
+	if (m_Endpoint)
+		maxMessageLength = -1; /* no limit */
+
 	String message;
 
-	StreamReadStatus srs = JsonRpc::ReadMessage(m_Stream, &message, m_Context, false);
+	StreamReadStatus srs = JsonRpc::ReadMessage(m_Stream, &message, m_Context, false, maxMessageLength);
 
 	if (srs != StatusNewItem)
 		return false;