diff --git a/protocols/autonat/src/v2/client/behaviour.rs b/protocols/autonat/src/v2/client/behaviour.rs index 7d558f6434a..a7836add87b 100644 --- a/protocols/autonat/src/v2/client/behaviour.rs +++ b/protocols/autonat/src/v2/client/behaviour.rs @@ -10,15 +10,14 @@ use futures_timer::Delay; use libp2p_core::{multiaddr::Protocol, transport::PortUse, Endpoint, Multiaddr}; use libp2p_identity::PeerId; use libp2p_swarm::{ - behaviour::{ConnectionEstablished, ExternalAddrConfirmed}, - ConnectionClosed, ConnectionDenied, ConnectionHandler, ConnectionId, FromSwarm, - NetworkBehaviour, NewExternalAddrCandidate, NotifyHandler, ToSwarm, + behaviour::ConnectionEstablished, ConnectionClosed, ConnectionDenied, ConnectionHandler, + ConnectionId, FromSwarm, NetworkBehaviour, NewExternalAddrCandidate, NotifyHandler, ToSwarm, }; use rand::prelude::*; use rand_core::OsRng; use std::fmt::{Debug, Display, Formatter}; -use crate::v2::{global_only::IpExt, protocol::DialRequest}; +use crate::v2::{global_only::IpExt, protocol::DialRequest, Nonce}; use super::handler::{dial_back, dial_request}; @@ -60,7 +59,6 @@ pub struct Behaviour where R: RngCore + 'static, { - pending_nonces: HashMap, rng: R, config: Config, pending_events: VecDeque< @@ -111,11 +109,6 @@ where .or_default() .score += 1; } - FromSwarm::ExternalAddrConfirmed(ExternalAddrConfirmed { addr }) => { - if let Some(info) = self.address_candidates.get_mut(addr) { - info.status = TestStatus::Tested; - } - } FromSwarm::ConnectionEstablished(ConnectionEstablished { peer_id, connection_id, @@ -157,13 +150,18 @@ where ) { let (nonce, outcome) = match event { Either::Right(nonce) => { - let Some(status) = self.pending_nonces.get_mut(&nonce) else { + let Some((_, info)) = self + .address_candidates + .iter_mut() + .find(|(_, info)| info.is_pending_with_nonce(nonce)) + else { tracing::warn!(%peer_id, %nonce, "Received unexpected nonce"); return; }; - *status = NonceStatus::Received; + info.status = TestStatus::Received(nonce); tracing::debug!(%peer_id, %nonce, "Successful dial-back"); + return; } Either::Left(dial_request::ToBehaviour::PeerHasServerSupport) => { @@ -178,11 +176,14 @@ where } }; - let nonce_status = self.pending_nonces.remove(&nonce); - let ((tested_addr, bytes_sent), result) = match outcome { Ok(address) => { - if nonce_status != Some(NonceStatus::Received) { + let received_dial_back = self + .address_candidates + .iter_mut() + .any(|(_, info)| info.is_received_with_nonce(nonce)); + + if !received_dial_back { tracing::warn!( %peer_id, %nonce, @@ -198,6 +199,9 @@ where .get_mut(&connection_id) .expect("inconsistent state") .supports_autonat = false; + + self.reset_status_to(nonce, TestStatus::Untested); // Reset so it will be tried again. + return; } Err(dial_request::Error::Io(e)) => { @@ -206,13 +210,20 @@ where %nonce, "Failed to complete AutoNAT probe: {e}" ); + + self.reset_status_to(nonce, TestStatus::Untested); // Reset so it will be tried again. + return; } Err(dial_request::Error::AddressNotReachable { address, bytes_sent, error, - }) => ((address, bytes_sent), Err(error)), + }) => { + self.reset_status_to(nonce, TestStatus::Failed); + + ((address, bytes_sent), Err(error)) + } }; self.pending_events.push_back(ToSwarm::GenerateEvent(Event { @@ -251,7 +262,6 @@ where { pub fn new(rng: R, config: Config) -> Self { Self { - pending_nonces: HashMap::new(), rng, next_tick: Delay::new(config.probe_interval), config, @@ -273,11 +283,10 @@ where }; let nonce = self.rng.gen(); - self.pending_nonces.insert(nonce, NonceStatus::Pending); self.address_candidates .get_mut(&addr) .expect("only emit candidates") - .status = TestStatus::Pending; + .status = TestStatus::Pending(nonce); self.pending_events.push_back(ToSwarm::NotifyHandler { peer_id, @@ -325,9 +334,22 @@ where Some((*conn_id, info.peer_id)) } + fn reset_status_to(&mut self, nonce: Nonce, new_status: TestStatus) { + let Some((_, info)) = self + .address_candidates + .iter_mut() + .find(|(_, i)| i.is_pending_with_nonce(nonce) || i.is_received_with_nonce(nonce)) + else { + return; + }; + + info.status = new_status; + } + + // FIXME: We don't want test-only APIs in our public API. pub fn validate_addr(&mut self, addr: &Multiaddr) { if let Some(info) = self.address_candidates.get_mut(addr) { - info.status = TestStatus::Tested; + info.status = TestStatus::Received(self.rng.next_u64()); } } } @@ -377,12 +399,6 @@ fn addr_is_local(addr: &Multiaddr) -> bool { }) } -#[derive(Debug, PartialEq)] -enum NonceStatus { - Pending, - Received, -} - struct ConnectionInfo { peer_id: PeerId, supports_autonat: bool, @@ -395,10 +411,27 @@ struct AddressInfo { status: TestStatus, } +impl AddressInfo { + fn is_pending_with_nonce(&self, nonce: Nonce) -> bool { + match self.status { + TestStatus::Pending(c) => c == nonce, + _ => false, + } + } + + fn is_received_with_nonce(&self, nonce: Nonce) -> bool { + match self.status { + TestStatus::Received(c) => c == nonce, + _ => false, + } + } +} + #[derive(Clone, Copy, Default, PartialEq)] enum TestStatus { #[default] Untested, - Pending, - Tested, + Pending(Nonce), + Failed, + Received(Nonce), }