address reviwer comments

Author: willscott

src/custom.rs

@@ -22,12 +22,16 @@ pub trait Transcoder: Send + Sync {
 /// A custom protocol definition.
 #[derive(Clone)]
 pub struct CustomProtocolDef {
+    /// The string identifier for the protocol (e.g. `tcp`, `http`, or `my-custom`).
     pub name: &'static str,
+    /// The unique unsigned integer code for the protocol.
     pub code: u32,
     /// The length of the binary payload.
-    /// `0` means no data. `> 0` means a fixed data length. `-1` denotes a length-prefixed protocol.
+    /// `0` means no data. `> 0` means a fixed data length. `-1` denotes a length-prefixed protocol or custom encoding.
     pub size: i32,
+    /// Whether the protocol's string representation is a file path (e.g., `/unix/tmp/socket` instead of a base-encoded payload).
     pub path: bool,
+    /// An optional transcoder used to encode and decode the protocol's data between binary and human-readable string formats.
     pub transcoder: Option<Arc<dyn Transcoder>>,
 }
 
@@ -94,7 +98,7 @@ impl Registry {
         Self::default()
     }
 
-    /// Add all built-in standard protocols to the registry
+    /// Add all built-in standard protocols to the registry.
     fn register_builtins(&mut self) {
         for &(name, code, size, path) in crate::protocol::BUILT_IN_PROTOCOLS.iter() {
             self.register(CustomProtocolDef {
@@ -143,7 +147,7 @@ impl Registry {
     }
 
     /// Iterate over the protocols in a `Multiaddr` using this registry.
-    pub fn iter<'a>(&'a self, ma: &'a Multiaddr) -> RegistryIter<'a> {
+    pub fn parse_addr<'a>(&'a self, ma: &'a Multiaddr) -> RegistryIter<'a> {
         RegistryIter {
             registry: self,
             data: ma.as_ref(),
@@ -187,56 +191,68 @@ impl Registry {
             }
         }
 
-        if let Ok((id, rest)) = id_res {
-            if let Some(def) = self.get_by_code(id) {
-                let (data, out_rest) = if def.size == 0 {
-                    (std::borrow::Cow::Borrowed(&rest[..0]), rest)
-                } else if def.size > 0 {
-                    let fixed = def.size as usize;
-                    if rest.len() < fixed {
-                        return Err(Error::DataLessThanLen);
-                    }
-                    let (d, r) = rest.split_at(fixed);
-                    (std::borrow::Cow::Borrowed(d), r)
-                } else {
-                    let (len, r) =
-                        unsigned_varint::decode::usize(rest).map_err(|_| Error::DataLessThanLen)?;
-                    if r.len() < len {
-                        return Err(Error::DataLessThanLen);
-                    }
-                    let (d, r2) = r.split_at(len);
-                    (std::borrow::Cow::Borrowed(d), r2)
-                };
+        let (id, rest) = match id_res {
+            Ok((id, rest)) => (id, rest),
+            Err(_) => return Err(Error::UnknownProtocolId(0)),
+        };
+
+        let def = match self.get_by_code(id) {
+            Some(def) => def,
+            None => {
+                // If the protocol isn't registered, we just return it as Unknown
+                // so that we can gracefully iterate over or reserialize it later.
                 return Ok((
-                    Protocol::Custom {
-                        def: Arc::new(def),
-                        data,
-                    },
-                    out_rest,
+                    Protocol::Unknown(id, std::borrow::Cow::Borrowed(rest)),
+                    [].as_ref(),
                 ));
             }
+        };
 
-            return Ok((
-                Protocol::Unknown(id, std::borrow::Cow::Borrowed(rest)),
-                [].as_ref(),
-            ));
-        }
+        // Extract the protocol data based on the registered size definition
+        let (data, out_rest) = if def.size == 0 {
+            // Protocol has no data payload expected
+            (std::borrow::Cow::Borrowed(&rest[..0]), rest)
+        } else if def.size > 0 {
+            // Protocol has a fixed-length data payload
+            let fixed = def.size as usize;
+            if rest.len() < fixed {
+                return Err(Error::DataLessThanLen);
+            }
+            let (d, r) = rest.split_at(fixed);
+            (std::borrow::Cow::Borrowed(d), r)
+        } else {
+            // Protocol size is -1, meaning it is length-prefixed.
+            // Decode the varint representing the length of the upcoming data.
+            let (len, r) =
+                unsigned_varint::decode::usize(rest).map_err(|_| Error::DataLessThanLen)?;
+            if r.len() < len {
+                return Err(Error::DataLessThanLen);
+            }
+            let (d, r2) = r.split_at(len);
+            (std::borrow::Cow::Borrowed(d), r2)
+        };
 
-        Err(Error::UnknownProtocolId(
-            id_res.map(|(i, _)| i).unwrap_or(0),
+        Ok((
+            Protocol::Custom {
+                def: Arc::new(def),
+                data,
+            },
+            out_rest,
         ))
     }
 
     /// Try parsing a single Protocol from string parts using the registry.
-    pub fn parse_protocol_from_str_parts<'a, I>(&self, iter: &mut I) -> Result<Protocol<'a>>
+    pub fn parse_protocol_from_str_parts<'a, I>(
+        &self,
+        iter: &mut std::iter::Peekable<I>,
+    ) -> Result<Protocol<'a>>
     where
         I: Iterator<Item = &'a str> + Clone,
     {
-        let mut peek_iter = iter.clone();
-        if let Some(tag) = peek_iter.next() {
-            if !self.by_name.contains_key(tag) {
-                return Err(Error::UnknownProtocolString(tag.to_string()));
-            }
+        let &tag = iter.peek().ok_or(Error::InvalidProtocolString)?;
+
+        if !self.by_name.contains_key(tag) {
+            return Err(Error::UnknownProtocolString(tag.to_string()));
         }
 
         let mut native_iter = iter.clone();
@@ -245,50 +261,39 @@ impl Registry {
             return Ok(p);
         }
 
-        let mut peek_iter = iter.clone();
-        if let Some(tag) = peek_iter.next() {
-            if let Some(def) = self.get_by_name(tag) {
-                iter.next(); // consume the tag
-                let data = if def.size == 0 {
-                    vec![]
-                } else if let Some(t) = &def.transcoder {
-                    let part = iter.next().ok_or(Error::InvalidProtocolString)?;
-                    t.string_to_bytes(part)
-                        .map_err(|_| Error::InvalidProtocolString)?
-                } else if def.path {
-                    let part = iter.next().ok_or(Error::InvalidProtocolString)?;
-                    percent_encoding::percent_decode(part.as_bytes()).collect::<Vec<u8>>()
-                } else {
-                    let part = iter.next().ok_or(Error::InvalidProtocolString)?;
-                    multibase::Base::Base58Btc
-                        .decode(part)
-                        .map_err(|_| Error::InvalidProtocolString)?
-                };
-                return Ok(Protocol::Custom {
-                    def: Arc::new(def),
-                    data: std::borrow::Cow::Owned(data),
-                });
-            }
-        }
-
-        let mut final_try = iter.clone();
-        if let Some(tag) = final_try.next() {
-            Err(Error::UnknownProtocolString(tag.to_string()))
+        let def = self.get_by_name(tag).unwrap();
+        iter.next(); // consume the tag
+        let data = if def.size == 0 {
+            vec![]
+        } else if let Some(t) = &def.transcoder {
+            let part = iter.next().ok_or(Error::InvalidProtocolString)?;
+            t.string_to_bytes(part)
+                .map_err(|_| Error::InvalidProtocolString)?
+        } else if def.path {
+            let part = iter.next().ok_or(Error::InvalidProtocolString)?;
+            percent_encoding::percent_decode(part.as_bytes()).collect::<Vec<u8>>()
         } else {
-            Err(Error::InvalidProtocolString)
-        }
+            let part = iter.next().ok_or(Error::InvalidProtocolString)?;
+            multibase::Base::Base58Btc
+                .decode(part)
+                .map_err(|_| Error::InvalidProtocolString)?
+        };
+        Ok(Protocol::Custom {
+            def: Arc::new(def),
+            data: std::borrow::Cow::Owned(data),
+        })
     }
 
     /// Parse a Multiaddr string using this registry
     pub fn try_from_str(&self, input: &str) -> Result<Multiaddr> {
         let mut addr = Multiaddr::empty();
-        let mut parts = input.split('/');
+        let mut parts = input.split('/').peekable();
 
         if Some("") != parts.next() {
             return Err(Error::InvalidMultiaddr);
         }
 
-        while parts.clone().peekable().peek().is_some() {
+        while parts.peek().is_some() {
             let p = self.parse_protocol_from_str_parts(&mut parts)?;
             addr = addr.with(p);
         }
@@ -310,7 +315,7 @@ impl Registry {
     /// Format a Multiaddr into a string using this registry
     pub fn to_string(&self, addr: &Multiaddr) -> String {
         let mut s = String::new();
-        for p in self.iter(addr) {
+        for p in self.parse_addr(addr) {
             s.push_str(&p.to_string());
         }
         s
@@ -360,7 +365,7 @@ mod tests {
         // Parse back from vec
         let parsed = registry.try_from_bytes(&vec).unwrap();
 
-        let mut iter = registry.iter(&parsed);
+        let mut iter = registry.parse_addr(&parsed);
         if let Some(Protocol::Ip4(ip)) = iter.next() {
             assert_eq!(ip, std::net::Ipv4Addr::new(127, 0, 0, 1));
         } else {
@@ -382,7 +387,7 @@ mod tests {
 
         // Assert tcp works
         let addr = registry.try_from_str("/ip4/127.0.0.1/tcp/80").unwrap();
-        let mut iter = registry.iter(&addr);
+        let mut iter = registry.parse_addr(&addr);
         assert!(matches!(iter.next(), Some(Protocol::Ip4(_))));
         assert!(matches!(iter.next(), Some(Protocol::Tcp(80))));
 
@@ -395,7 +400,7 @@ mod tests {
         // And similarly from bytes, it will now parse as Tcp since we natively fallback to standard protocols
         let vec = addr.to_vec();
         let parsed_unknown = registry.try_from_bytes(&vec).unwrap();
-        let mut parsed_iter = registry.iter(&parsed_unknown);
+        let mut parsed_iter = registry.parse_addr(&parsed_unknown);
         assert!(matches!(parsed_iter.next(), Some(Protocol::Ip4(_))));
         assert!(matches!(parsed_iter.next(), Some(Protocol::Tcp(80))));
     }
@@ -435,4 +440,45 @@ mod tests {
             "Round-trip multiaddr bytes must match the original instance precisely"
         );
     }
+
+    #[test]
+    fn test_custom_protocol_size_zero() {
+        let mut registry = Registry::new();
+        // Register a custom protocol with size 0 (no data payload expected)
+        registry.register(CustomProtocolDef::new(
+            "my-empty",
+            1000,
+            0,
+            false,
+            None::<SimpleTranscoder>,
+        ));
+
+        // It should parse without needing an additional value
+        let addr = registry.try_from_str("/ip4/127.0.0.1/my-empty").unwrap();
+
+        let vec = addr.to_vec();
+        let parsed = registry.try_from_bytes(&vec).unwrap();
+
+        let mut iter = registry.parse_addr(&parsed);
+        if let Some(Protocol::Ip4(ip)) = iter.next() {
+            assert_eq!(ip, std::net::Ipv4Addr::new(127, 0, 0, 1));
+        } else {
+            panic!("expected ip4");
+        }
+
+        if let Some(Protocol::Custom { def, data }) = iter.next() {
+            assert_eq!(def.code, 1000);
+            assert_eq!(def.name, "my-empty");
+            assert!(data.is_empty());
+        } else {
+            panic!("expected custom protocol");
+        }
+
+        // Ensure that a subsequent protocol is parsed correctly, not consumed as data
+        let addr2 = registry.try_from_str("/my-empty/tcp/80").unwrap();
+        let mut iter2 = registry.parse_addr(&addr2);
+        
+        assert!(matches!(iter2.next(), Some(Protocol::Custom { def, .. }) if def.name == "my-empty"));
+        assert!(matches!(iter2.next(), Some(Protocol::Tcp(80))));
+    }
 }