From fdf27928335482eae9877087a9534aaefa93c23c Mon Sep 17 00:00:00 2001 From: Matthias Seitz Date: Sun, 26 Mar 2023 19:07:59 +0200 Subject: [PATCH] fix(discv4): track whether endpoint is proven (#1984) --- crates/net/discv4/src/lib.rs | 50 +++++++++++++++++++++++++----------- 1 file changed, 35 insertions(+), 15 deletions(-) diff --git a/crates/net/discv4/src/lib.rs b/crates/net/discv4/src/lib.rs index ac4c28ad5d..99cd0e0f3d 100644 --- a/crates/net/discv4/src/lib.rs +++ b/crates/net/discv4/src/lib.rs @@ -563,7 +563,7 @@ impl Discv4Service { &key, entry, NodeStatus { - state: ConnectionState::Connected, + state: ConnectionState::Disconnected, direction: ConnectionDirection::Outgoing, }, ) { @@ -622,7 +622,7 @@ impl Discv4Service { /// This takes an optional Sender through which all successfully discovered nodes are sent once /// the request has finished. fn lookup_with(&mut self, target: PeerId, tx: Option) { - trace!(target : "net::discv4", ?target, "Starting lookup"); + trace!(target : "discv4", ?target, "Starting lookup"); let target_key = kad_key(target); // Start a lookup context with the 16 (MAX_NODES_PER_BUCKET) closest nodes @@ -630,7 +630,10 @@ impl Discv4Service { target_key.clone(), self.kbuckets .closest_values(&target_key) - .filter(|node| !self.pending_find_nodes.contains_key(&node.key.preimage().0)) + .filter(|node| { + node.value.has_endpoint_proof && + !self.pending_find_nodes.contains_key(&node.key.preimage().0) + }) .take(MAX_NODES_PER_BUCKET) .map(|n| (target_key.distance(&n.key), n.value.record)), tx, @@ -648,7 +651,7 @@ impl Discv4Service { return } - trace!(target : "net::discv4", ?target, num = closest.len(), "Start lookup closest nodes"); + trace!(target : "discv4", ?target, num = closest.len(), "Start lookup closest nodes"); for node in closest { self.find_node(&node, ctx.clone()); @@ -667,11 +670,6 @@ impl Discv4Service { self.pending_find_nodes.insert(node.id, FindNodeRequest::new(ctx)); } - /// Gets the number of entries that are considered connected. - pub fn num_connected(&self) -> usize { - self.kbuckets.buckets_iter().fold(0, |count, bucket| count + bucket.num_connected()) - } - /// Notifies all listeners fn notify(&mut self, update: DiscoveryUpdate) { self.update_listeners.retain_mut(|listener| match listener.try_send(update.clone()) { @@ -719,6 +717,11 @@ impl Discv4Service { removed } + /// Gets the number of entries that are considered connected. + pub fn num_connected(&self) -> usize { + self.kbuckets.buckets_iter().fold(0, |count, bucket| count + bucket.num_connected()) + } + /// Update the entry on RE-ping /// /// On re-ping we check for a changed enr_seq if eip868 is enabled and when it changed we sent a @@ -756,6 +759,7 @@ impl Discv4Service { }; } + /// Callback invoked when we receive a pong from the peer. fn update_on_pong(&mut self, record: NodeRecord, mut last_enr_seq: Option) { if record.id == *self.local_peer_id() { return @@ -773,7 +777,10 @@ impl Discv4Service { let key = kad_key(record.id); match self.kbuckets.entry(&key) { kbucket::Entry::Present(mut entry, old_status) => { + // endpoint is now proven + entry.value_mut().has_endpoint_proof = true; entry.value_mut().update_with_enr(last_enr_seq); + if !old_status.is_connected() { let _ = entry.update(ConnectionState::Connected, Some(old_status.direction)); debug!(target : "discv4", ?record, "added after successful endpoint proof"); @@ -786,7 +793,10 @@ impl Discv4Service { } } kbucket::Entry::Pending(mut entry, mut status) => { + // endpoint is now proven + entry.value().has_endpoint_proof = true; entry.value().update_with_enr(last_enr_seq); + if !status.is_connected() { status.state = ConnectionState::Connected; let _ = entry.update(status); @@ -1869,6 +1879,8 @@ struct NodeEntry { fork_id: Option, /// Counter for failed findNode requests find_node_failures: usize, + /// whether the endpoint of the peer is proven + has_endpoint_proof: bool, } // === impl NodeEntry === @@ -1882,9 +1894,17 @@ impl NodeEntry { last_enr_seq: None, fork_id: None, find_node_failures: 0, + has_endpoint_proof: false, } } + #[cfg(test)] + fn new_proven(record: NodeRecord) -> Self { + let mut node = Self::new(record); + node.has_endpoint_proof = true; + node + } + /// Updates the last timestamp and sets the enr seq fn update_with_enr(&mut self, last_enr_seq: Option) -> Option { self.update_now(|s| std::mem::replace(&mut s.last_enr_seq, last_enr_seq)) @@ -2097,7 +2117,7 @@ mod tests { let _ = service.kbuckets.insert_or_update( &key, - NodeEntry::new(record), + NodeEntry::new_proven(record), NodeStatus { direction: ConnectionDirection::Incoming, state: ConnectionState::Connected, @@ -2164,7 +2184,7 @@ mod tests { let _ = service.kbuckets.insert_or_update( &key, - NodeEntry::new(record), + NodeEntry::new_proven(record), NodeStatus { direction: ConnectionDirection::Incoming, state: ConnectionState::Connected, @@ -2283,11 +2303,11 @@ mod tests { // Since the endpoint was already proven from 1 POV it can already send a FindNode so the // next event is either the PONG or Find Node match event { - Discv4Event::FindNode => { + Discv4Event::FindNode | Discv4Event::EnrRequest => { // since we support enr in the ping it may also request the enr let event = poll_fn(|cx| service_2.poll(cx)).await; match event { - Discv4Event::EnrRequest => { + Discv4Event::FindNode | Discv4Event::EnrRequest => { let event = poll_fn(|cx| service_2.poll(cx)).await; assert_eq!(event, Discv4Event::Pong); } @@ -2298,7 +2318,7 @@ mod tests { } } Discv4Event::Pong => {} - _ => unreachable!(), + ev => unreachable!("{ev:?}"), } // endpoint is proven @@ -2306,7 +2326,7 @@ mod tests { kbucket::Entry::Present(_entry, status) => { assert!(status.is_connected()); } - _ => unreachable!(), + ev => unreachable!("{ev:?}"), } }