From d747cfd2166eb451b6aa57b41926835adbc4aa51 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20RIBEAU?=
 <francois.ribeau.external@stormshield.eu>
Date: Thu, 28 Sep 2023 16:50:58 +0200
Subject: [PATCH] fix(multiaddr): handle multiaddr mismatch with/without p2p
 proto

---
 protocols/identify/src/behaviour.rs         | 1 +
 protocols/kad/src/addresses.rs              | 1 +
 protocols/kad/src/behaviour.rs              | 7 +++++++
 protocols/mdns/src/behaviour/iface/query.rs | 1 +
 swarm/src/lib.rs                            | 2 +-
 5 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/protocols/identify/src/behaviour.rs b/protocols/identify/src/behaviour.rs
index f572b937d386..b052e3804b82 100644
--- a/protocols/identify/src/behaviour.rs
+++ b/protocols/identify/src/behaviour.rs
@@ -472,6 +472,7 @@ impl PeerCache {
             Some(cache) => cache,
         };
 
+        let addresses = addresses.filter_map(|a| libp2p_swarm::p2p_addr(Some(peer), a).ok());
         cache.put(peer, HashSet::from_iter(addresses));
     }
 
diff --git a/protocols/kad/src/addresses.rs b/protocols/kad/src/addresses.rs
index 3c2af4173fde..c82e0d82ab93 100644
--- a/protocols/kad/src/addresses.rs
+++ b/protocols/kad/src/addresses.rs
@@ -23,6 +23,7 @@ use smallvec::SmallVec;
 use std::fmt;
 
 /// A non-empty list of (unique) addresses of a peer in the routing table.
+/// Every address must be a fully-qualified /p2p address.
 #[derive(Clone)]
 pub struct Addresses {
     addrs: SmallVec<[Multiaddr; 6]>,
diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
index 262962cbd1ff..49f2b46d455b 100644
--- a/protocols/kad/src/behaviour.rs
+++ b/protocols/kad/src/behaviour.rs
@@ -525,6 +525,10 @@ where
     /// If the routing table has been updated as a result of this operation,
     /// a [`Event::RoutingUpdated`] event is emitted.
     pub fn add_address(&mut self, peer: &PeerId, address: Multiaddr) -> RoutingUpdate {
+        // ensuring address is a fully-qualified /p2p multiaddr
+        let Ok(address) = libp2p_swarm::p2p_addr(Some(*peer), address) else {
+            return RoutingUpdate::Failed;
+        };
         let key = kbucket::Key::from(*peer);
         match self.kbuckets.entry(&key) {
             kbucket::Entry::Present(mut entry, _) => {
@@ -605,6 +609,9 @@ where
         peer: &PeerId,
         address: &Multiaddr,
     ) -> Option<kbucket::EntryView<kbucket::Key<PeerId>, Addresses>> {
+        let Ok(address) = &libp2p_swarm::p2p_addr(Some(*peer), address.clone()) else {
+            return None;
+        };
         let key = kbucket::Key::from(*peer);
         match self.kbuckets.entry(&key) {
             kbucket::Entry::Present(mut entry, _) => {
diff --git a/protocols/mdns/src/behaviour/iface/query.rs b/protocols/mdns/src/behaviour/iface/query.rs
index 0185028f6ff3..0caa6aa0d63a 100644
--- a/protocols/mdns/src/behaviour/iface/query.rs
+++ b/protocols/mdns/src/behaviour/iface/query.rs
@@ -181,6 +181,7 @@ impl MdnsResponse {
 
                 peer.addresses().iter().filter_map(move |address| {
                     let new_addr = address_translation(address, &observed)?;
+                    let new_addr = libp2p_swarm::p2p_addr(Some(*peer.id()), new_addr).ok()?;
 
                     Some((*peer.id(), new_addr, new_expiration))
                 })
diff --git a/swarm/src/lib.rs b/swarm/src/lib.rs
index 52498156f1f0..c0f97814d99d 100644
--- a/swarm/src/lib.rs
+++ b/swarm/src/lib.rs
@@ -1947,7 +1947,7 @@ impl NetworkInfo {
 ///
 /// If the given address is not yet a `p2p` address for the given peer,
 /// the `/p2p/<peer-id>` protocol is appended to the returned address.
-fn p2p_addr(peer: Option<PeerId>, addr: Multiaddr) -> Result<Multiaddr, Multiaddr> {
+pub fn p2p_addr(peer: Option<PeerId>, addr: Multiaddr) -> Result<Multiaddr, Multiaddr> {
     let peer = match peer {
         Some(p) => p,
         None => return Ok(addr),