1use crate::{
4 errors::{P2PHandshakeError, P2PStreamError},
5 p2pstream::MAX_RESERVED_MESSAGE_ID,
6 protocol::{ProtoVersion, Protocol},
7 version::ParseVersionError,
8 Capability, EthMessageID, EthVersion,
9};
10use derive_more::{Deref, DerefMut};
11use std::{
12 borrow::Cow,
13 collections::{BTreeSet, HashMap},
14};
15
16#[derive(Debug, Clone, PartialEq, Eq, Hash)]
23pub enum SharedCapability {
24 Eth {
26 version: EthVersion,
28 offset: u8,
33 },
34 UnknownCapability {
36 cap: Capability,
38 offset: u8,
43 messages: u8,
46 },
47}
48
49impl SharedCapability {
50 pub(crate) fn new(
55 name: &str,
56 version: u8,
57 offset: u8,
58 messages: u8,
59 ) -> Result<Self, SharedCapabilityError> {
60 if offset <= MAX_RESERVED_MESSAGE_ID {
61 return Err(SharedCapabilityError::ReservedMessageIdOffset(offset))
62 }
63
64 match name {
65 "eth" => Ok(Self::eth(EthVersion::try_from(version)?, offset)),
66 _ => Ok(Self::UnknownCapability {
67 cap: Capability::new(name.to_string(), version as usize),
68 offset,
69 messages,
70 }),
71 }
72 }
73
74 pub(crate) const fn eth(version: EthVersion, offset: u8) -> Self {
76 Self::Eth { version, offset }
77 }
78
79 pub const fn capability(&self) -> Cow<'_, Capability> {
81 match self {
82 Self::Eth { version, .. } => Cow::Owned(Capability::eth(*version)),
83 Self::UnknownCapability { cap, .. } => Cow::Borrowed(cap),
84 }
85 }
86
87 #[inline]
89 pub fn name(&self) -> &str {
90 match self {
91 Self::Eth { .. } => "eth",
92 Self::UnknownCapability { cap, .. } => cap.name.as_ref(),
93 }
94 }
95
96 #[inline]
98 pub const fn is_eth(&self) -> bool {
99 matches!(self, Self::Eth { .. })
100 }
101
102 pub const fn version(&self) -> u8 {
104 match self {
105 Self::Eth { version, .. } => *version as u8,
106 Self::UnknownCapability { cap, .. } => cap.version as u8,
107 }
108 }
109
110 pub const fn eth_version(&self) -> Option<EthVersion> {
112 match self {
113 Self::Eth { version, .. } => Some(*version),
114 _ => None,
115 }
116 }
117
118 pub const fn message_id_offset(&self) -> u8 {
123 match self {
124 Self::Eth { offset, .. } | Self::UnknownCapability { offset, .. } => *offset,
125 }
126 }
127
128 pub const fn relative_message_id_offset(&self) -> u8 {
131 self.message_id_offset() - MAX_RESERVED_MESSAGE_ID - 1
132 }
133
134 pub const fn num_messages(&self) -> u8 {
136 match self {
137 Self::Eth { version, .. } => EthMessageID::message_count(*version),
138 Self::UnknownCapability { messages, .. } => *messages,
139 }
140 }
141}
142
143#[derive(Debug, Clone, Deref, DerefMut, PartialEq, Eq)]
147pub struct SharedCapabilities(Vec<SharedCapability>);
148
149impl SharedCapabilities {
150 #[inline]
152 pub fn try_new(
153 local_protocols: Vec<Protocol>,
154 peer_capabilities: Vec<Capability>,
155 ) -> Result<Self, P2PStreamError> {
156 shared_capability_offsets(local_protocols, peer_capabilities).map(Self)
157 }
158
159 #[inline]
161 pub fn iter_caps(&self) -> impl Iterator<Item = &SharedCapability> {
162 self.0.iter()
163 }
164
165 #[inline]
167 pub fn eth(&self) -> Result<&SharedCapability, P2PStreamError> {
168 self.iter_caps().find(|c| c.is_eth()).ok_or(P2PStreamError::CapabilityNotShared)
169 }
170
171 #[inline]
173 pub fn eth_version(&self) -> Result<EthVersion, P2PStreamError> {
174 self.iter_caps()
175 .find_map(SharedCapability::eth_version)
176 .ok_or(P2PStreamError::CapabilityNotShared)
177 }
178
179 #[inline]
182 pub fn is_exact_eth_snap_v2(&self) -> bool {
183 self.len() == 2 && self.ensure_matching_capability(&Capability::snap_2()).is_ok()
184 }
185
186 #[inline]
188 pub fn contains(&self, cap: &Capability) -> bool {
189 self.find(cap).is_some()
190 }
191
192 #[inline]
194 pub fn find(&self, cap: &Capability) -> Option<&SharedCapability> {
195 self.0.iter().find(|c| c.version() == cap.version as u8 && c.name() == cap.name)
196 }
197
198 #[inline]
204 pub fn relative_message_id(&self, cap: &Capability, message_id: u8) -> Option<u8> {
205 let shared = self.find(cap)?;
206 if message_id >= shared.num_messages() {
207 return None
208 }
209
210 shared.relative_message_id_offset().checked_add(message_id)
211 }
212
213 #[inline]
218 pub fn capability_message_id(&self, cap: &Capability, relative_message_id: u8) -> Option<u8> {
219 let shared = self.find(cap)?;
220 let start = shared.relative_message_id_offset();
221 let end = start.checked_add(shared.num_messages())?;
222
223 (start..end).contains(&relative_message_id).then(|| relative_message_id - start)
224 }
225
226 #[inline]
235 pub fn find_by_relative_offset(&self, offset: u8) -> Option<&SharedCapability> {
236 self.find_by_offset(offset.saturating_add(MAX_RESERVED_MESSAGE_ID + 1))
237 }
238
239 #[inline]
247 pub fn find_by_offset(&self, offset: u8) -> Option<&SharedCapability> {
248 let mut iter = self.0.iter();
249 let mut cap = iter.next()?;
250 if offset < cap.message_id_offset() {
251 return None
253 }
254
255 for next in iter {
256 if offset < next.message_id_offset() {
257 return Some(cap)
258 }
259 cap = next
260 }
261
262 let end = u16::from(cap.message_id_offset()) + u16::from(cap.num_messages());
263 (u16::from(offset) < end).then_some(cap)
264 }
265
266 #[inline]
268 pub fn ensure_matching_capability(
269 &self,
270 cap: &Capability,
271 ) -> Result<&SharedCapability, UnsupportedCapabilityError> {
272 self.find(cap).ok_or_else(|| UnsupportedCapabilityError { capability: cap.clone() })
273 }
274
275 #[inline]
277 pub const fn len(&self) -> usize {
278 self.0.len()
279 }
280
281 #[inline]
283 pub const fn is_empty(&self) -> bool {
284 self.0.is_empty()
285 }
286}
287
288#[inline]
297pub fn shared_capability_offsets(
298 local_protocols: Vec<Protocol>,
299 peer_capabilities: Vec<Capability>,
300) -> Result<Vec<SharedCapability>, P2PStreamError> {
301 let our_capabilities =
303 local_protocols.into_iter().map(Protocol::split).collect::<HashMap<_, _>>();
304
305 let mut shared_capabilities: HashMap<_, ProtoVersion> = HashMap::default();
307
308 let mut shared_capability_names = BTreeSet::new();
320
321 for peer_capability in peer_capabilities {
323 if let Some(messages) = our_capabilities.get(&peer_capability).copied() {
325 if shared_capabilities
328 .get(&peer_capability.name)
329 .is_none_or(|v| peer_capability.version > v.version)
330 {
331 shared_capabilities.insert(
332 peer_capability.name.clone(),
333 ProtoVersion { version: peer_capability.version, messages },
334 );
335 shared_capability_names.insert(peer_capability.name);
336 }
337 }
338 }
339
340 if shared_capabilities.is_empty() {
342 return Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
343 }
344
345 let mut shared_with_offsets = Vec::new();
348
349 let mut offset = MAX_RESERVED_MESSAGE_ID + 1;
353 for name in shared_capability_names {
354 let proto_version = &shared_capabilities[&name];
355 let shared_capability = SharedCapability::new(
356 &name,
357 proto_version.version as u8,
358 offset,
359 proto_version.messages,
360 )?;
361 offset += shared_capability.num_messages();
362 shared_with_offsets.push(shared_capability);
363 }
364
365 if shared_with_offsets.is_empty() {
366 return Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
367 }
368
369 Ok(shared_with_offsets)
370}
371
372#[derive(Debug, thiserror::Error)]
374pub enum SharedCapabilityError {
375 #[error(transparent)]
377 UnsupportedVersion(#[from] ParseVersionError),
378 #[error("message id offset `{0}` is reserved")]
381 ReservedMessageIdOffset(u8),
382}
383
384#[derive(Debug, thiserror::Error)]
386#[error("unsupported capability {capability}")]
387pub struct UnsupportedCapabilityError {
388 capability: Capability,
389}
390
391impl UnsupportedCapabilityError {
392 pub const fn new(capability: Capability) -> Self {
394 Self { capability }
395 }
396}
397
398#[cfg(test)]
399mod tests {
400 use super::*;
401 use crate::{Capabilities, Capability, SnapVersion};
402 use alloy_primitives::bytes::Bytes;
403 use alloy_rlp::{Decodable, Encodable};
404 use reth_eth_wire_types::RawCapabilityMessage;
405
406 #[test]
407 fn from_eth_68() {
408 let capability = SharedCapability::new("eth", 68, MAX_RESERVED_MESSAGE_ID + 1, 13).unwrap();
409
410 assert_eq!(capability.name(), "eth");
411 assert_eq!(capability.version(), 68);
412 assert_eq!(
413 capability,
414 SharedCapability::Eth {
415 version: EthVersion::Eth68,
416 offset: MAX_RESERVED_MESSAGE_ID + 1
417 }
418 );
419 }
420
421 #[test]
422 fn from_eth_67() {
423 let capability = SharedCapability::new("eth", 67, MAX_RESERVED_MESSAGE_ID + 1, 13).unwrap();
424
425 assert_eq!(capability.name(), "eth");
426 assert_eq!(capability.version(), 67);
427 assert_eq!(
428 capability,
429 SharedCapability::Eth {
430 version: EthVersion::Eth67,
431 offset: MAX_RESERVED_MESSAGE_ID + 1
432 }
433 );
434 }
435
436 #[test]
437 fn from_eth_66() {
438 let capability = SharedCapability::new("eth", 66, MAX_RESERVED_MESSAGE_ID + 1, 15).unwrap();
439
440 assert_eq!(capability.name(), "eth");
441 assert_eq!(capability.version(), 66);
442 assert_eq!(
443 capability,
444 SharedCapability::Eth {
445 version: EthVersion::Eth66,
446 offset: MAX_RESERVED_MESSAGE_ID + 1
447 }
448 );
449 }
450
451 #[test]
452 fn capabilities_supports_eth() {
453 let capabilities: Capabilities = vec![
454 Capability::new_static("eth", 66),
455 Capability::new_static("eth", 67),
456 Capability::new_static("eth", 68),
457 Capability::new_static("eth", 69),
458 Capability::new_static("eth", 70),
459 ]
460 .into();
461
462 assert!(capabilities.supports_eth());
463 assert!(capabilities.supports_eth_v66());
464 assert!(capabilities.supports_eth_v67());
465 assert!(capabilities.supports_eth_v68());
466 assert!(capabilities.supports_eth_v69());
467 assert!(capabilities.supports_eth_v70());
468 }
469
470 #[test]
471 fn lookup_rejects_message_ids_past_last_capability() {
472 let cap = Capability::new_static("test", 1);
473 let shared =
474 SharedCapabilities::try_new(vec![Protocol::new(cap.clone(), 1)], vec![cap.clone()])
475 .unwrap();
476 let offset = shared.find(&cap).unwrap().message_id_offset();
477
478 assert_eq!(shared.find_by_offset(offset).unwrap().capability().as_ref(), &cap);
479 assert!(shared.find_by_offset(offset + 1).is_none());
480 }
481
482 #[test]
483 fn test_peer_capability_version_zero() {
484 let cap = Capability::new_static("TestName", 0);
485 let local_capabilities: Vec<Protocol> =
486 vec![Protocol::new(cap.clone(), 0), EthVersion::Eth67.into(), EthVersion::Eth68.into()];
487 let peer_capabilities = vec![cap.clone()];
488
489 let shared = shared_capability_offsets(local_capabilities, peer_capabilities).unwrap();
490 assert_eq!(shared.len(), 1);
491 assert_eq!(shared[0], SharedCapability::UnknownCapability { cap, offset: 16, messages: 0 })
492 }
493
494 #[test]
495 fn test_peer_lower_capability_version() {
496 let local_capabilities: Vec<Protocol> =
497 vec![EthVersion::Eth66.into(), EthVersion::Eth67.into(), EthVersion::Eth68.into()];
498 let peer_capabilities: Vec<Capability> = vec![EthVersion::Eth66.into()];
499
500 let shared_capability =
501 shared_capability_offsets(local_capabilities, peer_capabilities).unwrap()[0].clone();
502
503 assert_eq!(
504 shared_capability,
505 SharedCapability::Eth {
506 version: EthVersion::Eth66,
507 offset: MAX_RESERVED_MESSAGE_ID + 1
508 }
509 )
510 }
511
512 #[test]
513 fn test_peer_capability_version_too_low() {
514 let local: Vec<Protocol> = vec![EthVersion::Eth67.into()];
515 let peer_capabilities: Vec<Capability> = vec![EthVersion::Eth66.into()];
516
517 let shared_capability = shared_capability_offsets(local, peer_capabilities);
518
519 assert!(matches!(
520 shared_capability,
521 Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
522 ))
523 }
524
525 #[test]
526 fn test_peer_capability_version_too_high() {
527 let local_capabilities = vec![EthVersion::Eth66.into()];
528 let peer_capabilities = vec![EthVersion::Eth67.into()];
529
530 let shared_capability = shared_capability_offsets(local_capabilities, peer_capabilities);
531
532 assert!(matches!(
533 shared_capability,
534 Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
535 ))
536 }
537
538 #[test]
539 fn test_find_by_offset() {
540 let local_capabilities = vec![EthVersion::Eth66.into()];
541 let peer_capabilities = vec![EthVersion::Eth66.into()];
542
543 let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
544
545 let shared_eth = shared.find_by_relative_offset(0).unwrap();
546 assert_eq!(shared_eth.name(), "eth");
547
548 let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 1).unwrap();
549 assert_eq!(shared_eth.name(), "eth");
550
551 assert!(shared.find_by_offset(MAX_RESERVED_MESSAGE_ID).is_none());
553 }
554
555 #[test]
556 fn test_find_by_offset_many() {
557 let cap = Capability::new_static("aaa", 1);
558 let proto = Protocol::new(cap.clone(), 5);
559 let local_capabilities = vec![proto.clone(), EthVersion::Eth66.into()];
560 let peer_capabilities = vec![cap, EthVersion::Eth66.into()];
561
562 let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
563
564 let shared_eth = shared.find_by_relative_offset(0).unwrap();
565 assert_eq!(shared_eth.name(), proto.cap.name);
566
567 let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 1).unwrap();
568 assert_eq!(shared_eth.name(), proto.cap.name);
569
570 let shared_eth = shared.find_by_relative_offset(4).unwrap();
572 assert_eq!(shared_eth.name(), proto.cap.name);
573 let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 5).unwrap();
574 assert_eq!(shared_eth.name(), proto.cap.name);
575
576 let shared_eth = shared.find_by_relative_offset(1 + proto.messages()).unwrap();
578 assert_eq!(shared_eth.name(), "eth");
579 }
580
581 #[test]
582 fn relative_message_id_accounts_for_intermediate_capabilities() {
583 let intermediate_cap = Capability::new_static("foo", 1);
584 let intermediate = Protocol::new(intermediate_cap.clone(), 3);
585 let snap = Capability::snap(SnapVersion::V2);
586 let eth = Capability::eth(EthVersion::Eth69);
587 let local_capabilities =
588 vec![EthVersion::Eth69.into(), intermediate, Protocol::snap(SnapVersion::V2)];
589 let peer_capabilities = vec![eth, intermediate_cap, snap.clone()];
590
591 let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
592 let snap_id = shared.relative_message_id(&snap, 2).unwrap();
593
594 assert_eq!(snap_id, EthMessageID::message_count(EthVersion::Eth69) + 3 + 2);
595 assert_eq!(shared.capability_message_id(&snap, snap_id), Some(2));
596 }
597
598 #[test]
599 fn capability_message_id_rejects_other_capability_range() {
600 let intermediate_cap = Capability::new_static("foo", 1);
601 let intermediate = Protocol::new(intermediate_cap.clone(), 3);
602 let snap = Capability::snap(SnapVersion::V2);
603 let local_capabilities =
604 vec![EthVersion::Eth69.into(), intermediate, Protocol::snap(SnapVersion::V2)];
605 let peer_capabilities =
606 vec![Capability::eth(EthVersion::Eth69), intermediate_cap.clone(), snap.clone()];
607
608 let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
609 let intermediate_id = shared.relative_message_id(&intermediate_cap, 1).unwrap();
610
611 assert_eq!(shared.capability_message_id(&snap, intermediate_id), None);
612 assert_eq!(shared.relative_message_id(&snap, SnapVersion::V2.message_count()), None);
613 }
614
615 #[test]
616 fn test_raw_capability_rlp() {
617 let msg = RawCapabilityMessage { id: 1, payload: Bytes::from(vec![0x01, 0x02, 0x03]) };
618
619 let mut encoded = Vec::new();
621 msg.encode(&mut encoded);
622
623 let decoded = RawCapabilityMessage::decode(&mut &encoded[..]).unwrap();
625
626 assert_eq!(msg, decoded);
628 }
629
630 #[test]
631 fn is_exact_eth_snap_v2_accepts_eth_and_snap() {
632 let shared = SharedCapabilities::try_new(
633 vec![EthVersion::Eth68.into(), Protocol::snap_2()],
634 vec![EthVersion::Eth68.into(), Capability::snap_2()],
635 )
636 .unwrap();
637 assert!(shared.is_exact_eth_snap_v2());
638 }
639
640 #[test]
641 fn is_exact_eth_snap_v2_rejects_eth_only() {
642 let shared = SharedCapabilities::try_new(
643 vec![EthVersion::Eth68.into()],
644 vec![EthVersion::Eth68.into()],
645 )
646 .unwrap();
647 assert!(!shared.is_exact_eth_snap_v2());
648 }
649
650 #[test]
651 fn is_exact_eth_snap_v2_rejects_eth_without_snap() {
652 let cap = Capability::new_static("les", 1);
654 let shared = SharedCapabilities::try_new(
655 vec![EthVersion::Eth68.into(), Protocol::new(cap.clone(), 5)],
656 vec![EthVersion::Eth68.into(), cap],
657 )
658 .unwrap();
659 assert!(!shared.is_exact_eth_snap_v2());
660 }
661
662 #[test]
663 fn is_exact_eth_snap_v2_rejects_eth_snap_plus_extra() {
664 let cap = Capability::new_static("les", 1);
666 let shared = SharedCapabilities::try_new(
667 vec![EthVersion::Eth68.into(), Protocol::snap_2(), Protocol::new(cap.clone(), 5)],
668 vec![EthVersion::Eth68.into(), Capability::snap_2(), cap],
669 )
670 .unwrap();
671 assert!(!shared.is_exact_eth_snap_v2());
672 }
673}