From cb58759e660e1b518179d4bed650cee1eeaa2c7d Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Fri, 23 Dec 2022 12:15:07 +0100
Subject: [PATCH] feat(net): validate discovered enr forkid (#586)

---
 crates/net/network/src/session/mod.rs |  8 +++++++-
 crates/net/network/src/state.rs       | 11 +++++++++--
 crates/net/network/src/swarm.rs       |  7 +++++++
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/crates/net/network/src/session/mod.rs b/crates/net/network/src/session/mod.rs
index dfb1403ef4..daee59de50 100644
--- a/crates/net/network/src/session/mod.rs
+++ b/crates/net/network/src/session/mod.rs
@@ -19,7 +19,7 @@ use reth_eth_wire::{
     error::EthStreamError,
     DisconnectReason, HelloMessage, Status, UnauthedEthStream, UnauthedP2PStream,
 };
-use reth_primitives::{ForkFilter, ForkTransition, PeerId, H256, U256};
+use reth_primitives::{ForkFilter, ForkId, ForkTransition, PeerId, H256, U256};
 use secp256k1::SecretKey;
 use std::{
     collections::HashMap,
@@ -127,6 +127,12 @@ impl SessionManager {
         }
     }
 
+    /// Check whether the provided [`ForkId`] is compatible based on the validation rules in
+    /// `EIP-2124`.
+    pub(crate) fn is_valid_fork_id(&self, fork_id: ForkId) -> bool {
+        self.fork_filter.validate(fork_id).is_ok()
+    }
+
     /// Returns the next unique [`SessionId`].
     fn next_id(&mut self) -> SessionId {
         let id = self.next_id;
diff --git a/crates/net/network/src/state.rs b/crates/net/network/src/state.rs
index ed8eaec7f1..474a39f1ff 100644
--- a/crates/net/network/src/state.rs
+++ b/crates/net/network/src/state.rs
@@ -262,8 +262,9 @@ where
             DiscoveryEvent::Discovered(peer, addr) => {
                 self.peers_manager.add_discovered_node(peer, addr);
             }
-            DiscoveryEvent::EnrForkId(peer, fork_id) => {
-                self.peers_manager.set_discovered_fork_id(peer, fork_id);
+            DiscoveryEvent::EnrForkId(peer_id, fork_id) => {
+                self.queued_messages
+                    .push_back(StateAction::DiscoveredEnrForkId { peer_id, fork_id });
             }
         }
     }
@@ -455,4 +456,10 @@ pub(crate) enum StateAction {
         /// Why the disconnect was initiated
         reason: Option<DisconnectReason>,
     },
+    /// Retrieved a [`ForkId`] from the peer via ENR request, See <https://eips.ethereum.org/EIPS/eip-868>
+    DiscoveredEnrForkId {
+        peer_id: PeerId,
+        /// The reported [`ForkId`] by this peer.
+        fork_id: ForkId,
+    },
 }
diff --git a/crates/net/network/src/swarm.rs b/crates/net/network/src/swarm.rs
index 9c751890b6..c5b0620e09 100644
--- a/crates/net/network/src/swarm.rs
+++ b/crates/net/network/src/swarm.rs
@@ -227,6 +227,13 @@ where
                 let msg = PeerMessage::NewBlockHashes(hashes);
                 self.sessions.send_message(&peer_id, msg);
             }
+            StateAction::DiscoveredEnrForkId { peer_id, fork_id } => {
+                if self.sessions.is_valid_fork_id(fork_id) {
+                    self.state_mut().peers_mut().set_discovered_fork_id(peer_id, fork_id);
+                } else {
+                    self.state_mut().peers_mut().remove_discovered_node(peer_id);
+                }
+            }
         }
         None
     }