From b45540ca08cd0fad66736cd5f9de2e225c2dd97c Mon Sep 17 00:00:00 2001
From: SChernykh <sergey.v.chernykh@gmail.com>
Date: Fri, 29 Oct 2021 14:24:05 +0200
Subject: [PATCH] TCPServer: clean up old IP bans

---
 src/common.h       | 28 ++++++++++++++++++++++++++++
 src/p2p_server.h   |  2 +-
 src/tcp_server.h   | 36 ++----------------------------------
 src/tcp_server.inl |  8 ++++++--
 src/util.h         | 21 ++++++++++++++++++++-
 src/uv_util.h      |  3 +++
 6 files changed, 60 insertions(+), 38 deletions(-)

diff --git a/src/common.h b/src/common.h
index f9f2aef..d35ede1 100644
--- a/src/common.h
+++ b/src/common.h
@@ -279,6 +279,34 @@ enum class NetworkType {
 	Stagenet,
 };
 
+struct raw_ip
+{
+	alignas(8) uint8_t data[16];
+
+	FORCEINLINE bool operator<(const raw_ip& other) const
+	{
+		const uint64_t* a = reinterpret_cast<const uint64_t*>(data);
+		const uint64_t* b = reinterpret_cast<const uint64_t*>(other.data);
+
+		if (a[1] < b[1]) return true;
+		if (a[1] > b[1]) return false;
+
+		return a[0] < b[0];
+	}
+
+	FORCEINLINE bool operator==(const raw_ip& other) const
+	{
+		const uint64_t* a = reinterpret_cast<const uint64_t*>(data);
+		const uint64_t* b = reinterpret_cast<const uint64_t*>(other.data);
+
+		return (a[0] == b[0]) && (a[1] == b[1]);
+	}
+
+	FORCEINLINE bool operator!=(const raw_ip& other) const { return !operator==(other); }
+};
+
+static_assert(sizeof(raw_ip) == 16, "struct raw_ip has invalid size");
+
 void* malloc_hook(size_t n) noexcept;
 void* realloc_hook(void* ptr, size_t size) noexcept;
 void* calloc_hook(size_t count, size_t size) noexcept;
diff --git a/src/p2p_server.h b/src/p2p_server.h
index 88ab127..0bcbc60 100644
--- a/src/p2p_server.h
+++ b/src/p2p_server.h
@@ -192,7 +192,7 @@ private:
 	std::vector<Broadcast*> m_broadcastQueue;
 
 	uv_mutex_t m_missingBlockRequestsLock;
-	std::set<std::pair<uint64_t, uint64_t>> m_missingBlockRequests;
+	unordered_set<std::pair<uint64_t, uint64_t>> m_missingBlockRequests;
 
 	static void on_broadcast(uv_async_t* handle) { reinterpret_cast<P2PServer*>(handle->data)->on_broadcast(); }
 	void on_broadcast();
diff --git a/src/tcp_server.h b/src/tcp_server.h
index 4c12dd6..0d30124 100644
--- a/src/tcp_server.h
+++ b/src/tcp_server.h
@@ -18,8 +18,6 @@
 #pragma once
 
 #include "uv_util.h"
