feat: add find by offset functions (#5541)

crates/net/eth-wire/src/capability.rs

@@ -374,6 +374,38 @@ impl SharedCapabilities {
         self.0.iter().find(|c| c.version() == cap.version as u8 && c.name() == cap.name)
     }
 
+    /// Returns the matching shared capability for the given capability offset.
+    ///
+    /// `offset` is the multiplexed message id offset of the capability relative to
+    /// [`MAX_RESERVED_MESSAGE_ID`].
+    #[inline]
+    pub fn find_by_relative_offset(&self, offset: u8) -> Option<&SharedCapability> {
+        self.find_by_offset(offset.saturating_add(MAX_RESERVED_MESSAGE_ID))
+    }
+
+    /// Returns the matching shared capability for the given capability offset.
+    ///
+    /// `offset` is the multiplexed message id offset of the capability that includes the reserved
+    /// message id space.
+    #[inline]
+    pub fn find_by_offset(&self, offset: u8) -> Option<&SharedCapability> {
+        let mut iter = self.0.iter();
+        let mut cap = iter.next()?;
+        if offset < cap.message_id_offset() {
+            // reserved message id space
+            return None
+        }
+
+        for next in iter {
+            if offset < next.message_id_offset() {
+                return Some(cap)
+            }
+            cap = next
+        }
+
+        Some(cap)
+    }
+
     /// Returns the shared capability for the given capability or an error if it's not compatible.
     #[inline]
     pub fn ensure_matching_capability(
@@ -597,4 +629,46 @@ mod tests {
             Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
         ))
     }
+
+    #[test]
+    fn test_find_by_offset() {
+        let local_capabilities = vec![EthVersion::Eth66.into()];
+        let peer_capabilities = vec![EthVersion::Eth66.into()];
+
+        let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
+
+        assert!(shared.find_by_relative_offset(0).is_none());
+        let shared_eth = shared.find_by_relative_offset(1).unwrap();
+        assert_eq!(shared_eth.name(), "eth");
+
+        let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 1).unwrap();
+        assert_eq!(shared_eth.name(), "eth");
+    }
+
+    #[test]
+    fn test_find_by_offset_many() {
+        let cap = Capability::new_static("aaa", 1);
+        let proto = Protocol::new(cap.clone(), 5);
+        let local_capabilities = vec![proto.clone(), EthVersion::Eth66.into()];
+        let peer_capabilities = vec![cap, EthVersion::Eth66.into()];
+
+        let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
+
+        assert!(shared.find_by_relative_offset(0).is_none());
+        let shared_eth = shared.find_by_relative_offset(1).unwrap();
+        assert_eq!(shared_eth.name(), proto.cap.name);
+
+        let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 1).unwrap();
+        assert_eq!(shared_eth.name(), proto.cap.name);
+
+        // the 5th shared message is the last message of the aaa capability
+        let shared_eth = shared.find_by_relative_offset(5).unwrap();
+        assert_eq!(shared_eth.name(), proto.cap.name);
+        let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 5).unwrap();
+        assert_eq!(shared_eth.name(), proto.cap.name);
+
+        // the 6th shared message is the first message of the eth capability
+        let shared_eth = shared.find_by_relative_offset(1 + proto.messages).unwrap();
+        assert_eq!(shared_eth.name(), "eth");
+    }
 }