fix: validate fork_id before adding peer from discovery (paradigmxyz#871

)
adpthegreat · Jan 13, 2023 · 5c80bc9 · 5c80bc9
1 parent a5bdfe5
commit 5c80bc9
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git 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
@@ -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);