diff --git a/crates/net/eth-wire/src/capability.rs b/crates/net/eth-wire/src/capability.rs index 9270122a17a5..3a72504e93b5 100644 --- a/crates/net/eth-wire/src/capability.rs +++ b/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"); + } }