protocols/identify Check multiaddr has valid peer id component prior to caching

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`.

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]
 

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));
+    }
 }