diff --git a/src/protocol/libp2p/kademlia/routing_table.rs b/src/protocol/libp2p/kademlia/routing_table.rs index e012318e..8d31dedd 100644 --- a/src/protocol/libp2p/kademlia/routing_table.rs +++ b/src/protocol/libp2p/kademlia/routing_table.rs @@ -22,9 +22,12 @@ //! Kademlia routing table implementation. use crate::{ - protocol::libp2p::kademlia::{ - bucket::{KBucket, KBucketEntry}, - types::{ConnectionType, Distance, KademliaPeer, Key, U256}, + protocol::{ + ensure_address_with_peer, + libp2p::kademlia::{ + bucket::{KBucket, KBucketEntry}, + types::{ConnectionType, Distance, KademliaPeer, Key, U256}, + }, }, transport::{ manager::address::{scores, AddressRecord}, @@ -33,8 +36,7 @@ use crate::{ PeerId, }; -use multiaddr::{Multiaddr, Protocol}; -use multihash::Multihash; +use multiaddr::Multiaddr; /// Number of k-buckets. const NUM_BUCKETS: usize = 256; @@ -187,18 +189,7 @@ impl RoutingTable { "add known peer" ); - // TODO: https://github.com/paritytech/litep2p/issues/337 this has to be moved elsewhere at some point - let addresses: Vec = addresses - .into_iter() - .filter_map(|address| { - let last = address.iter().last(); - if std::matches!(last, Some(Protocol::P2p(_))) { - Some(address) - } else { - Some(address.with(Protocol::P2p(Multihash::from_bytes(&peer.to_bytes()).ok()?))) - } - }) - .collect(); + let addresses: Vec = ensure_address_with_peer(addresses.into_iter(), peer); if addresses.is_empty() { tracing::debug!( diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 36591912..758ac553 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -29,7 +29,8 @@ use crate::{ PeerId, }; -use multiaddr::Multiaddr; +use multiaddr::{Multiaddr, Protocol}; +use multihash::Multihash; use std::fmt::Debug; @@ -141,3 +142,21 @@ pub trait UserProtocol: Send { /// Start the the user protocol event loop. async fn run(self: Box, service: TransportService) -> crate::Result<()>; } + +pub fn ensure_address_with_peer( + addresses: impl Iterator, + peer_id: PeerId, +) -> Vec { + addresses + .filter_map(|address| { + let last = address.iter().last(); + if std::matches!(last, Some(Protocol::P2p(_))) { + Some(address) + } else { + Some(address.with(Protocol::P2p( + Multihash::from_bytes(&peer_id.to_bytes()).ok()?, + ))) + } + }) + .collect() +} diff --git a/src/protocol/transport_service.rs b/src/protocol/transport_service.rs index b729e931..f889c8fd 100644 --- a/src/protocol/transport_service.rs +++ b/src/protocol/transport_service.rs @@ -28,8 +28,7 @@ use crate::{ }; use futures::{future::BoxFuture, stream::FuturesUnordered, Stream, StreamExt}; -use multiaddr::{Multiaddr, Protocol}; -use multihash::Multihash; +use multiaddr::Multiaddr; use tokio::sync::mpsc::{channel, Receiver, Sender}; use std::{ @@ -44,6 +43,8 @@ use std::{ time::{Duration, Instant}, }; +use super::ensure_address_with_peer; + /// Logging target for the file. const LOG_TARGET: &str = "litep2p::transport-service"; @@ -467,15 +468,10 @@ impl TransportService { /// /// The list is filtered for duplicates and unsupported transports. pub fn add_known_address(&mut self, peer: &PeerId, addresses: impl Iterator) { - let addresses: HashSet = addresses - .filter_map(|address| { - if !std::matches!(address.iter().last(), Some(Protocol::P2p(_))) { - Some(address.with(Protocol::P2p(Multihash::from_bytes(&peer.to_bytes()).ok()?))) - } else { - Some(address) - } - }) - .collect(); + let mut addresses = ensure_address_with_peer(addresses.into_iter(), *peer); + + addresses.sort(); + addresses.dedup(); self.transport_handle.add_known_address(peer, addresses.into_iter()); } diff --git a/src/transport/manager/mod.rs b/src/transport/manager/mod.rs index f44a07a6..93562292 100644 --- a/src/transport/manager/mod.rs +++ b/src/transport/manager/mod.rs @@ -1381,6 +1381,7 @@ mod tests { rx: tokio::sync::mpsc::Receiver, } + #[allow(dead_code)] impl MockTransport { fn new(rx: tokio::sync::mpsc::Receiver) -> Self { Self { rx }