From 451e476d586228a179baefe15305cb0ff7f5a58e Mon Sep 17 00:00:00 2001
From: Andrey Tvorozhkov <andrey@head-labs.com>
Date: Tue, 10 Sep 2024 00:00:23 +0300
Subject: [PATCH] Calculate connections per ip

---
 adnl/adnl-ext-server.cpp     | 12 ++++++++++++
 adnl/adnl-ext-server.hpp     |  4 ++++
 tdnet/td/net/TcpListener.cpp |  8 --------
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/adnl/adnl-ext-server.cpp b/adnl/adnl-ext-server.cpp
index dda60fbd3..c1247dd9f 100644
--- a/adnl/adnl-ext-server.cpp
+++ b/adnl/adnl-ext-server.cpp
@@ -19,6 +19,7 @@
 #include "adnl-ext-server.hpp"
 #include "keys/encryptor.h"
 #include "utils.hpp"
+#include "td/utils/port/IPAddress.h"
 
 namespace ton {
 
@@ -155,7 +156,18 @@ void AdnlExtServerImpl::add_local_id(AdnlNodeIdShort id) {
   local_ids_.insert(id);
 }
 
+
 void AdnlExtServerImpl::accepted(td::SocketFd fd) {
+  td::IPAddress my;
+  auto s = my.init_peer_address(fd);
+  if (s.is_ok()){
+    auto addr = my.get_ip_host();
+
+    auto &connection_count = ip_connection_count_[addr];
+    connection_count++;
+    LOG(INFO) << "Accept from: " << addr << " connections: " << connection_count;
+  }
+
   td::actor::create_actor<AdnlInboundConnection>(td::actor::ActorOptions().with_name("inconn").with_poll(),
                                                  std::move(fd), peer_table_, actor_id(this))
       .release();
diff --git a/adnl/adnl-ext-server.hpp b/adnl/adnl-ext-server.hpp
index b4604a76d..59f7cde60 100644
--- a/adnl/adnl-ext-server.hpp
+++ b/adnl/adnl-ext-server.hpp
@@ -45,6 +45,7 @@ class AdnlInboundConnection : public AdnlExtConnection {
   td::Status process_init_packet(td::BufferSlice data) override;
   td::Status process_custom_packet(td::BufferSlice &data, bool &processed) override;
   void inited_crypto(td::Result<td::BufferSlice> R);
+  void init_stop();
 
  private:
   td::actor::ActorId<AdnlPeerTable> peer_table_;
@@ -60,6 +61,7 @@ class AdnlExtServerImpl : public AdnlExtServer {
   void add_tcp_port(td::uint16 port) override;
   void add_local_id(AdnlNodeIdShort id) override;
   void accepted(td::SocketFd fd);
+  void stop(std::string ip_addr);
   void decrypt_init_packet(AdnlNodeIdShort dst, td::BufferSlice data, td::Promise<td::BufferSlice> promise);
 
   void start_up() override {
@@ -88,6 +90,8 @@ class AdnlExtServerImpl : public AdnlExtServer {
   std::set<AdnlNodeIdShort> local_ids_;
   std::set<td::uint16> ports_;
   std::map<td::uint16, td::actor::ActorOwn<td::TcpInfiniteListener>> listeners_;
+  std::unordered_map<std::string, int> ip_connection_count_;
+  int max_connections_per_ip_ = 3;
 };
 
 }  // namespace adnl
diff --git a/tdnet/td/net/TcpListener.cpp b/tdnet/td/net/TcpListener.cpp
index be7518504..a08817a57 100644
--- a/tdnet/td/net/TcpListener.cpp
+++ b/tdnet/td/net/TcpListener.cpp
@@ -17,7 +17,6 @@
     Copyright 2017-2020 Telegram Systems LLP
 */
 #include "td/net/TcpListener.h"
-#include "td/utils/port/IPAddress.h"
 
 namespace td {
 TcpListener::TcpListener(int port, std::unique_ptr<Callback> callback) : port_(port), callback_(std::move(callback)) {
@@ -62,13 +61,6 @@ void TcpListener::loop() {
         break;
       }
       TRY_RESULT(client_socket, std::move(r_socket));
-
-      IPAddress my;
-      auto s = my.init_peer_address(client_socket);
-      if (s.is_ok()){
-        LOG(INFO) << "Accept from: " << my.get_ip_host();
-      }
-
       callback_->accept(std::move(client_socket));
     }
     if (td::can_close(server_socket_fd_)) {