From 5c80bc9122911c846b400f21466dbaec2bbb5690 Mon Sep 17 00:00:00 2001
From: Sanket Shanbhag <TechieBoy@users.noreply.github.com>
Date: Fri, 13 Jan 2023 23:00:34 +0530
Subject: [PATCH] fix: validate fork_id before adding peer from discovery
 (#871)

---
 crates/net/network/src/state.rs | 8 +++++++-
 crates/net/network/src/swarm.rs | 6 ++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/crates/net/network/src/state.rs b/crates/net/network/src/state.rs
index a79e9e635b..4d46eb4319 100644
--- a/crates/net/network/src/state.rs
+++ b/crates/net/network/src/state.rs
@@ -270,7 +270,11 @@ where
     fn on_discovery_event(&mut self, event: DiscoveryEvent) {
         match event {
             DiscoveryEvent::Discovered { peer_id, socket_addr, fork_id } => {
-                self.peers_manager.add_peer(peer_id, socket_addr, fork_id);
+                self.queued_messages.push_back(StateAction::DiscoveredNode {
+                    peer_id,
+                    socket_addr,
+                    fork_id,
+                });
             }
             DiscoveryEvent::EnrForkId(peer_id, fork_id) => {
                 self.queued_messages
@@ -482,6 +486,8 @@ pub(crate) enum StateAction {
         /// The reported [`ForkId`] by this peer.
         fork_id: ForkId,
     },
+    /// A new node was found through the discovery, possibly with a ForkId
+    DiscoveredNode { peer_id: PeerId, socket_addr: SocketAddr, fork_id: Option<ForkId> },
     /// A peer was added
     PeerAdded(PeerId),
     /// A peer was dropped
diff --git a/crates/net/network/src/swarm.rs b/crates/net/network/src/swarm.rs
index dc562b88d9..af0bb7ab78 100644
--- a/crates/net/network/src/swarm.rs
+++ b/crates/net/network/src/swarm.rs
@@ -235,6 +235,12 @@ where
             }
             StateAction::PeerAdded(peer_id) => return Some(SwarmEvent::PeerAdded(peer_id)),
             StateAction::PeerRemoved(peer_id) => return Some(SwarmEvent::PeerRemoved(peer_id)),
+            StateAction::DiscoveredNode { peer_id, socket_addr, fork_id } => {
+                // Insert peer only if no fork id or a valid fork id
+                if fork_id.map_or_else(|| true, |f| self.sessions.is_valid_fork_id(f)) {
+                    self.state_mut().peers_mut().add_peer(peer_id, socket_addr, fork_id);
+                }
+            }
             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);