Command max_bytes moved from dynamic map to static switch

8 files changed, 115 insertions, 48 deletions

diff --git a/src/cryptonote_basic/CMakeLists.txt b/src/cryptonote_basic/CMakeLists.txt
index 113fd9d86..5286256c7 100644
--- a/src/cryptonote_basic/CMakeLists.txt
+++ b/src/cryptonote_basic/CMakeLists.txt
@@ -38,6 +38,7 @@ endif()
 
 set(cryptonote_basic_sources
   account.cpp
+  connection_context.cpp
   cryptonote_basic_impl.cpp
   cryptonote_format_utils.cpp
   difficulty.cpp
diff --git a/src/cryptonote_basic/connection_context.cpp b/src/cryptonote_basic/connection_context.cpp
new file mode 100644
index 000000000..a0b8ca1f1
--- /dev/null
+++ b/src/cryptonote_basic/connection_context.cpp
@@ -0,0 +1,71 @@
+// Copyright (c) 2020, The Monero Project
+//
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without modification, are
+// permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this list of
+//    conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice, this list
+//    of conditions and the following disclaimer in the documentation and/or other
+//    materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its contributors may be
+//    used to endorse or promote products derived from this software without specific
+//    prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
+// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL
+// THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
+// THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "connection_context.h"
+
+#include "cryptonote_protocol/cryptonote_protocol_defs.h"
+#include "p2p/p2p_protocol_defs.h"
+
+namespace cryptonote
+{
+  std::size_t cryptonote_connection_context::get_max_bytes(const int command) noexcept
+  {
+    switch (command)
+    {
+    case nodetool::COMMAND_HANDSHAKE_T<cryptonote::CORE_SYNC_DATA>::ID:
+      return 65536;
+    case nodetool::COMMAND_TIMED_SYNC_T<cryptonote::CORE_SYNC_DATA>::ID:
+      return 65536;
+    case nodetool::COMMAND_PING::ID:
+      return 4096;
+    case nodetool::COMMAND_REQUEST_SUPPORT_FLAGS::ID:
+      return 4096;
+    case cryptonote::NOTIFY_NEW_BLOCK::ID:
+      return 1024 * 1024 * 128; // 128 MB (max packet is a bit less than 100 MB though)
+    case cryptonote::NOTIFY_NEW_TRANSACTIONS::ID:
+      return 1024 * 1024 * 128; // 128 MB (max packet is a bit less than 100 MB though)
+    case cryptonote::NOTIFY_REQUEST_GET_OBJECTS::ID:
+      return 1024 * 1024 * 2; // 2 MB
+    case cryptonote::NOTIFY_RESPONSE_GET_OBJECTS::ID:
+      return 1024 * 1024 * 128; // 128 MB (max packet is a bit less than 100 MB though)
+    case cryptonote::NOTIFY_REQUEST_CHAIN::ID:
+      return 512 * 1024; // 512 kB
+    case cryptonote::NOTIFY_RESPONSE_CHAIN_ENTRY::ID:
+      return 1024 * 1024 * 4; // 4 MB
+    case cryptonote::NOTIFY_NEW_FLUFFY_BLOCK::ID:
+      return 1024 * 1024 * 4; // 4 MB, but it does not includes transaction data
+    case cryptonote::NOTIFY_REQUEST_FLUFFY_MISSING_TX::ID:
+      return 1024 * 1024; // 1 MB
+    case cryptonote::NOTIFY_GET_TXPOOL_COMPLEMENT::ID:
+      return 1024 * 1024 * 4; // 4 MB
+    default:
+      break;
+    };
+    return std::numeric_limits<size_t>::max();
+  }
+} // cryptonote
diff --git a/src/cryptonote_basic/connection_context.h b/src/cryptonote_basic/connection_context.h
index 8cd896a36..a7d688300 100644
--- a/src/cryptonote_basic/connection_context.h
+++ b/src/cryptonote_basic/connection_context.h
@@ -39,7 +39,6 @@
 
 namespace cryptonote
 {
-
   struct cryptonote_connection_context: public epee::net_utils::connection_context_base
   {
     cryptonote_connection_context(): m_state(state_before_handshake), m_remote_blockchain_height(0), m_last_response_height(0),
@@ -56,28 +55,11 @@ namespace cryptonote
       state_normal
     };
 
+    static constexpr int handshake_command() noexcept { return 1001; }
     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;
-    }
+    //! \return Maximum number of bytes permissible for `command`.
+    static size_t get_max_bytes(int command) noexcept;
 
     state m_state;
     std::vector<std::pair<crypto::hash, uint64_t>> m_needed_objects;
@@ -95,7 +77,6 @@ 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 5368df7b9..28530f3e7 100644
--- a/src/cryptonote_protocol/cryptonote_protocol_handler.h
+++ b/src/cryptonote_protocol/cryptonote_protocol_handler.h
@@ -110,7 +110,6 @@ 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 a3020495e..d6a08e21b 100644
--- a/src/cryptonote_protocol/cryptonote_protocol_handler.inl
+++ b/src/cryptonote_protocol/cryptonote_protocol_handler.inl
@@ -2876,25 +2876,6 @@ 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.h b/src/p2p/net_node.h
index 1f9fa46a4..59a6e5091 100644
--- a/src/p2p/net_node.h
+++ b/src/p2p/net_node.h
@@ -118,8 +118,6 @@ namespace nodetool
         m_in_timedsync(false)
     {}
 
-    static constexpr int handshake_command() noexcept { return 1001; }
-
     std::vector<cryptonote::blobdata> fluff_txs;
     std::chrono::steady_clock::time_point flush_time;
     peerid_type peer_id;
diff --git a/src/p2p/net_node.inl b/src/p2p/net_node.inl
index 2cea30288..b8bd7b2a7 100644
--- a/src/p2p/net_node.inl
+++ b/src/p2p/net_node.inl
@@ -2636,7 +2636,6 @@ 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>
diff --git a/tests/unit_tests/levin.cpp b/tests/unit_tests/levin.cpp
index dc854ef0c..53a7f7b67 100644
--- a/tests/unit_tests/levin.cpp
+++ b/tests/unit_tests/levin.cpp
@@ -183,13 +183,12 @@ namespace
         }
 
         //\return Number of messages processed
