From 911b18a51fafdee44a445dc491192f85430fbbbf Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Fri, 19 Aug 2022 06:52:05 +0400
Subject: [PATCH] src/protocol.rs: Add `/webrtc` and `/certhash` (#59)

---
 CHANGELOG.md    |  4 ++++
 src/protocol.rs | 31 +++++++++++++++++++++++++++++++
 tests/lib.rs    | 21 +++++++++++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 61f4458..d522b87 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,11 +1,15 @@
 # 0.15.0 [unreleased]
 
+- Add `WebRTC` instance for `Multiaddr`. See [PR 59].
+- Add `Certhash` instance for `Multiaddr`. See [PR 59].
+
 - Add support for Noise protocol. See [PR 53].
 
 - Use `multibase` instead of `bs58` for base58 encoding. See [PR 56].
 
 [PR 53]: https://github.com/multiformats/rust-multiaddr/pull/53
 [PR 56]: https://github.com/multiformats/rust-multiaddr/pull/56
+[PR 59]: https://github.com/multiformats/rust-multiaddr/pull/59
 
 # 0.14.0 [2022-02-02]
 
diff --git a/src/protocol.rs b/src/protocol.rs
index 67d3cf6..59b9537 100644
--- a/src/protocol.rs
+++ b/src/protocol.rs
@@ -28,6 +28,8 @@ const IP4: u32 = 4;
 const IP6: u32 = 41;
 const P2P_WEBRTC_DIRECT: u32 = 276;
 const P2P_WEBRTC_STAR: u32 = 275;
+const WEBRTC: u32 = 280;
+const CERTHASH: u32 = 466;
 const P2P_WEBSOCKET_STAR: u32 = 479;
 const MEMORY: u32 = 777;
 const ONION: u32 = 444;
@@ -80,6 +82,8 @@ pub enum Protocol<'a> {
     Ip6(Ipv6Addr),
     P2pWebRtcDirect,
     P2pWebRtcStar,
+    WebRTC,
+    Certhash(Multihash),
     P2pWebSocketStar,
     /// Contains the "port" to contact. Similar to TCP or UDP, 0 means "assign me a port".
     Memory(u64),
@@ -192,6 +196,12 @@ impl<'a> Protocol<'a> {
             }
             "p2p-websocket-star" => Ok(Protocol::P2pWebSocketStar),
             "p2p-webrtc-star" => Ok(Protocol::P2pWebRtcStar),
+            "webrtc" => Ok(Protocol::WebRTC),
+            "certhash" => {
+                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
+                let (_base, decoded) = multibase::decode(s)?;
+                Ok(Protocol::Certhash(Multihash::from_bytes(&decoded)?))
+            }
             "p2p-webrtc-direct" => Ok(Protocol::P2pWebRtcDirect),
             "p2p-circuit" => Ok(Protocol::P2pCircuit),
             "memory" => {
@@ -268,6 +278,12 @@ impl<'a> Protocol<'a> {
             }
             P2P_WEBRTC_DIRECT => Ok((Protocol::P2pWebRtcDirect, input)),
             P2P_WEBRTC_STAR => Ok((Protocol::P2pWebRtcStar, input)),
+            WEBRTC => Ok((Protocol::WebRTC, input)),
+            CERTHASH => {
+                let (n, input) = decode::usize(input)?;
+                let (data, rest) = split_at(n, input)?;
+                Ok((Protocol::Certhash(Multihash::from_bytes(data)?), rest))
+            }
             P2P_WEBSOCKET_STAR => Ok((Protocol::P2pWebSocketStar, input)),
             MEMORY => {
                 let (data, rest) = split_at(8, input)?;
@@ -441,6 +457,13 @@ impl<'a> Protocol<'a> {
             }
             Protocol::P2pWebSocketStar => w.write_all(encode::u32(P2P_WEBSOCKET_STAR, &mut buf))?,
             Protocol::P2pWebRtcStar => w.write_all(encode::u32(P2P_WEBRTC_STAR, &mut buf))?,
+            Protocol::WebRTC => w.write_all(encode::u32(WEBRTC, &mut buf))?,
+            Protocol::Certhash(hash) => {
+                w.write_all(encode::u32(CERTHASH, &mut buf))?;
+                let bytes = hash.to_bytes();
+                w.write_all(encode::usize(bytes.len(), &mut encode::usize_buffer()))?;
+                w.write_all(&bytes)?
+            }
             Protocol::P2pWebRtcDirect => w.write_all(encode::u32(P2P_WEBRTC_DIRECT, &mut buf))?,
             Protocol::P2pCircuit => w.write_all(encode::u32(P2P_CIRCUIT, &mut buf))?,
             Protocol::Memory(port) => {
@@ -466,6 +489,8 @@ impl<'a> Protocol<'a> {
             Ip6(a) => Ip6(a),
             P2pWebRtcDirect => P2pWebRtcDirect,
             P2pWebRtcStar => P2pWebRtcStar,
+            WebRTC => WebRTC,
+            Certhash(hash) => Certhash(hash),
             P2pWebSocketStar => P2pWebSocketStar,
             Memory(a) => Memory(a),
             Onion(addr, port) => Onion(Cow::Owned(addr.into_owned()), port),
@@ -502,6 +527,12 @@ impl<'a> fmt::Display for Protocol<'a> {
             Ip6(addr) => write!(f, "/ip6/{}", addr),
             P2pWebRtcDirect => f.write_str("/p2p-webrtc-direct"),
             P2pWebRtcStar => f.write_str("/p2p-webrtc-star"),
+            WebRTC => f.write_str("/webrtc"),
+            Certhash(hash) => write!(
+                f,
+                "/certhash/{}",
+                multibase::encode(multibase::Base::Base64Url, hash.to_bytes())
+            ),
             P2pWebSocketStar => f.write_str("/p2p-websocket-star"),
             Memory(port) => write!(f, "/memory/{}", port),
             Onion(addr, port) => {
diff --git a/tests/lib.rs b/tests/lib.rs
index 2631b65..f31a955 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -338,6 +338,25 @@ fn construct_success() {
         "047F00000106007FC603",
         vec![Ip4(local), Tcp(127), Noise],
     );
+
+    ma_valid(
+        "/ip4/127.0.0.1/udp/1234/webrtc",
+        "047F000001910204D29802",
+        vec![Ip4(local), Udp(1234), WebRTC],
+    );
+
+    let (_base, decoded) =
+        multibase::decode("uEiDDq4_xNyDorZBH3TlGazyJdOWSwvo4PUo5YHFMrvDE8g").unwrap();
+    ma_valid(
+        "/ip4/127.0.0.1/udp/1234/webrtc/certhash/uEiDDq4_xNyDorZBH3TlGazyJdOWSwvo4PUo5YHFMrvDE8g",
+        "047F000001910204D29802D203221220C3AB8FF13720E8AD9047DD39466B3C8974E592C2FA383D4A3960714CAEF0C4F2",
+        vec![
+            Ip4(local),
+            Udp(1234),
+            WebRTC,
+            Certhash(Multihash::from_bytes(&decoded).unwrap()),
+        ],
+    );
 }
 
 #[test]
@@ -374,6 +393,8 @@ fn construct_fail() {
         "/ip4/127.0.0.1/p2p",
         "/ip4/127.0.0.1/p2p/tcp",
         "/p2p-circuit/50",
+        "/ip4/127.0.0.1/udp/1234/webrtc/certhash",
+        "/ip4/127.0.0.1/udp/1234/webrtc/certhash/b2uaraocy6yrdblb4sfptaddgimjmmp", // 1 character missing from certhash
     ];
 
     for address in &addresses {