Add aggressive restrictions to pre-handshake p2p buffer limit

1 files changed, 23 insertions, 9 deletions

diff --git a/contrib/epee/include/net/levin_protocol_handler_async.h b/contrib/epee/include/net/levin_protocol_handler_async.h
index 8cb2be3e1..9e10747c2 100644
--- a/contrib/epee/include/net/levin_protocol_handler_async.h
+++ b/contrib/epee/include/net/levin_protocol_handler_async.h
@@ -84,7 +84,8 @@ class async_protocol_handler_config
 
 public:
   typedef t_connection_context connection_context;
-  uint64_t m_max_packet_size; 
+  uint64_t m_initial_max_packet_size;
+  uint64_t m_max_packet_size;
   uint64_t m_invoke_timeout;
 
   int invoke(int command, const epee::span<const uint8_t> in_buff, std::string& buff_out, boost::uuids::uuid connection_id);
@@ -105,7 +106,7 @@ public:
   size_t get_in_connections_count();
   void set_handler(levin_commands_handler<t_connection_context>* handler, void (*destroy)(levin_commands_handler<t_connection_context>*) = NULL);
 
-  async_protocol_handler_config():m_pcommands_handler(NULL), m_pcommands_handler_destroy(NULL), m_max_packet_size(LEVIN_DEFAULT_MAX_PACKET_SIZE), m_invoke_timeout(LEVIN_DEFAULT_TIMEOUT_PRECONFIGURED)
+  async_protocol_handler_config():m_pcommands_handler(NULL), m_pcommands_handler_destroy(NULL), m_initial_max_packet_size(LEVIN_INITIAL_MAX_PACKET_SIZE), m_max_packet_size(LEVIN_DEFAULT_MAX_PACKET_SIZE), m_invoke_timeout(LEVIN_DEFAULT_TIMEOUT_PRECONFIGURED)
   {}
   ~async_protocol_handler_config() { set_handler(NULL, NULL); }
   void del_out_connections(size_t count);
@@ -162,6 +163,7 @@ public:
   net_utils::i_service_endpoint* m_pservice_endpoint; 
   config_type& m_config;
   t_connection_context& m_connection_context;
+  std::atomic<uint64_t> m_max_packet_size;
 
   net_utils::buffer m_cache_in_buffer;
   stream_state m_state;
@@ -289,7 +291,8 @@ public:
             m_current_head(bucket_head2()),
             m_pservice_endpoint(psnd_hndlr), 
             m_config(config), 
-            m_connection_context(conn_context), 
+            m_connection_context(conn_context),
+            m_max_packet_size(config.m_initial_max_packet_size),
             m_cache_in_buffer(4 * 1024),
             m_state(stream_state_head)
   {
@@ -399,13 +402,14 @@ public:
     }
 
     // these should never fail, but do runtime check for safety
-    CHECK_AND_ASSERT_MES(m_config.m_max_packet_size >= m_cache_in_buffer.size(), false, "Bad m_cache_in_buffer.size()");
-    CHECK_AND_ASSERT_MES(m_config.m_max_packet_size - m_cache_in_buffer.size() >= m_fragment_buffer.size(), false, "Bad m_cache_in_buffer.size() + m_fragment_buffer.size()");
+    const uint64_t max_packet_size = m_max_packet_size;
+    CHECK_AND_ASSERT_MES(max_packet_size >= m_cache_in_buffer.size(), false, "Bad m_cache_in_buffer.size()");
+    CHECK_AND_ASSERT_MES(max_packet_size - m_cache_in_buffer.size() >= m_fragment_buffer.size(), false, "Bad m_cache_in_buffer.size() + m_fragment_buffer.size()");
 
     // flipped to subtraction; prevent overflow since m_max_packet_size is variable and public
-    if(cb > m_config.m_max_packet_size - m_cache_in_buffer.size() - m_fragment_buffer.size())
+    if(cb > max_packet_size - m_cache_in_buffer.size() - m_fragment_buffer.size())
     {
-      MWARNING(m_connection_context << "Maximum packet size exceed!, m_max_packet_size = " << m_config.m_max_packet_size
+      MWARNING(m_connection_context << "Maximum packet size exceed!, m_max_packet_size = " << max_packet_size
                           << ", packet received " << m_cache_in_buffer.size() +  cb 
                           << ", connection will be closed.");
       return false;
@@ -519,6 +523,10 @@ public:
                 m_current_head.m_command, buff_to_invoke, return_buff, m_connection_context
               );
 
+              // peer_id remains unset if dropped
+              if (m_current_head.m_command == m_connection_context.handshake_command() && m_connection_context.handshake_complete())
+                m_max_packet_size = m_config.m_max_packet_size;
+
               bucket_head2 head = make_header(m_current_head.m_command, return_buff.size(), LEVIN_PACKET_RESPONSE, false);
               head.m_return_code = SWAP32LE(return_code);
 
@@ -576,9 +584,9 @@ public:
           m_cache_in_buffer.erase(sizeof(bucket_head2));
           m_state = stream_state_body;
           m_oponent_protocol_ver = m_current_head.m_protocol_version;
-          if(m_current_head.m_cb > m_config.m_max_packet_size)
+          if(m_current_head.m_cb > max_packet_size)
           {
-            LOG_ERROR_CC(m_connection_context, "Maximum packet size exceed!, m_max_packet_size = " << m_config.m_max_packet_size 
+            LOG_ERROR_CC(m_connection_context, "Maximum packet size exceed!, m_max_packet_size = " << max_packet_size
               << ", packet header received " << m_current_head.m_cb 
               << ", connection will be closed.");
             return false;
@@ -633,6 +641,9 @@ public:
       boost::interprocess::ipcdetail::atomic_write32(&m_invoke_buf_ready, 0);
       CRITICAL_REGION_BEGIN(m_invoke_response_handlers_lock);
 
+      if (command == m_connection_context.handshake_command())
+        m_max_packet_size = m_config.m_max_packet_size;
+
       if(!send_message(command, in_buff, LEVIN_PACKET_REQUEST, true))
       {
         LOG_ERROR_CC(m_connection_context, "Failed to do_send");
@@ -674,6 +685,9 @@ public:
 
     boost::interprocess::ipcdetail::atomic_write32(&m_invoke_buf_ready, 0);
 
+    if (command == m_connection_context.handshake_command())
+      m_max_packet_size = m_config.m_max_packet_size;
+
     if (!send_message(command, in_buff, LEVIN_PACKET_REQUEST, true))
     {
       LOG_ERROR_CC(m_connection_context, "Failed to send request");