4 files changed, 44 insertions, 0 deletions
diff --git a/src/cryptonote_basic/connection_context.h b/src/cryptonote_basic/connection_context.h
index e5c00d4f3..8cd896a36 100644
--- a/src/cryptonote_basic/connection_context.h
+++ b/src/cryptonote_basic/connection_context.h
@@ -31,6 +31,7 @@
 #pragma once
 #include <unordered_set>
 #include <atomic>
+#include <algorithm>
 #include <boost/date_time/posix_time/posix_time.hpp>
 #include "net/net_utils_base.h"
 #include "copyable_atomic.h"
@@ -57,6 +58,27 @@ namespace cryptonote
 
     bool handshake_complete() const noexcept { return m_state != state_before_handshake; }
 
+    void set_max_bytes(int command, size_t bytes) {
+      const auto i = std::lower_bound(m_max_bytes.begin(), m_max_bytes.end(), std::make_pair(command, bytes), [](const std::pair<int, size_t> &e0, const std::pair<int, size_t> &e1){
+        return e0.first < e1.first;
+      });
+      if (i == m_max_bytes.end())
+        m_max_bytes.push_back(std::make_pair(command, bytes));
+      else if (i->first == command)
+        i->second = bytes;
+      else
+        m_max_bytes.insert(i, std::make_pair(command, bytes));
+    }
+    size_t get_max_bytes(int command) const {
+      const auto i = std::lower_bound(m_max_bytes.begin(), m_max_bytes.end(), std::make_pair(command, 0), [](const std::pair<int, size_t> &e0, const std::pair<int, size_t> &e1){
+        return e0.first < e1.first;
+      });
+      if (i == m_max_bytes.end() || i->first != command)
+        return std::numeric_limits<size_t>::max();
+      else
+        return i->second;
+    }
+
     state m_state;
     std::vector<std::pair<crypto::hash, uint64_t>> m_needed_objects;
     std::unordered_set<crypto::hash> m_requested_objects;
@@ -73,6 +95,7 @@ namespace cryptonote
     int m_expect_response;
     uint64_t m_expect_height;
     size_t m_num_requested;
+    std::vector<std::pair<int, size_t>> m_max_bytes;
   };
 
   inline std::string get_protocol_state_string(cryptonote_connection_context::state s)
diff --git a/src/cryptonote_protocol/cryptonote_protocol_handler.h b/src/cryptonote_protocol/cryptonote_protocol_handler.h
index 28530f3e7..5368df7b9 100644
--- a/src/cryptonote_protocol/cryptonote_protocol_handler.h
+++ b/src/cryptonote_protocol/cryptonote_protocol_handler.h
@@ -110,6 +110,7 @@ namespace cryptonote
     std::list<connection_info> get_connections();
     const block_queue &get_block_queue() const { return m_block_queue; }
     void stop();
+    void on_connection_new(cryptonote_connection_context &context);
     void on_connection_close(cryptonote_connection_context &context);
     void set_max_out_peers(unsigned int max) { m_max_out_peers = max; }
     bool no_sync() const { return m_no_sync; }
diff --git a/src/cryptonote_protocol/cryptonote_protocol_handler.inl b/src/cryptonote_protocol/cryptonote_protocol_handler.inl
index bd39077ae..c54974387 100644
--- a/src/cryptonote_protocol/cryptonote_protocol_handler.inl
+++ b/src/cryptonote_protocol/cryptonote_protocol_handler.inl
@@ -2873,6 +2873,25 @@ skip:
   }
   //------------------------------------------------------------------------------------------------------------------------
   template<class t_core>
+  void t_cryptonote_protocol_handler<t_core>::on_connection_new(cryptonote_connection_context &context)
+  {
+    context.set_max_bytes(nodetool::COMMAND_HANDSHAKE_T<cryptonote::CORE_SYNC_DATA>::ID, 65536);
+    context.set_max_bytes(nodetool::COMMAND_TIMED_SYNC_T<cryptonote::CORE_SYNC_DATA>::ID, 65536);
+    context.set_max_bytes(nodetool::COMMAND_PING::ID, 4096);
+    context.set_max_bytes(nodetool::COMMAND_REQUEST_SUPPORT_FLAGS::ID, 4096);
+
+    context.set_max_bytes(cryptonote::NOTIFY_NEW_BLOCK::ID, 1024 * 1024 * 128); // 128 MB (max packet is a bit less than 100 MB though)
+    context.set_max_bytes(cryptonote::NOTIFY_NEW_TRANSACTIONS::ID, 1024 * 1024 * 128); // 128 MB (max packet is a bit less than 100 MB though)
+    context.set_max_bytes(cryptonote::NOTIFY_REQUEST_GET_OBJECTS::ID, 1024 * 1024 * 2); // 2 MB
+    context.set_max_bytes(cryptonote::NOTIFY_RESPONSE_GET_OBJECTS::ID, 1024 * 1024 * 128); // 128 MB (max packet is a bit less than 100 MB though)
+    context.set_max_bytes(cryptonote::NOTIFY_REQUEST_CHAIN::ID, 512 * 1024); // 512 kB
+    context.set_max_bytes(cryptonote::NOTIFY_RESPONSE_CHAIN_ENTRY::ID, 1024 * 1024 * 4); // 4 MB
+    context.set_max_bytes(cryptonote::NOTIFY_NEW_FLUFFY_BLOCK::ID, 1024 * 1024 * 4); // 4 MB, but it does not includes transaction data
+    context.set_max_bytes(cryptonote::NOTIFY_REQUEST_FLUFFY_MISSING_TX::ID, 1024 * 1024); // 1 MB
+    context.set_max_bytes(cryptonote::NOTIFY_GET_TXPOOL_COMPLEMENT::ID, 1024 * 1024 * 4); // 4 MB
+  }
+  //------------------------------------------------------------------------------------------------------------------------
+  template<class t_core>
   void t_cryptonote_protocol_handler<t_core>::on_connection_close(cryptonote_connection_context &context)
   {
     uint64_t target = 0;
diff --git a/src/p2p/net_node.inl b/src/p2p/net_node.inl
index 487e89923..72f0dc221 100644
--- a/src/p2p/net_node.inl
+++ b/src/p2p/net_node.inl
@@ -2633,6 +2633,7 @@ namespace nodetool
   void node_server<t_payload_net_handler>::on_connection_new(p2p_connection_context& context)
   {
     MINFO("["<< epee::net_utils::print_connection_context(context) << "] NEW CONNECTION");
+    m_payload_handler.on_connection_new(context);
   }
   //-----------------------------------------------------------------------------------
   template<class t_payload_net_handler>