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;