-        std::size_t process_send_queue()
+        std::size_t process_send_queue(const bool valid = true)
         {
             std::size_t count = 0;
             for ( ; !endpoint_.send_queue_.empty(); ++count, endpoint_.send_queue_.pop_front())
             {
-                // invalid messages shoudn't be possible in this test;
-                EXPECT_TRUE(handler_.handle_recv(endpoint_.send_queue_.front().data(), endpoint_.send_queue_.front().size()));
+                EXPECT_EQ(valid, handler_.handle_recv(endpoint_.send_queue_.front().data(), endpoint_.send_queue_.front().size()));
             }
             return count;
         }
@@ -239,6 +238,13 @@ namespace
             return {connection, std::move(request)};
         }
 
+        static received_message get_raw_message(std::deque<received_message>& queue)
+        {
+            received_message out{std::move(queue.front())};
+            queue.pop_front();
+            return out;
+        }
+
         virtual int invoke(int command, const epee::span<const uint8_t> in_buff, epee::byte_slice& buff_out, cryptonote::levin::detail::p2p_context& context) override final
         {
             buff_out = nullptr;
@@ -295,6 +301,11 @@ namespace
         {
             return get_message<T>(notified_);
         }
+
+        received_message get_raw_notification()
+        {
+            return get_raw_message(notified_);
+        }
     };
 
     class levin_notify : public ::testing::Test
@@ -323,6 +334,8 @@ namespace
             EXPECT_EQ(0u, events_.relayed_method_size());
         }
 
+        cryptonote::levin::connections& get_connections() noexcept { return *connections_; }
+
         void add_connection(const bool is_incoming)
         {
             contexts_.emplace_back(io_service_, *connections_, random_generator_, is_incoming);
@@ -2141,3 +2154,27 @@ TEST_F(levin_notify, noise_stem)
         }
     }
 }
+
+TEST_F(levin_notify, command_max_bytes)
+{
+    static constexpr int ping_command = nodetool::COMMAND_PING::ID;
+
+    add_connection(true);
+
+    std::string bytes(4096, 'h');
+
+    EXPECT_EQ(1, get_connections().notify(ping_command, epee::strspan<std::uint8_t>(bytes), contexts_.front().get_id()));
+    EXPECT_EQ(1u, contexts_.front().process_send_queue(true));
+    EXPECT_EQ(1u, receiver_.notified_size());
+
+    const received_message msg = receiver_.get_raw_notification();
+    EXPECT_EQ(ping_command, msg.command);
+    EXPECT_EQ(contexts_.front().get_id(), msg.connection);
+    EXPECT_EQ(bytes, msg.payload);
+
+    bytes.push_back('e');
+
+    EXPECT_EQ(1, get_connections().notify(ping_command, epee::strspan<std::uint8_t>(bytes), contexts_.front().get_id()));
+    EXPECT_EQ(1u, contexts_.front().process_send_queue(false));
+    EXPECT_EQ(0u, receiver_.notified_size());
+}