Cleanup TCP throttling code (performance) + move connection checks

3 files changed, 31 insertions, 8 deletions

diff --git a/contrib/epee/include/net/abstract_tcp_server2.h b/contrib/epee/include/net/abstract_tcp_server2.h
index bc0da66e2..be9999203 100644
--- a/contrib/epee/include/net/abstract_tcp_server2.h
+++ b/contrib/epee/include/net/abstract_tcp_server2.h
@@ -76,6 +76,13 @@ namespace net_utils
   protected:
     virtual ~i_connection_filter(){}
   };
+
+  struct i_connection_limit
+  {
+    virtual bool is_host_limit(const epee::net_utils::network_address &address)=0;
+  protected:
+    virtual ~i_connection_limit(){}
+  };
   
 
   /************************************************************************/
@@ -260,10 +267,11 @@ namespace net_utils
     struct shared_state : connection_basic_shared_state, t_protocol_handler::config_type
     {
       shared_state()
-        : connection_basic_shared_state(), t_protocol_handler::config_type(), pfilter(nullptr), stop_signal_sent(false)
+        : connection_basic_shared_state(), t_protocol_handler::config_type(), pfilter(nullptr), plimit(nullptr), stop_signal_sent(false)
       {}
 
       i_connection_filter* pfilter;
+      i_connection_limit* plimit;
       bool stop_signal_sent;
     };
 
@@ -369,6 +377,7 @@ namespace net_utils
     size_t get_threads_count(){return m_threads_count;}
 
     void set_connection_filter(i_connection_filter* pfilter);
+    void set_connection_limit(i_connection_limit* plimit);
 
     void set_default_remote(epee::net_utils::network_address remote)
     {
diff --git a/contrib/epee/include/net/abstract_tcp_server2.inl b/contrib/epee/include/net/abstract_tcp_server2.inl
index d88f18194..8a3a8299c 100644
--- a/contrib/epee/include/net/abstract_tcp_server2.inl
+++ b/contrib/epee/include/net/abstract_tcp_server2.inl
@@ -328,7 +328,7 @@ namespace net_utils
       return;
     }
     auto self = connection<T>::shared_from_this();
-    if (m_connection_type != e_connection_type_RPC) {
+    if (speed_limit_is_enabled()) {
       auto calc_duration = []{
         CRITICAL_REGION_LOCAL(
           network_throttle_manager_t::m_lock_get_global_throttle_in
@@ -382,7 +382,7 @@ namespace net_utils
             m_conn_context.m_max_speed_down,
             speed
           );
-          {
+          if (speed_limit_is_enabled()) {
             CRITICAL_REGION_LOCAL(
               network_throttle_manager_t::m_lock_get_global_throttle_in
             );
@@ -454,7 +454,7 @@ namespace net_utils
       return;
     }
     auto self = connection<T>::shared_from_this();
-    if (m_connection_type != e_connection_type_RPC) {
+    if (speed_limit_is_enabled()) {
       auto calc_duration = [this]{
         CRITICAL_REGION_LOCAL(
           network_throttle_manager_t::m_lock_get_global_throttle_out
@@ -513,7 +513,7 @@ namespace net_utils
             m_conn_context.m_max_speed_down,
             speed
           );
-          {
+          if (speed_limit_is_enabled()) {
             CRITICAL_REGION_LOCAL(
               network_throttle_manager_t::m_lock_get_global_throttle_out
             );
@@ -873,6 +873,13 @@ namespace net_utils
     ).pfilter;
     if (filter && !filter->is_remote_host_allowed(*real_remote))
       return false;
+
+    auto *limit = static_cast<shared_state&>(
+      connection_basic::get_state()
+    ).plimit;
+    if (limit && limit->is_host_limit(*real_remote))
+      return false;
+
     ec_t ec;
     #if !defined(_WIN32) || !defined(__i686)
     connection_basic::socket_.next_layer().set_option(
@@ -1022,7 +1029,7 @@ namespace net_utils
   template<typename T>
   bool connection<T>::speed_limit_is_enabled() const
   {
-    return m_connection_type != e_connection_type_RPC;
+    return m_connection_type == e_connection_type_P2P;
   }
 
   template<typename T>
@@ -1349,6 +1356,13 @@ namespace net_utils
   }
   //---------------------------------------------------------------------------------
   template<class t_protocol_handler>
+  void boosted_tcp_server<t_protocol_handler>::set_connection_limit(i_connection_limit* plimit)
+  {
+    assert(m_state != nullptr); // always set in constructor
+    m_state->plimit = plimit;
+  }
+  //---------------------------------------------------------------------------------
+  template<class t_protocol_handler>
   bool boosted_tcp_server<t_protocol_handler>::run_server(size_t threads_count, bool wait, const boost::thread::attributes& attrs)
   {
     TRY_ENTRY();
diff --git a/contrib/epee/include/net/network_throttle-detail.hpp b/contrib/epee/include/net/network_throttle-detail.hpp
index 0a6dc4a20..3a88105c7 100644
--- a/contrib/epee/include/net/network_throttle-detail.hpp
+++ b/contrib/epee/include/net/network_throttle-detail.hpp
@@ -46,13 +46,13 @@ namespace net_utils
 
 
 class network_throttle : public i_network_throttle {
-	private:
+	public:
 		struct packet_info {
 			size_t m_size; // octets sent. Summary for given small-window (e.g. for all packaged in 1 second)
 			packet_info();
 		};
 
-
+	private:
 		network_speed_bps m_target_speed;
 		size_t m_network_add_cost; // estimated add cost of headers 
 		size_t m_network_minimal_segment; // estimated minimal cost of sending 1 byte to round up to