From c4f7877853d9792e1bf277ae475ffb2caf93b0c0 Mon Sep 17 00:00:00 2001
From: Marco Munizaga <marco@marcopolo.io>
Date: Tue, 16 Nov 2021 04:05:47 -0800
Subject: [PATCH] protocols/identify: Check multiaddr has valid peer id
 component prior to caching (#2338)

Check multiaddr has valid peer id component or none at all. We don't want to
cache a multiaddr with a purposely wrong multiaddr. e.g. something that ends
with .../p2p/some-other-peer. While this should fail to dial because we [check
this before
dialing](https://github.com/libp2p/rust-libp2p/blob/master/core/src/connection/pool/concurrent_dial.rs#L144),
it's better to not cache this in the first place.
---
 CHANGELOG.md                       |  1 +
 protocols/identify/CHANGELOG.md    |  2 ++
 protocols/identify/src/identify.rs | 39 +++++++++++++++++++++++++++++-
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5e691d03..68b5fe28 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,6 +45,7 @@
 ## Version 0.41.0 [unreleased]
 
 - Update individual crates.
+    - `libp2p-identify`
     - `libp2p-kad`
     - `libp2p-websocket`
 - Forward `wasm-bindgen` feature to `futures-timer`, `instant`, `parking_lot`, `getrandom/js` and `rand/wasm-bindgen`.
diff --git a/protocols/identify/CHANGELOG.md b/protocols/identify/CHANGELOG.md
index b6cb0762..849e517e 100644
--- a/protocols/identify/CHANGELOG.md
+++ b/protocols/identify/CHANGELOG.md
@@ -1,10 +1,12 @@
 # 0.32.0 [unreleased]
 
 - Use `futures-timer` instead of `wasm-timer` (see [PR 2245]).
+- Filter invalid peers from cache used in `addresses_of_peer` – [PR 2338].
 
 - Update dependencies.
 
 [PR 2245]: https://github.com/libp2p/rust-libp2p/pull/2245
+[PR 2338]: https://github.com/libp2p/rust-libp2p/pull/2338
 
 # 0.31.0 [2021-11-01]
 
diff --git a/protocols/identify/src/identify.rs b/protocols/identify/src/identify.rs
index 09859216..e3600e0d 100644
--- a/protocols/identify/src/identify.rs
+++ b/protocols/identify/src/identify.rs
@@ -23,6 +23,7 @@ use crate::protocol::{IdentifyInfo, ReplySubstream};
 use futures::prelude::*;
 use libp2p_core::{
     connection::{ConnectionId, ListenerId},
+    multiaddr::Protocol,
     upgrade::UpgradeError,
     ConnectedPoint, Multiaddr, PeerId, PublicKey,
 };
@@ -304,7 +305,11 @@ impl NetworkBehaviour for Identify {
         event: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
     ) {
         match event {
-            IdentifyHandlerEvent::Identified(info) => {
+            IdentifyHandlerEvent::Identified(mut info) => {
+                // Remove invalid multiaddrs.
+                info.listen_addrs
+                    .retain(|addr| multiaddr_matches_peer_id(addr, &peer_id));
+
                 // Replace existing addresses to prevent other peer from filling up our memory.
                 self.discovered_peers
                     .put(peer_id, HashSet::from_iter(info.listen_addrs.clone()));
@@ -499,6 +504,16 @@ fn listen_addrs(params: &impl PollParameters) -> Vec<Multiaddr> {
     listen_addrs
 }
 
+/// If there is a given peer_id in the multiaddr, make sure it is the same as
+/// the given peer_id. If there is no peer_id for the peer in the mutiaddr, this returns true.
+fn multiaddr_matches_peer_id(addr: &Multiaddr, peer_id: &PeerId) -> bool {
+    let last_component = addr.iter().last();
+    if let Some(Protocol::P2p(multi_addr_peer_id)) = last_component {
+        return multi_addr_peer_id == *peer_id.as_ref();
+    }
+    return true;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -771,4 +786,26 @@ mod tests {
 
         assert_eq!(connected_peer, swarm1_peer_id);
     }
+
+    #[test]
+    fn check_multiaddr_matches_peer_id() {
+        let peer_id = PeerId::random();
+        let other_peer_id = PeerId::random();
+        let mut addr: Multiaddr = "/ip4/147.75.69.143/tcp/4001"
+            .parse()
+            .expect("failed to parse multiaddr");
+
+        let addr_without_peer_id: Multiaddr = addr.clone();
+        let mut addr_with_other_peer_id = addr.clone();
+
+        addr.push(Protocol::P2p(peer_id.clone().into()));
+        addr_with_other_peer_id.push(Protocol::P2p(other_peer_id.into()));
+
+        assert!(multiaddr_matches_peer_id(&addr, &peer_id));
+        assert!(!multiaddr_matches_peer_id(
+            &addr_with_other_peer_id,
+            &peer_id
+        ));
+        assert!(multiaddr_matches_peer_id(&addr_without_peer_id, &peer_id));
+    }
 }