-#include <map>
-#include <set>
 
 namespace p2pool {
 
@@ -46,36 +44,6 @@ public:
 
 	int listen_port() const { return m_listenPort; }
 
-	struct raw_ip
-	{
-		alignas(8) uint8_t data[16];
-
-		FORCEINLINE bool operator<(const raw_ip& other) const
-		{
-			const uint64_t* a = reinterpret_cast<const uint64_t*>(data);
-			const uint64_t* b = reinterpret_cast<const uint64_t*>(other.data);
-
-			if (a[1] < b[1]) return true;
-			if (a[1] > b[1]) return false;
-
-			return a[0] < b[0];
-		}
-
-		FORCEINLINE bool operator==(const raw_ip& other) const
-		{
-			const uint64_t* a = reinterpret_cast<const uint64_t*>(data);
-			const uint64_t* b = reinterpret_cast<const uint64_t*>(other.data);
-
-			return (a[0] == b[0]) && (a[1] == b[1]);
-		}
-
-		FORCEINLINE bool operator!=(const raw_ip& other) const { return !operator==(other); }
-	};
-
-	static_assert(sizeof(raw_ip) == 16, "struct raw_ip has invalid size");
-	static_assert(sizeof(in6_addr) == 16, "struct in6_addr has invalid size");
-	static_assert(sizeof(in_addr) == 4, "struct in_addr has invalid size");
-
 	bool connect_to_peer(bool is_v6, const raw_ip& ip, int port);
 	virtual void on_connect_failed(bool is_v6, const raw_ip& ip, int port);
 
@@ -193,12 +161,12 @@ protected:
 	uint32_t m_numIncomingConnections;
 
 	uv_mutex_t m_bansLock;
-	std::map<raw_ip, time_t> m_bans;
+	unordered_map<raw_ip, time_t> m_bans;
 
 	bool is_banned(const raw_ip& ip);
 
 	uv_mutex_t m_pendingConnectionsLock;
-	std::set<raw_ip> m_pendingConnections;
+	unordered_set<raw_ip> m_pendingConnections;
 
 	uv_async_t m_dropConnectionsAsync;
 	static void on_drop_connections(uv_async_t* async) { reinterpret_cast<TCPServer*>(async->data)->close_sockets(false); }
diff --git a/src/tcp_server.inl b/src/tcp_server.inl
index b2a8563..73d61ba 100644
--- a/src/tcp_server.inl
+++ b/src/tcp_server.inl
@@ -319,8 +319,12 @@ bool TCPServer<READ_BUF_SIZE, WRITE_BUF_SIZE>::is_banned(const raw_ip& ip)
 	MutexLock lock(m_bansLock);
 
 	auto it = m_bans.find(ip);
-	if ((it != m_bans.end()) && (time(nullptr) < it->second)) {
-		return true;
+	if (it != m_bans.end()) {
+		const bool banned = (time(nullptr) < it->second);
+		if (!banned) {
+			m_bans.erase(it);
+		}
+		return banned;
 	}
 
 	return false;
diff --git a/src/util.h b/src/util.h
index d4524d8..073a351 100644
--- a/src/util.h
+++ b/src/util.h
@@ -19,7 +19,7 @@
 
 #ifdef _MSC_VER
 #pragma warning(push)
-#pragma warning(disable : 5027)
+#pragma warning(disable : 4623 5026 5027)
 #endif
 
 #define ROBIN_HOOD_MALLOC(size) p2pool::malloc_hook(size)
@@ -165,4 +165,23 @@ struct hash<std::array<uint8_t, N>>
 	}
 };
 
+template<>
+struct hash<p2pool::raw_ip>
+{
+	FORCEINLINE size_t operator()(const p2pool::raw_ip& value) const noexcept
+	{
+		return hash_bytes(value.data, sizeof(value.data));
+	}
+};
+
+template<>
+struct hash<std::pair<uint64_t, uint64_t>>
+{
+	FORCEINLINE size_t operator()(const std::pair<uint64_t, uint64_t>& value) const noexcept
+	{
+		static_assert(sizeof(value) == sizeof(uint64_t) * 2, "Invalid std::pair<uint64_t, uint64_t> size");
+		return hash_bytes(&value, sizeof(value));
+	}
+};
+
 } // namespace robin_hood
diff --git a/src/uv_util.h b/src/uv_util.h
index 3d11d13..01d0812 100644
--- a/src/uv_util.h
+++ b/src/uv_util.h
@@ -19,6 +19,9 @@
 
 #include <uv.h>
 
+static_assert(sizeof(in6_addr) == 16, "struct in6_addr has invalid size");
+static_assert(sizeof(in_addr) == 4, "struct in_addr has invalid size");
+
 namespace p2pool {
 
 struct MutexLock : public nocopy_nomove