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 Some(cap)
263 }
264
265 #[inline]
267 pub fn ensure_matching_capability(
268 &self,
269 cap: &Capability,
270 ) -> Result<&SharedCapability, UnsupportedCapabilityError> {
271 self.find(cap).ok_or_else(|| UnsupportedCapabilityError { capability: cap.clone() })
272 }
273
274 #[inline]
276 pub const fn len(&self) -> usize {
277 self.0.len()
278 }
279
280 #[inline]
282 pub const fn is_empty(&self) -> bool {
283 self.0.is_empty()
284 }
285}
286
287#[inline]
296pub fn shared_capability_offsets(
297 local_protocols: Vec<Protocol>,
298 peer_capabilities: Vec<Capability>,
299) -> Result<Vec<SharedCapability>, P2PStreamError> {
300 let our_capabilities =
302 local_protocols.into_iter().map(Protocol::split).collect::<HashMap<_, _>>();
303
304 let mut shared_capabilities: HashMap<_, ProtoVersion> = HashMap::default();
306
307 let mut shared_capability_names = BTreeSet::new();
319
320 for peer_capability in peer_capabilities {
322 if let Some(messages) = our_capabilities.get(&peer_capability).copied() {
324 if shared_capabilities
327 .get(&peer_capability.name)
328 .is_none_or(|v| peer_capability.version > v.version)
329 {
330 shared_capabilities.insert(
331 peer_capability.name.clone(),
332 ProtoVersion { version: peer_capability.version, messages },
333 );
334 shared_capability_names.insert(peer_capability.name);
335 }
336 }
337 }
338
339 if shared_capabilities.is_empty() {
341 return Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
342 }
343
344 let mut shared_with_offsets = Vec::new();
347
348 let mut offset = MAX_RESERVED_MESSAGE_ID + 1;
352 for name in shared_capability_names {
353 let proto_version = &shared_capabilities[&name];
354 let shared_capability = SharedCapability::new(
355 &name,
356 proto_version.version as u8,
357 offset,
358 proto_version.messages,
359 )?;
360 offset += shared_capability.num_messages();
361 shared_with_offsets.push(shared_capability);
362 }
363
364 if shared_with_offsets.is_empty() {
365 return Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
366 }
367
368 Ok(shared_with_offsets)
369}
370
371#[derive(Debug, thiserror::Error)]
373pub enum SharedCapabilityError {
374 #[error(transparent)]
376 UnsupportedVersion(#[from] ParseVersionError),
377 #[error("message id offset `{0}` is reserved")]
380 ReservedMessageIdOffset(u8),
381}
382
383#[derive(Debug, thiserror::Error)]
385#[error("unsupported capability {capability}")]
386pub struct UnsupportedCapabilityError {
387 capability: Capability,
388}
389
390impl UnsupportedCapabilityError {
391 pub const fn new(capability: Capability) -> Self {
393 Self { capability }
394 }
395}
396
397#[cfg(test)]
398mod tests {
399 use super::*;
400 use crate::{Capabilities, Capability, SnapVersion};
401 use alloy_primitives::bytes::Bytes;
402 use alloy_rlp::{Decodable, Encodable};
403 use reth_eth_wire_types::RawCapabilityMessage;
404
405 #[test]
406 fn from_eth_68() {
407 let capability = SharedCapability::new("eth", 68, MAX_RESERVED_MESSAGE_ID + 1, 13).unwrap();
408
409 assert_eq!(capability.name(), "eth");
410 assert_eq!(capability.version(), 68);
411 assert_eq!(
412 capability,
413 SharedCapability::Eth {
414 version: EthVersion::Eth68,
415 offset: MAX_RESERVED_MESSAGE_ID + 1
416 }
417 );
418 }
419
420 #[test]
421 fn from_eth_67() {
422 let capability = SharedCapability::new("eth", 67, MAX_RESERVED_MESSAGE_ID + 1, 13).unwrap();
423
424 assert_eq!(capability.name(), "eth");
425 assert_eq!(capability.version(), 67);
426 assert_eq!(
427 capability,
428 SharedCapability::Eth {
429 version: EthVersion::Eth67,
430 offset: MAX_RESERVED_MESSAGE_ID + 1
431 }
432 );
433 }
434
435 #[test]
436 fn from_eth_66() {
437 let capability = SharedCapability::new("eth", 66, MAX_RESERVED_MESSAGE_ID + 1, 15).unwrap();
438
439 assert_eq!(capability.name(), "eth");
440 assert_eq!(capability.version(), 66);
441 assert_eq!(
442 capability,
443 SharedCapability::Eth {
444 version: EthVersion::Eth66,
445 offset: MAX_RESERVED_MESSAGE_ID + 1
446 }
447 );
448 }
449
450 #[test]
451 fn capabilities_supports_eth() {
452 let capabilities: Capabilities = vec![
453 Capability::new_static("eth", 66),
454 Capability::new_static("eth", 67),
455 Capability::new_static("eth", 68),
456 Capability::new_static("eth", 69),
457 Capability::new_static("eth", 70),
458 ]
459 .into();
460
461 assert!(capabilities.supports_eth());
462 assert!(capabilities.supports_eth_v66());
463 assert!(capabilities.supports_eth_v67());
464 assert!(capabilities.supports_eth_v68());
465 assert!(capabilities.supports_eth_v69());
466 assert!(capabilities.supports_eth_v70());
467 }
468
469 #[test]
470 fn test_peer_capability_version_zero() {
471 let cap = Capability::new_static("TestName", 0);
472 let local_capabilities: Vec<Protocol> =
473 vec![Protocol::new(cap.clone(), 0), EthVersion::Eth67.into(), EthVersion::Eth68.into()];
474 let peer_capabilities = vec![cap.clone()];
475
476 let shared = shared_capability_offsets(local_capabilities, peer_capabilities).unwrap();
477 assert_eq!(shared.len(), 1);
478 assert_eq!(shared[0], SharedCapability::UnknownCapability { cap, offset: 16, messages: 0 })
479 }
480
481 #[test]
482 fn test_peer_lower_capability_version() {
483 let local_capabilities: Vec<Protocol> =
484 vec![EthVersion::Eth66.into(), EthVersion::Eth67.into(), EthVersion::Eth68.into()];
485 let peer_capabilities: Vec<Capability> = vec![EthVersion::Eth66.into()];
486
487 let shared_capability =
488 shared_capability_offsets(local_capabilities, peer_capabilities).unwrap()[0].clone();
489
490 assert_eq!(
491 shared_capability,
492 SharedCapability::Eth {
493 version: EthVersion::Eth66,
494 offset: MAX_RESERVED_MESSAGE_ID + 1
495 }
496 )
497 }
498
499 #[test]
500 fn test_peer_capability_version_too_low() {
501 let local: Vec<Protocol> = vec![EthVersion::Eth67.into()];
502 let peer_capabilities: Vec<Capability> = vec![EthVersion::Eth66.into()];
503
504 let shared_capability = shared_capability_offsets(local, peer_capabilities);
505
506 assert!(matches!(
507 shared_capability,
508 Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
509 ))
510 }
511
512 #[test]
513 fn test_peer_capability_version_too_high() {
514 let local_capabilities = vec![EthVersion::Eth66.into()];
515 let peer_capabilities = vec![EthVersion::Eth67.into()];
516
517 let shared_capability = shared_capability_offsets(local_capabilities, peer_capabilities);
518
519 assert!(matches!(
520 shared_capability,
521 Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
522 ))
523 }
524
525 #[test]
526 fn test_find_by_offset() {
527 let local_capabilities = vec![EthVersion::Eth66.into()];
528 let peer_capabilities = vec![EthVersion::Eth66.into()];
529
530 let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
531
532 let shared_eth = shared.find_by_relative_offset(0).unwrap();
533 assert_eq!(shared_eth.name(), "eth");
534
535 let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 1).unwrap();
536 assert_eq!(shared_eth.name(), "eth");
537
538 assert!(shared.find_by_offset(MAX_RESERVED_MESSAGE_ID).is_none());
540 }
541
542 #[test]
543 fn test_find_by_offset_many() {
544 let cap = Capability::new_static("aaa", 1);
545 let proto = Protocol::new(cap.clone(), 5);
546 let local_capabilities = vec![proto.clone(), EthVersion::Eth66.into()];
547 let peer_capabilities = vec![cap, EthVersion::Eth66.into()];
548
549 let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
550
551 let shared_eth = shared.find_by_relative_offset(0).unwrap();
552 assert_eq!(shared_eth.name(), proto.cap.name);
553
554 let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 1).unwrap();
555 assert_eq!(shared_eth.name(), proto.cap.name);
556
557 let shared_eth = shared.find_by_relative_offset(4).unwrap();
559 assert_eq!(shared_eth.name(), proto.cap.name);
560 let shared_eth = shared.find_by_offset(MAX_RESERVED_MESSAGE_ID + 5).unwrap();
561 assert_eq!(shared_eth.name(), proto.cap.name);
562
563 let shared_eth = shared.find_by_relative_offset(1 + proto.messages()).unwrap();
565 assert_eq!(shared_eth.name(), "eth");
566 }
567
568 #[test]
569 fn relative_message_id_accounts_for_intermediate_capabilities() {
570 let intermediate_cap = Capability::new_static("foo", 1);
571 let intermediate = Protocol::new(intermediate_cap.clone(), 3);
572 let snap = Capability::snap(SnapVersion::V2);
573 let eth = Capability::eth(EthVersion::Eth69);
574 let local_capabilities =
575 vec![EthVersion::Eth69.into(), intermediate, Protocol::snap(SnapVersion::V2)];
576 let peer_capabilities = vec![eth, intermediate_cap, snap.clone()];
577
578 let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
579 let snap_id = shared.relative_message_id(&snap, 2).unwrap();
580
581 assert_eq!(snap_id, EthMessageID::message_count(EthVersion::Eth69) + 3 + 2);
582 assert_eq!(shared.capability_message_id(&snap, snap_id), Some(2));
583 }
584
585 #[test]
586 fn capability_message_id_rejects_other_capability_range() {
587 let intermediate_cap = Capability::new_static("foo", 1);
588 let intermediate = Protocol::new(intermediate_cap.clone(), 3);
589 let snap = Capability::snap(SnapVersion::V2);
590 let local_capabilities =
591 vec![EthVersion::Eth69.into(), intermediate, Protocol::snap(SnapVersion::V2)];
592 let peer_capabilities =
593 vec![Capability::eth(EthVersion::Eth69), intermediate_cap.clone(), snap.clone()];
594
595 let shared = SharedCapabilities::try_new(local_capabilities, peer_capabilities).unwrap();
596 let intermediate_id = shared.relative_message_id(&intermediate_cap, 1).unwrap();
597
598 assert_eq!(shared.capability_message_id(&snap, intermediate_id), None);
599 assert_eq!(shared.relative_message_id(&snap, SnapVersion::V2.message_count()), None);
600 }
601
602 #[test]
603 fn test_raw_capability_rlp() {
604 let msg = RawCapabilityMessage { id: 1, payload: Bytes::from(vec![0x01, 0x02, 0x03]) };
605
606 let mut encoded = Vec::new();
608 msg.encode(&mut encoded);
609
610 let decoded = RawCapabilityMessage::decode(&mut &encoded[..]).unwrap();
612
613 assert_eq!(msg, decoded);
615 }
616
617 #[test]
618 fn is_exact_eth_snap_v2_accepts_eth_and_snap() {
619 let shared = SharedCapabilities::try_new(
620 vec![EthVersion::Eth68.into(), Protocol::snap_2()],
621 vec![EthVersion::Eth68.into(), Capability::snap_2()],
622 )
623 .unwrap();
624 assert!(shared.is_exact_eth_snap_v2());
625 }
626
627 #[test]
628 fn is_exact_eth_snap_v2_rejects_eth_only() {
629 let shared = SharedCapabilities::try_new(
630 vec![EthVersion::Eth68.into()],
631 vec![EthVersion::Eth68.into()],
632 )
633 .unwrap();
634 assert!(!shared.is_exact_eth_snap_v2());
635 }
636
637 #[test]
638 fn is_exact_eth_snap_v2_rejects_eth_without_snap() {
639 let cap = Capability::new_static("les", 1);
641 let shared = SharedCapabilities::try_new(
642 vec![EthVersion::Eth68.into(), Protocol::new(cap.clone(), 5)],
643 vec![EthVersion::Eth68.into(), cap],
644 )
645 .unwrap();
646 assert!(!shared.is_exact_eth_snap_v2());
647 }
648
649 #[test]
650 fn is_exact_eth_snap_v2_rejects_eth_snap_plus_extra() {
651 let cap = Capability::new_static("les", 1);
653 let shared = SharedCapabilities::try_new(
654 vec![EthVersion::Eth68.into(), Protocol::snap_2(), Protocol::new(cap.clone(), 5)],
655 vec![EthVersion::Eth68.into(), Capability::snap_2(), cap],
656 )
657 .unwrap();
658 assert!(!shared.is_exact_eth_snap_v2());
659 }
660}