limit the max message size

3 files changed, 38 insertions, 7 deletions

diff --git a/include/salticidae/conn.h b/include/salticidae/conn.h
index 2f72376..4e2f0f9 100644
--- a/include/salticidae/conn.h
+++ b/include/salticidae/conn.h
@@ -358,10 +358,13 @@ class ConnPool {
     public:
 
     class Config {
-        friend ConnPool;
+        friend class ConnPool;
+        template<typename OpcodeType> friend class MsgNetwork;
         int _max_listen_backlog;
         double _conn_server_timeout;
         size_t _seg_buff_size;
+        size_t _max_msg_size;
+        size_t _max_msg_queue_size;
         size_t _max_recv_buff_size;
         size_t _nworker;
         size_t _queue_capacity;
@@ -378,6 +381,8 @@ class ConnPool {
             _max_listen_backlog(10),
             _conn_server_timeout(2),
             _seg_buff_size(4096),
+            _max_msg_size(1024),
+            _max_msg_queue_size(65536),
             _max_recv_buff_size(4096),
             _nworker(1),
             _queue_capacity(0),
@@ -404,6 +409,16 @@ class ConnPool {
             return *this;
         }
 
+        Config &max_msg_size(size_t x) {
+            _max_msg_size = x;
+            return *this;
+        }
+
+        Config &max_msg_queue_size(size_t x) {
+            _max_msg_queue_size = x;
+            return *this;
+        }
+
         Config &max_recv_buff_size(size_t x) {
             _max_recv_buff_size = x;
             return *this;
diff --git a/include/salticidae/network.h b/include/salticidae/network.h
index 48b555f..18f3a42 100644
--- a/include/salticidae/network.h
+++ b/include/salticidae/network.h
@@ -120,6 +120,8 @@ class MsgNetwork: public ConnPool {
 #endif
 
     private:
+    const size_t max_msg_size;
+    const size_t max_msg_queue_size;
     std::unordered_map<
         typename Msg::opcode_t,
         std::function<void(const Msg &msg, const conn_t &)>> handler_map;
@@ -170,8 +172,10 @@ class MsgNetwork: public ConnPool {
     virtual ~MsgNetwork() { stop(); }
 
     MsgNetwork(const EventContext &ec, const Config &config):
-            ConnPool(ec, config) {
-        incoming_msgs.set_capacity(65536);
+            ConnPool(ec, config),
+            max_msg_size(config._max_msg_size),
+            max_msg_queue_size(config._max_msg_queue_size) {
+        incoming_msgs.set_capacity(max_msg_queue_size);
         incoming_msgs.reg_handler(ec, [this, burst_size=config._burst_size](queue_t &q) {
             std::pair<Msg, conn_t> item;
             size_t cnt = 0;
@@ -560,6 +564,13 @@ void MsgNetwork<OpcodeType>::on_read(const ConnPool::conn_t &_conn) {
             if (recv_buffer.size() < Msg::header_size) break;
             /* new header available */
             msg = Msg(recv_buffer.pop(Msg::header_size));
+            if (msg.get_length() > max_msg_size)
+            {
+                SALTICIDAE_LOG_WARN(
+                        "oversized message from %s, terminating the connection",
+                        std::string(*conn).c_str());
+                throw MsgNetworkError(SALTI_ERROR_CONN_OVERSIZED_MSG);
+            }
             msg_state = Conn::PAYLOAD;
         }
         if (msg_state == Conn::PAYLOAD)
diff --git a/include/salticidae/util.h b/include/salticidae/util.h
index 3196042..b2ddd8e 100644
--- a/include/salticidae/util.h
+++ b/include/salticidae/util.h
@@ -106,7 +106,8 @@ enum SalticidaeErrorCode {
     SALTI_ERROR_RAND_SOURCE,
     SALTI_ERROR_CONN_NOT_READY,
     SALTI_ERROR_NOT_AVAIL,
-    SALTI_ERROR_UNKNOWN
+    SALTI_ERROR_UNKNOWN,
+    SALTI_ERROR_CONN_OVERSIZED_MSG
 };
 
 extern const char *SALTICIDAE_ERROR_STRINGS[];
@@ -148,12 +149,16 @@ struct ConnPoolError: public SalticidaeError {
     using SalticidaeError::SalticidaeError;
 };
 
-class PeerNetworkError: public ConnPoolError {
+class MsgNetworkError: public ConnPoolError {
     using ConnPoolError::ConnPoolError;
 };
 
-class ClientNetworkError: public ConnPoolError {
-    using ConnPoolError::ConnPoolError;
+class PeerNetworkError: public MsgNetworkError {
+    using MsgNetworkError::MsgNetworkError;
+};
+
+class ClientNetworkError: public MsgNetworkError {
+    using MsgNetworkError::MsgNetworkError;
 };
 
 extern const char *TTY_COLOR_RED;