feat: add with_p2p method on multiaddr

CHANGELOG.md

@@ -1,3 +1,9 @@
+# 0.18.1 - unreleased
+
+- Add `to_p2p_terminated` on `Multiaddr`. See [PR 102].
+
+[PR 102]: https://github.com/multiformats/rust-multiaddr/pull/102
+
 # 0.18.0
 
 - Add `WebTransport` instance for `Multiaddr`. See [PR 70].

src/lib.rs

@@ -26,6 +26,8 @@ use std::{
     sync::Arc,
 };
 
+use libp2p_identity::PeerId;
+
 #[cfg(feature = "url")]
 pub use self::from_url::{from_url, from_url_lossy, FromUrlErr};
 
@@ -127,6 +129,42 @@ impl Multiaddr {
         self
     }
 
+    /// Ensures the `Multiaddr` is a `/p2p/...` address for the given peer.
+    ///
+    /// If the given address is already a `p2p` address for the given peer,
+    /// i.e. the last encapsulated protocol is `/p2p/<peer-id>`, this is a no-op.
+    ///
+    /// If the given address is already a `p2p` address for a different peer
+    /// than the one given, the peer-id is replaced with the one given.
+    ///
+    /// If the given address is not yet a `p2p` address,the `/p2p/<peer-id>`
+    /// protocol is appended to the returned address.
+    pub fn to_p2p_terminated(mut self, peer_id: PeerId) -> Self {
+        let mut last = None;
+        let mut index = 0;
+        loop {
+            let data = &self.bytes.as_ref()[index..];
+            if data.is_empty() {
+                break;
+            }
+
+            let (p, next_data) =
+                Protocol::from_bytes(data).expect("`Multiaddr` is known to be valid.");
+
+            last = Some((p, index));
+            index += data.len() - next_data.len();
+        }
+
+        match last {
+            Some((Protocol::P2p(p), _)) if p == peer_id => self,
+            Some((Protocol::P2p(_), index)) => {
+                Arc::make_mut(&mut self.bytes).truncate(index);
+                self.with(Protocol::P2p(peer_id))
+            }
+            _ => self.with(Protocol::P2p(peer_id)),
+        }
+    }
+
     /// Returns the components of this multiaddress.
     ///
     /// # Example

tests/lib.rs

@@ -661,3 +661,36 @@ fn arbitrary_impl_for_all_proto_variants() {
     let variants = core::mem::variant_count::<Protocol>() as u8;
     assert_eq!(variants, Proto::IMPL_VARIANT_COUNT);
 }
+
+#[test]
+fn to_p2p_terminated() {
+    let peer_id = "QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN"
+        .parse::<PeerId>()
+        .unwrap();
+
+    const TEST_DATA: &[(&str, &str)] = &[
+        (
+            "/ip4/127.0.0.1",
+            "/ip4/127.0.0.1/p2p/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
+        ),
+        (
+            "/ip4/127.0.0.1/p2p/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
+            "/ip4/127.0.0.1/p2p/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
+        ),
+        (
+            "/ip4/127.0.0.1/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC",
+            "/ip4/127.0.0.1/p2p/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
+        ),
+        (
+            "/ip4/127.0.0.1/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC/ip4/127.0.0.1/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC/ip4/127.0.0.1/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC",
+            "/ip4/127.0.0.1/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC/ip4/127.0.0.1/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC/ip4/127.0.0.1/p2p/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
+        ),
+    ];
+
+    for (input, expected) in TEST_DATA {
+        let input = input.parse::<Multiaddr>().unwrap();
+        let expected = expected.parse::<Multiaddr>().unwrap();
+
+        assert_eq!(expected, input.to_p2p_terminated(peer_id));
+    }
+}