replace duplicate code with common function

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<Multiaddr> = 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<Multiaddr> = ensure_address_with_peer(addresses.into_iter(), peer);
 
         if addresses.is_empty() {
             tracing::debug!(

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<Self>, service: TransportService) -> crate::Result<()>;
 }
+
+pub fn ensure_address_with_peer(
+    addresses: impl Iterator<Item = Multiaddr>,
+    peer_id: PeerId,
+) -> Vec<Multiaddr> {
+    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()
+}

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<Item = Multiaddr>) {
-        let addresses: HashSet<Multiaddr> = 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());
     }

src/transport/manager/mod.rs

@@ -1381,6 +1381,7 @@ mod tests {
         rx: tokio::sync::mpsc::Receiver<TransportEvent>,
     }
 
+    #[allow(dead_code)]
     impl MockTransport {
         fn new(rx: tokio::sync::mpsc::Receiver<TransportEvent>) -> Self {
             Self { rx }