clear addresses when endpoints are known

iroh-net/src/magicsock.rs

@@ -386,6 +386,20 @@ impl MagicSock {
         }
     }
 
+    /// Updates our direct addresses.
+    ///
+    /// On a successful update, our address is published to discovery.
+    pub(super) fn update_endpoints(&self, eps: Vec<config::Endpoint>) {
+        let updated = self.endpoints.update(DiscoveredEndpoints::new(eps)).is_ok();
+        if updated {
+            let eps = self.endpoints.read();
+            eps.log_endpoint_change();
+            self.node_map
+                .on_direct_addr_discovered(eps.iter().map(|ep| ep.addr));
+            self.publish_my_addr();
+        }
+    }
+
     /// Get a reference to the DNS resolver used in this [`MagicSock`].
     pub fn dns_resolver(&self) -> &DnsResolver {
         &self.dns_resolver
@@ -2115,15 +2129,7 @@ impl Actor {
             // The STUN address(es) are always first.
             // Despite this sorting, clients are not relying on this sorting for decisions;
 
-            let updated = msock
-                .endpoints
-                .update(DiscoveredEndpoints::new(eps))
-                .is_ok();
-            if updated {
-                let eps = msock.endpoints.read();
-                eps.log_endpoint_change();
-                msock.publish_my_addr();
-            }
+            msock.update_endpoints(eps);
 
             // Regardless of whether our local endpoints changed, we now want to send any queued
             // call-me-maybe messages.

iroh-net/src/magicsock/node_map.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::HashMap,
+    collections::{hash_map::Entry, HashMap},
     hash::Hash,
     net::{IpAddr, SocketAddr},
     path::Path,
@@ -17,7 +17,10 @@ use stun_rs::TransactionId;
 use tokio::io::AsyncWriteExt;
 use tracing::{debug, info, instrument, trace, warn};
 
-use self::node_state::{NodeState, Options, PingHandled};
+use self::{
+    best_addr::ClearReason,
+    node_state::{NodeState, Options, PingHandled},
+};
 use super::{
     metrics::Metrics as MagicsockMetrics, ActorMessage, DiscoMessageSource, QuicMappedAddr,
 };
@@ -307,6 +310,13 @@ impl NodeMap {
     pub(super) fn prune_inactive(&self) {
         self.inner.lock().prune_inactive();
     }
+
+    pub(crate) fn on_direct_addr_discovered(
+        &self,
+        discovered: impl Iterator<Item = impl Into<IpPort>>,
+    ) {
+        self.inner.lock().on_direct_addr_discovered(discovered);
+    }
 }
 
 impl NodeMapInner {
@@ -355,6 +365,34 @@ impl NodeMapInner {
         }
     }
 
+    /// Prunes nodes that claim to share a address we know points to us,
+    pub(super) fn on_direct_addr_discovered(
+        &mut self,
+        discovered: impl Iterator<Item = impl Into<IpPort>>,
+    ) {
+        for addr in discovered {
+            self.remove_by_ipp(addr.into(), ClearReason::MatchesOurLocalAddr)
+        }
+    }
+
+    /// Removes a node by its IpPort
+    fn remove_by_ipp(&mut self, ipp: IpPort, reason: ClearReason) {
+        if let Some(id) = self.by_ip_port.remove(&ipp) {
+            if let Entry::Occupied(mut entry) = self.by_id.entry(id) {
+                let node = entry.get_mut();
+                node.remove_direct_addr(&ipp, reason);
+                if node.direct_addresses().count() == 0 {
+                    let node_id = node.public_key();
+                    let mapped_addr = node.quic_mapped_addr();
+                    self.by_node_key.remove(node_id);
+                    self.by_quic_mapped_addr.remove(mapped_addr);
+                    debug!(node_id=%node_id.fmt_short(), ?reason, "removing node");
+                    entry.remove();
+                }
+            }
+        }
+    }
+
     fn get_id(&self, id: NodeStateKey) -> Option<usize> {
         match id {
             NodeStateKey::Idx(id) => Some(*id),

iroh-net/src/magicsock/node_map/best_addr.rs

@@ -60,11 +60,12 @@ pub(super) enum State<'a> {
     Empty,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 pub enum ClearReason {
     Reset,
     Inactive,
     PongTimeout,
+    MatchesOurLocalAddr,
 }
 
 impl BestAddr {

iroh-net/src/magicsock/node_map/node_state.rs

@@ -311,7 +311,24 @@ impl NodeState {
         (best_addr, relay_url)
     }
 
-    /// Fixup best_addr from candidates.
+    /// Removes a direct address for this node.
+    ///
+    /// If this is also de best address, it will be cleared as well.
+    pub(super) fn remove_direct_addr(&mut self, ip_port: &IpPort, reason: ClearReason) {
+        let Some(state) = self.direct_addr_state.remove(ip_port) else {
+            return;
+        };
+
+        match state.last_alive().map(|instant| instant.elapsed()) {
+            Some(last_alive) => debug!(%ip_port, ?last_alive, ?reason, "pruning address"),
+            None => debug!(%ip_port, last_seen=%"never", ?reason, "pruning address"),
+        }
+
+        self.best_addr
+            .clear_if_equals((*ip_port).into(), reason, self.relay_url.is_some());
+    }
+
+    /// Fixup best_adrr from candidates.
     ///
     /// If somehow we end up in a state where we failed to set a best_addr, while we do have
     /// valid candidates, this will chose a candidate and set best_addr again.  Most likely
@@ -762,19 +779,8 @@ impl NodeState {
         // used ones) last
         prune_candidates.sort_unstable_by_key(|(_ip_port, last_alive)| *last_alive);
         prune_candidates.truncate(prune_count);
-        for (ip_port, last_alive) in prune_candidates.into_iter() {
-            self.direct_addr_state.remove(&ip_port);
-
-            match last_alive.map(|instant| instant.elapsed()) {
-                Some(last_alive) => debug!(%ip_port, ?last_alive, "pruning address"),
-                None => debug!(%ip_port, last_seen=%"never", "pruning address"),
-            }
-
-            self.best_addr.clear_if_equals(
-                ip_port.into(),
-                ClearReason::Inactive,
-                self.relay_url.is_some(),
-            );
+        for (ip_port, _last_alive) in prune_candidates.into_iter() {
+            self.remove_direct_addr(&ip_port, ClearReason::Inactive)
         }
         debug!(
             paths = %summarize_node_paths(&self.direct_addr_state),