From 8d9f16c387b34f697af92f3cca9808a940aad5a4 Mon Sep 17 00:00:00 2001
From: SChernykh <sergey.v.chernykh@gmail.com>
Date: Tue, 4 Jul 2023 16:47:14 +0200
Subject: [PATCH] P2PServer: sanity check message id

---
 src/p2p_server.cpp |  9 ++++++++-
 src/p2p_server.h   | 21 +++++++++++----------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/p2p_server.cpp b/src/p2p_server.cpp
index 0c07a96..d1d287c 100644
--- a/src/p2p_server.cpp
+++ b/src/p2p_server.cpp
@@ -1418,7 +1418,14 @@ bool P2PServer::P2PClient::on_read(char* data, uint32_t size)
 
 	uint32_t bytes_read;
 	do {
-		MessageId id = static_cast<MessageId>(buf[0]);
+		if (buf[0] > static_cast<uint8_t>(MessageId::LAST)) {
+			LOGWARN(5, "peer " << static_cast<char*>(m_addrString) << " sent an unknown message id " << buf[0]);
+			ban(DEFAULT_BAN_TIME);
+			server->remove_peer_from_list(this);
+			return false;
+		}
+
+		const MessageId id = static_cast<MessageId>(buf[0]);
 
 		// Peer must complete the handshake challenge before sending any other messages
 		if (!m_handshakeComplete && (id != m_expectedMessage)) {
diff --git a/src/p2p_server.h b/src/p2p_server.h
index 63657bf..b2a7f5e 100644
--- a/src/p2p_server.h
+++ b/src/p2p_server.h
@@ -45,16 +45,17 @@ class P2PServer : public TCPServer
 {
 public:
 	enum class MessageId {
-		HANDSHAKE_CHALLENGE = 0,
-		HANDSHAKE_SOLUTION = 1,
-		LISTEN_PORT = 2,
-		BLOCK_REQUEST = 3,
-		BLOCK_RESPONSE = 4,
-		BLOCK_BROADCAST = 5,
-		PEER_LIST_REQUEST = 6,
-		PEER_LIST_RESPONSE = 7,
-		BLOCK_BROADCAST_COMPACT = 8,
-		BLOCK_NOTIFY = 9,
+		HANDSHAKE_CHALLENGE,
+		HANDSHAKE_SOLUTION,
+		LISTEN_PORT,
+		BLOCK_REQUEST,
+		BLOCK_RESPONSE,
+		BLOCK_BROADCAST,
+		PEER_LIST_REQUEST,
+		PEER_LIST_RESPONSE,
+		BLOCK_BROADCAST_COMPACT,
+		BLOCK_NOTIFY,
+		LAST = BLOCK_NOTIFY,
 	};
 
 	explicit P2PServer(p2pool *pool);