From 7b0cc256c557c147c7fc6e6399dfd2faa5a993db Mon Sep 17 00:00:00 2001
From: SChernykh <sergey.v.chernykh@gmail.com>
Date: Mon, 6 Sep 2021 16:17:20 +0200
Subject: [PATCH] P2PServer: ban peers that disconnect without finishing
 handshake

---
 src/p2p_server.cpp | 13 +++++++++++++
 src/p2p_server.h   |  1 +
 src/tcp_server.h   |  1 +
 src/tcp_server.inl |  4 +++-
 4 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/p2p_server.cpp b/src/p2p_server.cpp
index 21f1dd9..956c94c 100644
--- a/src/p2p_server.cpp
+++ b/src/p2p_server.cpp
@@ -912,6 +912,19 @@ bool P2PServer::P2PClient::on_read(char* data, uint32_t size)
 	return true;
 }
 
+void P2PServer::P2PClient::on_read_failed(int /*err*/)
+{
+	if (!m_handshakeComplete) {
+		LOGWARN(5, "peer " << static_cast<char*>(m_addrString) << " disconnected before finishing handshake");
+
+		ban(DEFAULT_BAN_TIME);
+		P2PServer* server = static_cast<P2PServer*>(m_owner);
+		if (server) {
+			server->remove_peer_from_list(this);
+		}
+	}
+}
+
 bool P2PServer::P2PClient::send_handshake_challenge()
 {
 	P2PServer* owner = static_cast<P2PServer*>(m_owner);
diff --git a/src/p2p_server.h b/src/p2p_server.h
index 949aa56..04b6c90 100644
--- a/src/p2p_server.h
+++ b/src/p2p_server.h
@@ -63,6 +63,7 @@ public:
 		void reset() override;
 		bool on_connect() override;
 		bool on_read(char* data, uint32_t size) override;
+		void on_read_failed(int err) override;
 
 		// Both peers send handshake challenge immediately after a connection is established
 		// Both peers must have the same consensus ID for handshake to succeed
diff --git a/src/tcp_server.h b/src/tcp_server.h
index 438ebb3..997378d 100644
--- a/src/tcp_server.h
+++ b/src/tcp_server.h
@@ -89,6 +89,7 @@ public:
 		virtual void reset();
 		virtual bool on_connect() = 0;
 		virtual bool on_read(char* data, uint32_t size) = 0;
+		virtual void on_read_failed(int /*err*/) {}
 
 		static void on_alloc(uv_handle_t* handle, size_t suggested_size, uv_buf_t* buf);
 		static void on_read(uv_stream_t* stream, ssize_t nread, const uv_buf_t* buf);
diff --git a/src/tcp_server.inl b/src/tcp_server.inl
index c2409cc..4913ee8 100644
--- a/src/tcp_server.inl
+++ b/src/tcp_server.inl
@@ -835,7 +835,9 @@ void TCPServer<READ_BUF_SIZE, WRITE_BUF_SIZE>::Client::on_read(uv_stream_t* stre
 	}
 	else if (nread < 0) {
 		if (nread != UV_EOF) {
-			LOGWARN(5, "client " << static_cast<const char*>(pThis->m_addrString) << " failed to read response, err = " << uv_err_name(static_cast<int>(nread)));
+			const int err = static_cast<int>(nread);
+			LOGWARN(5, "client " << static_cast<const char*>(pThis->m_addrString) << " failed to read response, err = " << uv_err_name(err));
+			pThis->on_read_failed(err);
 		}
 		pThis->close();
 	}