diff --git a/crates/net/network/src/discovery.rs b/crates/net/network/src/discovery.rs index 07e3248b3..e11fa73c6 100644 --- a/crates/net/network/src/discovery.rs +++ b/crates/net/network/src/discovery.rs @@ -124,14 +124,18 @@ impl Discovery { } /// Processes an incoming [NodeRecord] update from a discovery service - fn on_node_record_update(&mut self, record: NodeRecord, _fork_id: Option) { + fn on_node_record_update(&mut self, record: NodeRecord, fork_id: Option) { let id = record.id; let addr = record.tcp_addr(); match self.discovered_nodes.entry(id) { Entry::Occupied(_entry) => {} Entry::Vacant(entry) => { entry.insert(addr); - self.queued_events.push_back(DiscoveryEvent::Discovered(id, addr)) + self.queued_events.push_back(DiscoveryEvent::Discovered { + peer_id: id, + socket_addr: addr, + fork_id, + }); } } } @@ -210,7 +214,7 @@ impl Discovery { /// Events produced by the [`Discovery`] manager. pub enum DiscoveryEvent { /// A new node was discovered - Discovered(PeerId, SocketAddr), + Discovered { peer_id: PeerId, socket_addr: SocketAddr, fork_id: Option }, /// Retrieved a [`ForkId`] from the peer via ENR request, See EnrForkId(PeerId, ForkId), } diff --git a/crates/net/network/src/peers/manager.rs b/crates/net/network/src/peers/manager.rs index 4c882657b..cfc383908 100644 --- a/crates/net/network/src/peers/manager.rs +++ b/crates/net/network/src/peers/manager.rs @@ -424,9 +424,9 @@ impl PeersManager { /// Called for a newly discovered peer. /// - /// If the peer already exists, then the address and kind will be updated. - pub(crate) fn add_peer(&mut self, peer_id: PeerId, addr: SocketAddr) { - self.add_peer_kind(peer_id, PeerKind::Basic, addr) + /// If the peer already exists, then the address, kind and fork_id will be updated. + pub(crate) fn add_peer(&mut self, peer_id: PeerId, addr: SocketAddr, fork_id: Option) { + self.add_peer_kind(peer_id, PeerKind::Basic, addr, fork_id) } /// Called for a newly discovered trusted peer. @@ -434,13 +434,19 @@ impl PeersManager { /// If the peer already exists, then the address and kind will be updated. #[allow(dead_code)] pub(crate) fn add_trusted_peer(&mut self, peer_id: PeerId, addr: SocketAddr) { - self.add_peer_kind(peer_id, PeerKind::Trusted, addr) + self.add_peer_kind(peer_id, PeerKind::Trusted, addr, None) } /// Called for a newly discovered peer. /// - /// If the peer already exists, then the address and kind will be updated. - pub(crate) fn add_peer_kind(&mut self, peer_id: PeerId, kind: PeerKind, addr: SocketAddr) { + /// If the peer already exists, then the address, kind and fork_id will be updated. + pub(crate) fn add_peer_kind( + &mut self, + peer_id: PeerId, + kind: PeerKind, + addr: SocketAddr, + fork_id: Option, + ) { if self.ban_list.is_banned(&peer_id, &addr.ip()) { return } @@ -450,11 +456,14 @@ impl PeersManager { let node = entry.get_mut(); node.kind = kind; node.addr = addr; + node.fork_id = fork_id; return } Entry::Vacant(entry) => { trace!(target : "net::peers", ?peer_id, ?addr, "discovered new node"); - entry.insert(Peer::with_kind(addr, kind)); + let mut peer = Peer::with_kind(addr, kind); + peer.fork_id = fork_id; + entry.insert(peer); self.queued_actions.push_back(PeerAction::PeerAdded(peer_id)); } } @@ -587,7 +596,7 @@ impl PeersManager { while let Poll::Ready(Some(cmd)) = self.handle_rx.poll_next_unpin(cx) { match cmd { PeerCommand::Add(peer_id, addr) => { - self.add_peer(peer_id, addr); + self.add_peer(peer_id, addr, None); } PeerCommand::Remove(peer) => self.remove_peer(peer), PeerCommand::ReputationChange(peer_id, rep) => { @@ -1080,7 +1089,7 @@ mod test { let peer = PeerId::random(); let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008); let mut peers = PeersManager::default(); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1103,7 +1112,7 @@ mod test { let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008); let mut peers = PeersManager::default(); peers.ban_peer(peer); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::BanPeer { peer_id } => { @@ -1128,7 +1137,7 @@ mod test { PeerBackoffDurations { low: Duration::from_secs(3), ..Default::default() }; let config = PeersConfig { backoff_durations, ..Default::default() }; let mut peers = PeersManager::new(config); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1192,7 +1201,7 @@ mod test { PeerBackoffDurations { high: Duration::from_secs(3), ..Default::default() }; let config = PeersConfig { backoff_durations, ..Default::default() }; let mut peers = PeersManager::new(config); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1252,7 +1261,7 @@ mod test { let peer = PeerId::random(); let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008); let mut peers = PeersManager::default(); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1308,7 +1317,7 @@ mod test { let peer = PeerId::random(); let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008); let mut peers = PeersManager::default(); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1420,7 +1429,7 @@ mod test { let peer = PeerId::random(); let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008); let mut peers = PeersManager::default(); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1464,7 +1473,7 @@ mod test { let peer = PeerId::random(); let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008); let mut peers = PeersManager::default(); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1501,7 +1510,7 @@ mod test { let p = peers.peers.get(&peer).unwrap(); assert_eq!(p.state, PeerConnectionState::DisconnectingOut); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); let p = peers.peers.get(&peer).unwrap(); assert_eq!(p.state, PeerConnectionState::DisconnectingOut); @@ -1514,7 +1523,7 @@ mod test { let peer = PeerId::random(); let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008); let mut peers = PeersManager::default(); - peers.add_peer(peer, socket_addr); + peers.add_peer(peer, socket_addr, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1551,7 +1560,7 @@ mod test { let ban_list = BanList::new(HashSet::new(), vec![ip]); let config = PeersConfig::default().with_ban_list(ban_list); let mut peer_manager = PeersManager::new(config); - peer_manager.add_peer(H512::default(), socket_addr); + peer_manager.add_peer(H512::default(), socket_addr, None); assert!(peer_manager.peers.is_empty()); } @@ -1643,7 +1652,7 @@ mod test { let basic_peer = PeerId::random(); let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009); - peers.add_peer(basic_peer, basic_sock); + peers.add_peer(basic_peer, basic_sock, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { @@ -1683,7 +1692,7 @@ mod test { let basic_peer = PeerId::random(); let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009); - peers.add_peer(basic_peer, basic_sock); + peers.add_peer(basic_peer, basic_sock, None); match event!(peers) { PeerAction::PeerAdded(peer_id) => { diff --git a/crates/net/network/src/state.rs b/crates/net/network/src/state.rs index cd9ea53b6..741da9629 100644 --- a/crates/net/network/src/state.rs +++ b/crates/net/network/src/state.rs @@ -255,7 +255,7 @@ where /// Adds a peer and its address with the given kind to the peerset. pub(crate) fn add_peer_kind(&mut self, peer_id: PeerId, kind: PeerKind, addr: SocketAddr) { - self.peers_manager.add_peer_kind(peer_id, kind, addr) + self.peers_manager.add_peer_kind(peer_id, kind, addr, None) } pub(crate) fn remove_peer(&mut self, peer_id: PeerId, kind: PeerKind) { @@ -268,8 +268,8 @@ where /// Event hook for events received from the discovery service. fn on_discovery_event(&mut self, event: DiscoveryEvent) { match event { - DiscoveryEvent::Discovered(peer, addr) => { - self.peers_manager.add_peer(peer, addr); + DiscoveryEvent::Discovered { peer_id, socket_addr, fork_id } => { + self.peers_manager.add_peer(peer_id, socket_addr, fork_id); } DiscoveryEvent::EnrForkId(peer_id, fork_id) => { self.queued_messages