fix(discv4): track whether endpoint is proven (#1984)

This commit is contained in:
Matthias Seitz
2023-03-26 19:07:59 +02:00
committed by GitHub
parent f0a95689a4
commit fdf2792833

View File

@ -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<NodeRecordSender>) {
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<u64>) {
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<ForkId>,
/// 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<u64>) -> Option<u64> {
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:?}"),
}
}