Merge pull request #7228 from vtnerd/fix/p2p_memory_usage_017

Add aggressive restrictions to pre-handshake p2p buffer limit

4 files changed, 8 insertions, 0 deletions

diff --git a/tests/fuzz/levin.cpp b/tests/fuzz/levin.cpp
index 012d05f36..0ba0ff7f3 100644
--- a/tests/fuzz/levin.cpp
+++ b/tests/fuzz/levin.cpp
@@ -52,6 +52,8 @@ namespace
 
   struct test_levin_connection_context : public epee::net_utils::connection_context_base
   {
+    static constexpr int handshake_command() noexcept { return 1001; }
+    static constexpr bool handshake_complete() noexcept { return true; }
   };
 
   typedef epee::levin::async_protocol_handler_config<test_levin_connection_context> test_levin_protocol_handler_config;
diff --git a/tests/net_load_tests/net_load_tests.h b/tests/net_load_tests/net_load_tests.h
index 882d42c02..1cc68746a 100644
--- a/tests/net_load_tests/net_load_tests.h
+++ b/tests/net_load_tests/net_load_tests.h
@@ -48,6 +48,8 @@ namespace net_load_tests
   struct test_connection_context : epee::net_utils::connection_context_base
   {
     test_connection_context(): epee::net_utils::connection_context_base(boost::uuids::nil_uuid(), {}, false, false), m_closed(false) {}
+    static constexpr int handshake_command() noexcept { return 1001; }
+    static constexpr bool handshake_complete() noexcept { return true; }
     volatile bool m_closed;
   };
 
diff --git a/tests/unit_tests/epee_levin_protocol_handler_async.cpp b/tests/unit_tests/epee_levin_protocol_handler_async.cpp
index 314fcf9f2..55092fd90 100644
--- a/tests/unit_tests/epee_levin_protocol_handler_async.cpp
+++ b/tests/unit_tests/epee_levin_protocol_handler_async.cpp
@@ -43,6 +43,8 @@ namespace
 {
   struct test_levin_connection_context : public epee::net_utils::connection_context_base
   {
+    static constexpr int handshake_command() noexcept { return 1001; }
+    static constexpr bool handshake_complete() noexcept { return true; }
   };
 
   typedef epee::levin::async_protocol_handler_config<test_levin_connection_context> test_levin_protocol_handler_config;
@@ -194,6 +196,7 @@ namespace
     {
       m_handler_config.set_handler(m_pcommands_handler, [](epee::levin::levin_commands_handler<test_levin_connection_context> *handler) { delete handler; });
       m_handler_config.m_invoke_timeout = invoke_timeout;
+      m_handler_config.m_initial_max_packet_size = max_packet_size;
       m_handler_config.m_max_packet_size = max_packet_size;
     }
 
diff --git a/tests/unit_tests/levin.cpp b/tests/unit_tests/levin.cpp
index 22638942d..76e3282e0 100644
--- a/tests/unit_tests/levin.cpp
+++ b/tests/unit_tests/levin.cpp
@@ -178,6 +178,7 @@ namespace
         {
             using base_type = epee::net_utils::connection_context_base;
             static_cast<base_type&>(context_) = base_type{random_generator(), {}, is_incoming, false};
+            context_.m_state = cryptonote::cryptonote_connection_context::state_normal;
             handler_.after_init_connection();
         }