1#![doc(
20 html_logo_url = "https://raw.githubusercontent.com/paradigmxyz/reth/main/assets/reth-docs.png",
21 html_favicon_url = "https://avatars0.githubusercontent.com/u/97369466?s=256",
22 issue_tracker_base_url = "https://github.com/paradigmxyz/reth/issues/"
23)]
24#![cfg_attr(not(test), warn(unused_crate_dependencies))]
25#![cfg_attr(docsrs, feature(doc_cfg))]
26
27use crate::{
28 error::{DecodePacketError, Discv4Error},
29 proto::{FindNode, Message, Neighbours, Packet, Ping, Pong},
30};
31use alloy_primitives::{bytes::Bytes, hex, B256};
32use discv5::{
33 kbucket,
34 kbucket::{
35 BucketInsertResult, Distance, Entry as BucketEntry, InsertResult, KBucketsTable,
36 NodeStatus, MAX_NODES_PER_BUCKET,
37 },
38 ConnectionDirection, ConnectionState,
39};
40use enr::Enr;
41use itertools::Itertools;
42use parking_lot::Mutex;
43use proto::{EnrRequest, EnrResponse};
44use reth_ethereum_forks::ForkId;
45use reth_network_peers::{pk2id, PeerId};
46use secp256k1::SecretKey;
47use std::{
48 cell::RefCell,
49 collections::{btree_map, hash_map::Entry, BTreeMap, HashMap, VecDeque},
50 fmt, io,
51 net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
52 pin::Pin,
53 rc::Rc,
54 sync::Arc,
55 task::{ready, Context, Poll},
56 time::{Duration, Instant, SystemTime, UNIX_EPOCH},
57};
58use tokio::{
59 net::UdpSocket,
60 sync::{mpsc, mpsc::error::TrySendError, oneshot, oneshot::Sender as OneshotSender},
61 task::{JoinHandle, JoinSet},
62 time::Interval,
63};
64use tokio_stream::{wrappers::ReceiverStream, Stream, StreamExt};
65use tracing::{debug, trace};
66
67pub mod error;
68pub mod proto;
69
70mod config;
71pub use config::{Discv4Config, Discv4ConfigBuilder};
72
73mod node;
74use node::{kad_key, NodeKey};
75
76mod table;
77
78pub use reth_network_peers::NodeRecord;
80
81#[cfg(any(test, feature = "test-utils"))]
82pub mod test_utils;
83
84use crate::table::PongTable;
85use reth_net_nat::ResolveNatInterval;
86pub use reth_net_nat::{external_ip, NatResolver};
88
89pub const DEFAULT_DISCOVERY_ADDR: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED);
93
94pub const DEFAULT_DISCOVERY_PORT: u16 = 30303;
98
99pub const DEFAULT_DISCOVERY_ADDRESS: SocketAddr =
103 SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, DEFAULT_DISCOVERY_PORT));
104
105const MAX_PACKET_SIZE: usize = 1280;
107
108const MIN_PACKET_SIZE: usize = 32 + 65 + 1;
110
111const ALPHA: usize = 3;
113
114const MAX_NODES_PING: usize = 2 * MAX_NODES_PER_BUCKET;
119
120const MAX_QUEUED_PINGS: usize = 2 * MAX_NODES_PER_BUCKET;
128
129const SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS: usize = (MAX_PACKET_SIZE - 109) / 91;
135
136const ENDPOINT_PROOF_EXPIRATION: Duration = Duration::from_secs(24 * 60 * 60);
140
141const EXPIRE_DURATION: Duration = Duration::from_secs(60 * 60);
143
144const UDP_MESSAGE_POLL_LOOP_BUDGET: i32 = 4;
149
150type EgressSender = mpsc::Sender<(Bytes, SocketAddr)>;
151type EgressReceiver = mpsc::Receiver<(Bytes, SocketAddr)>;
152
153pub(crate) type IngressSender = mpsc::Sender<IngressEvent>;
154pub(crate) type IngressReceiver = mpsc::Receiver<IngressEvent>;
155
156type NodeRecordSender = OneshotSender<Vec<NodeRecord>>;
157
158#[derive(Debug, Clone)]
165pub struct Discv4 {
166 local_addr: SocketAddr,
168 to_service: mpsc::UnboundedSender<Discv4Command>,
170 node_record: Arc<Mutex<NodeRecord>>,
174}
175
176impl Discv4 {
177 pub async fn spawn(
181 local_address: SocketAddr,
182 local_enr: NodeRecord,
183 secret_key: SecretKey,
184 config: Discv4Config,
185 ) -> io::Result<Self> {
186 let (discv4, service) = Self::bind(local_address, local_enr, secret_key, config).await?;
187
188 service.spawn();
189
190 Ok(discv4)
191 }
192
193 #[cfg(feature = "test-utils")]
197 pub fn noop() -> Self {
198 let (to_service, _rx) = mpsc::unbounded_channel();
199 let local_addr =
200 (IpAddr::from(std::net::Ipv4Addr::UNSPECIFIED), DEFAULT_DISCOVERY_PORT).into();
201 Self {
202 local_addr,
203 to_service,
204 node_record: Arc::new(Mutex::new(NodeRecord::new(
205 "127.0.0.1:3030".parse().unwrap(),
206 PeerId::random(),
207 ))),
208 }
209 }
210
211 pub async fn bind(
243 local_address: SocketAddr,
244 local_node_record: NodeRecord,
245 secret_key: SecretKey,
246 config: Discv4Config,
247 ) -> io::Result<(Self, Discv4Service)> {
248 let socket = Arc::new(UdpSocket::bind(local_address).await?);
249 trace!(target: "discv4", local_addr=?socket.local_addr(), "opened UDP socket");
250 let (tx, rx) = mpsc::channel(config.udp_ingress_message_buffer);
251
252 Self::bind_with_socket(socket, Some(tx), rx, local_node_record, secret_key, config)
253 }
254
255 pub fn bind_shared(
259 socket: Arc<UdpSocket>,
260 local_node_record: NodeRecord,
261 secret_key: SecretKey,
262 config: Discv4Config,
263 ) -> io::Result<(Self, Discv4Service, IngressHandler)> {
264 let (tx, rx) = mpsc::channel(config.udp_ingress_message_buffer);
265 let local_id = local_node_record.id;
266 let (discv4, service) =
267 Self::bind_with_socket(socket, None, rx, local_node_record, secret_key, config)?;
268
269 let handler = IngressHandler::new(tx, local_id);
270
271 Ok((discv4, service, handler))
272 }
273
274 fn bind_with_socket(
275 socket: Arc<UdpSocket>,
276 ingress_tx: Option<IngressSender>,
277 ingress_rx: IngressReceiver,
278 mut local_node_record: NodeRecord,
279 secret_key: SecretKey,
280 config: Discv4Config,
281 ) -> io::Result<(Self, Discv4Service)> {
282 let local_addr = socket.local_addr()?;
283 local_node_record.udp_port = local_addr.port();
284
285 let mut service = Discv4Service::new(
286 socket,
287 ingress_tx,
288 ingress_rx,
289 local_addr,
290 local_node_record,
291 secret_key,
292 config,
293 );
294
295 service.resolve_external_ip();
297
298 let discv4 = service.handle();
299 Ok((discv4, service))
300 }
301
302 pub const fn local_addr(&self) -> SocketAddr {
304 self.local_addr
305 }
306
307 pub fn node_record(&self) -> NodeRecord {
311 *self.node_record.lock()
312 }
313
314 pub fn external_ip(&self) -> IpAddr {
316 self.node_record.lock().address
317 }
318
319 pub fn set_lookup_interval(&self, duration: Duration) {
321 self.send_to_service(Discv4Command::SetLookupInterval(duration))
322 }
323
324 pub async fn lookup_self(&self) -> Result<Vec<NodeRecord>, Discv4Error> {
339 self.lookup_node(None).await
340 }
341
342 pub async fn lookup(&self, node_id: PeerId) -> Result<Vec<NodeRecord>, Discv4Error> {
346 self.lookup_node(Some(node_id)).await
347 }
348
349 pub async fn lookup_random(&self) -> Result<Vec<NodeRecord>, Discv4Error> {
351 let target = PeerId::random();
352 self.lookup_node(Some(target)).await
353 }
354
355 pub fn send_lookup(&self, node_id: PeerId) {
357 let cmd = Discv4Command::Lookup { node_id: Some(node_id), tx: None };
358 self.send_to_service(cmd);
359 }
360
361 async fn lookup_node(&self, node_id: Option<PeerId>) -> Result<Vec<NodeRecord>, Discv4Error> {
362 let (tx, rx) = oneshot::channel();
363 let cmd = Discv4Command::Lookup { node_id, tx: Some(tx) };
364 self.to_service.send(cmd)?;
365 Ok(rx.await?)
366 }
367
368 pub fn send_lookup_self(&self) {
370 let cmd = Discv4Command::Lookup { node_id: None, tx: None };
371 self.send_to_service(cmd);
372 }
373
374 pub fn remove_peer(&self, node_id: PeerId) {
376 let cmd = Discv4Command::Remove(node_id);
377 self.send_to_service(cmd);
378 }
379
380 pub fn add_node(&self, node_record: NodeRecord) {
382 let cmd = Discv4Command::Add(node_record);
383 self.send_to_service(cmd);
384 }
385
386 pub fn add_boot_node(&self, node_record: NodeRecord) {
391 let cmd = Discv4Command::AddBootNode(node_record);
392 self.send_to_service(cmd);
393 }
394
395 pub fn ban(&self, node_id: PeerId, ip: IpAddr) {
399 let cmd = Discv4Command::Ban(node_id, ip);
400 self.send_to_service(cmd);
401 }
402
403 pub fn ban_ip(&self, ip: IpAddr) {
407 let cmd = Discv4Command::BanIp(ip);
408 self.send_to_service(cmd);
409 }
410
411 pub fn ban_node(&self, node_id: PeerId) {
415 let cmd = Discv4Command::BanPeer(node_id);
416 self.send_to_service(cmd);
417 }
418
419 pub fn set_tcp_port(&self, port: u16) {
423 let cmd = Discv4Command::SetTcpPort(port);
424 self.send_to_service(cmd);
425 }
426
427 pub fn set_eip868_rlp_pair(&self, key: Vec<u8>, rlp: Bytes) {
433 let cmd = Discv4Command::SetEIP868RLPPair { key, rlp };
434 self.send_to_service(cmd);
435 }
436
437 pub fn set_eip868_rlp(&self, key: Vec<u8>, value: impl alloy_rlp::Encodable) {
441 self.set_eip868_rlp_pair(key, Bytes::from(alloy_rlp::encode(&value)))
442 }
443
444 #[inline]
445 fn send_to_service(&self, cmd: Discv4Command) {
446 let _ = self.to_service.send(cmd).map_err(|err| {
447 debug!(
448 target: "discv4",
449 %err,
450 "channel capacity reached, dropping command",
451 )
452 });
453 }
454
455 pub async fn update_stream(&self) -> Result<ReceiverStream<DiscoveryUpdate>, Discv4Error> {
457 let (tx, rx) = oneshot::channel();
458 let cmd = Discv4Command::Updates(tx);
459 self.to_service.send(cmd)?;
460 Ok(rx.await?)
461 }
462
463 pub fn terminate(&self) {
465 self.send_to_service(Discv4Command::Terminated);
466 }
467}
468
469#[must_use = "Stream does nothing unless polled"]
483pub struct Discv4Service {
484 local_address: SocketAddr,
486 local_eip_868_enr: Enr<SecretKey>,
488 local_node_record: NodeRecord,
490 shared_node_record: Arc<Mutex<NodeRecord>>,
492 secret_key: SecretKey,
494 _socket: Arc<UdpSocket>,
496 _tasks: JoinSet<()>,
500 kbuckets: KBucketsTable<NodeKey, NodeEntry>,
502 ingress: IngressReceiver,
506 egress: EgressSender,
510 queued_pings: VecDeque<(NodeRecord, PingReason)>,
517 pending_pings: HashMap<PeerId, PingRequest>,
519 pending_lookup: HashMap<PeerId, (Instant, LookupContext)>,
524 pending_find_nodes: HashMap<PeerId, FindNodeRequest>,
526 pending_enr_requests: HashMap<PeerId, EnrRequestState>,
528 to_service: mpsc::UnboundedSender<Discv4Command>,
530 commands_rx: mpsc::UnboundedReceiver<Discv4Command>,
532 update_listeners: Vec<mpsc::Sender<DiscoveryUpdate>>,
534 lookup_interval: Interval,
536 lookup_rotator: LookupTargetRotator,
538 pending_lookup_reset: bool,
540 evict_expired_requests_interval: Interval,
542 ping_interval: Interval,
544 resolve_external_ip_interval: Option<ResolveNatInterval>,
546 config: Discv4Config,
548 queued_events: VecDeque<Discv4Event>,
550 received_pongs: PongTable,
552 expire_interval: Interval,
554 cached_find_node: Option<CachedFindNode>,
556}
557
558impl Discv4Service {
559 pub(crate) fn new(
564 socket: Arc<UdpSocket>,
565 ingress_tx: Option<IngressSender>,
566 ingress_rx: IngressReceiver,
567 local_address: SocketAddr,
568 local_node_record: NodeRecord,
569 secret_key: SecretKey,
570 config: Discv4Config,
571 ) -> Self {
572 let (egress_tx, egress_rx) = mpsc::channel(config.udp_egress_message_buffer);
573 let mut tasks = JoinSet::<()>::new();
574
575 if let Some(ingress_tx) = ingress_tx {
576 let udp = Arc::clone(&socket);
577 tasks.spawn(receive_loop(udp, ingress_tx, local_node_record.id));
578 }
579
580 let udp = Arc::clone(&socket);
581 tasks.spawn(send_loop(udp, egress_rx));
582
583 let kbuckets = KBucketsTable::new(
584 NodeKey::from(&local_node_record).into(),
585 Duration::from_secs(60),
586 MAX_NODES_PER_BUCKET,
587 None,
588 None,
589 );
590
591 let self_lookup_interval = tokio::time::interval(config.lookup_interval);
592
593 let ping_interval = tokio::time::interval_at(
596 tokio::time::Instant::now() + config.ping_interval,
597 config.ping_interval,
598 );
599
600 let evict_expired_requests_interval = tokio::time::interval_at(
601 tokio::time::Instant::now() + config.request_timeout,
602 config.request_timeout,
603 );
604
605 let lookup_rotator = if config.enable_dht_random_walk {
606 LookupTargetRotator::default()
607 } else {
608 LookupTargetRotator::local_only()
609 };
610
611 let local_eip_868_enr = {
613 let mut builder = Enr::builder();
614 builder.ip(local_node_record.address);
615 if local_node_record.address.is_ipv4() {
616 builder.udp4(local_node_record.udp_port);
617 builder.tcp4(local_node_record.tcp_port);
618 } else {
619 builder.udp6(local_node_record.udp_port);
620 builder.tcp6(local_node_record.tcp_port);
621 }
622
623 for (key, val) in &config.additional_eip868_rlp_pairs {
624 builder.add_value_rlp(key, val.clone());
625 }
626 builder.build(&secret_key).expect("v4 is set")
627 };
628
629 let (to_service, commands_rx) = mpsc::unbounded_channel();
630
631 let shared_node_record = Arc::new(Mutex::new(local_node_record));
632
633 Self {
634 local_address,
635 local_eip_868_enr,
636 local_node_record,
637 shared_node_record,
638 _socket: socket,
639 kbuckets,
640 secret_key,
641 _tasks: tasks,
642 ingress: ingress_rx,
643 egress: egress_tx,
644 queued_pings: VecDeque::with_capacity(MAX_QUEUED_PINGS),
645 pending_pings: Default::default(),
646 pending_lookup: Default::default(),
647 pending_find_nodes: Default::default(),
648 pending_enr_requests: Default::default(),
649 commands_rx,
650 to_service,
651 update_listeners: Vec::with_capacity(1),
652 lookup_interval: self_lookup_interval,
653 ping_interval,
654 evict_expired_requests_interval,
655 lookup_rotator,
656 pending_lookup_reset: config.enable_lookup,
657 resolve_external_ip_interval: config.resolve_external_ip_interval(),
658 config,
659 queued_events: Default::default(),
660 received_pongs: Default::default(),
661 expire_interval: tokio::time::interval(EXPIRE_DURATION),
662 cached_find_node: None,
663 }
664 }
665
666 pub fn handle(&self) -> Discv4 {
668 Discv4 {
669 local_addr: self.local_address,
670 to_service: self.to_service.clone(),
671 node_record: self.shared_node_record.clone(),
672 }
673 }
674
675 fn enr_seq(&self) -> Option<u64> {
677 self.config.enable_eip868.then(|| self.local_eip_868_enr.seq())
678 }
679
680 pub fn set_lookup_interval(&mut self, duration: Duration) {
682 self.lookup_interval = tokio::time::interval(duration);
683 }
684
685 fn resolve_external_ip(&mut self) {
689 if let Some(r) = &self.resolve_external_ip_interval &&
690 let Some(external_ip) =
691 r.resolver().clone().as_external_ip(self.local_node_record.udp_port)
692 {
693 self.set_external_ip_addr(external_ip);
694 }
695 }
696
697 pub fn set_external_ip_addr(&mut self, external_ip: IpAddr) {
700 if self.local_node_record.address != external_ip {
701 debug!(target: "discv4", ?external_ip, "Updating external ip");
702 self.local_node_record.address = external_ip;
703 let _ = self.local_eip_868_enr.set_ip(external_ip, &self.secret_key);
704 let mut lock = self.shared_node_record.lock();
705 *lock = self.local_node_record;
706 debug!(target: "discv4", enr=?self.local_eip_868_enr, "Updated local ENR");
707 }
708 }
709
710 pub const fn local_peer_id(&self) -> &PeerId {
712 &self.local_node_record.id
713 }
714
715 pub const fn local_addr(&self) -> SocketAddr {
717 self.local_address
718 }
719
720 pub const fn local_enr(&self) -> NodeRecord {
724 self.local_node_record
725 }
726
727 #[cfg(test)]
729 pub const fn local_enr_mut(&mut self) -> &mut NodeRecord {
730 &mut self.local_node_record
731 }
732
733 pub fn contains_node(&self, id: PeerId) -> bool {
735 let key = kad_key(id);
736 self.kbuckets.get_index(&key).is_some()
737 }
738
739 pub fn bootstrap(&mut self) {
751 for record in self.config.bootstrap_nodes.clone() {
752 debug!(target: "discv4", ?record, "pinging boot node");
753 let key = kad_key(record.id);
754 let entry = NodeEntry::new(record);
755
756 match self.kbuckets.insert_or_update(
758 &key,
759 entry,
760 NodeStatus {
761 state: ConnectionState::Disconnected,
762 direction: ConnectionDirection::Outgoing,
763 },
764 ) {
765 InsertResult::Failed(_) => {}
766 _ => {
767 self.try_ping(record, PingReason::InitialInsert);
768 }
769 }
770 }
771 }
772
773 pub fn add_boot_node(&mut self, record: NodeRecord) -> bool {
778 self.config.bootstrap_nodes.insert(record);
779 self.add_node(record)
780 }
781
782 pub fn spawn(mut self) -> JoinHandle<()> {
786 tokio::task::spawn(async move {
787 self.bootstrap();
788
789 while let Some(event) = self.next().await {
790 trace!(target: "discv4", ?event, "processed");
791 }
792 trace!(target: "discv4", "service terminated");
793 })
794 }
795
796 pub fn update_stream(&mut self) -> ReceiverStream<DiscoveryUpdate> {
798 let (tx, rx) = mpsc::channel(512);
799 self.update_listeners.push(tx);
800 ReceiverStream::new(rx)
801 }
802
803 pub fn lookup_self(&mut self) {
805 self.lookup(self.local_node_record.id)
806 }
807
808 pub fn lookup(&mut self, target: PeerId) {
818 self.lookup_with(target, None)
819 }
820
821 fn lookup_with(&mut self, target: PeerId, tx: Option<NodeRecordSender>) {
831 trace!(target: "discv4", ?target, "Starting lookup");
832 let target_key = kad_key(target);
833
834 let ctx = LookupContext::new(
837 target_key.clone(),
838 self.kbuckets
839 .closest_values(&target_key)
840 .filter(|node| {
841 node.value.has_endpoint_proof &&
842 !self.pending_find_nodes.contains_key(&node.key.preimage().0)
843 })
844 .take(MAX_NODES_PER_BUCKET)
845 .map(|n| (target_key.distance(&n.key), n.value.record)),
846 tx,
847 );
848
849 let closest = ctx.closest(ALPHA);
851
852 if closest.is_empty() && self.pending_find_nodes.is_empty() {
853 self.bootstrap();
858 return
859 }
860
861 trace!(target: "discv4", ?target, num = closest.len(), "Start lookup closest nodes");
862
863 for node in closest {
864 self.find_node_checked(&node, ctx.clone());
868 }
869 }
870
871 fn find_node(&mut self, node: &NodeRecord, ctx: LookupContext) {
875 trace!(target: "discv4", ?node, lookup=?ctx.target(), "Sending FindNode");
876 ctx.mark_queried(node.id);
877 let (payload, hash) = self.find_node_packet(ctx.target());
878 let to = node.udp_addr();
879 trace!(target: "discv4", ?to, ?hash, "sending FindNode packet");
880 let _ = self.egress.try_send((payload, to)).map_err(|err| {
881 debug!(target: "discv4", %err, "dropped outgoing packet");
882 });
883 self.pending_find_nodes.insert(node.id, FindNodeRequest::new(ctx));
884 }
885
886 fn find_node_checked(&mut self, node: &NodeRecord, ctx: LookupContext) {
891 let max_failures = self.config.max_find_node_failures;
892 let needs_ping = self
893 .on_entry(node.id, |entry| entry.exceeds_find_node_failures(max_failures))
894 .unwrap_or(true);
895 if needs_ping {
896 self.try_ping(*node, PingReason::Lookup(*node, ctx))
897 } else {
898 self.find_node(node, ctx)
899 }
900 }
901
902 fn notify(&mut self, update: DiscoveryUpdate) {
906 self.update_listeners.retain_mut(|listener| match listener.try_send(update.clone()) {
907 Ok(()) => true,
908 Err(err) => match err {
909 TrySendError::Full(_) => true,
910 TrySendError::Closed(_) => false,
911 },
912 });
913 }
914
915 pub fn ban_ip(&mut self, ip: IpAddr) {
917 self.config.ban_list.ban_ip(ip);
918 }
919
920 pub fn ban_node(&mut self, node_id: PeerId) {
922 self.remove_node(node_id);
923 self.config.ban_list.ban_peer(node_id);
924 }
925
926 pub fn ban_ip_until(&mut self, ip: IpAddr, until: Instant) {
928 self.config.ban_list.ban_ip_until(ip, until);
929 }
930
931 pub fn ban_node_until(&mut self, node_id: PeerId, until: Instant) {
933 self.remove_node(node_id);
934 self.config.ban_list.ban_peer_until(node_id, until);
935 }
936
937 pub fn remove_node(&mut self, node_id: PeerId) -> bool {
942 let key = kad_key(node_id);
943 self.remove_key(node_id, key)
944 }
945
946 pub fn soft_remove_node(&mut self, node_id: PeerId) -> bool {
951 let key = kad_key(node_id);
952 let Some(bucket) = self.kbuckets.get_bucket(&key) else { return false };
953 if bucket.num_entries() < MAX_NODES_PER_BUCKET / 2 {
954 return false
956 }
957 self.remove_key(node_id, key)
958 }
959
960 fn remove_key(&mut self, node_id: PeerId, key: discv5::Key<NodeKey>) -> bool {
961 let removed = self.kbuckets.remove(&key);
962 if removed {
963 trace!(target: "discv4", ?node_id, "removed node");
964 self.notify(DiscoveryUpdate::Removed(node_id));
965 }
966 removed
967 }
968
969 pub fn num_connected(&self) -> usize {
971 self.kbuckets.buckets_iter().fold(0, |count, bucket| count + bucket.num_connected())
972 }
973
974 fn has_bond(&self, remote_id: PeerId, remote_ip: IpAddr) -> bool {
976 if let Some(timestamp) = self.received_pongs.last_pong(remote_id, remote_ip) &&
977 timestamp.elapsed() < self.config.bond_expiration
978 {
979 return true
980 }
981 false
982 }
983
984 fn on_entry<F, R>(&mut self, peer_id: PeerId, f: F) -> Option<R>
986 where
987 F: FnOnce(&NodeEntry) -> R,
988 {
989 let key = kad_key(peer_id);
990 match self.kbuckets.entry(&key) {
991 BucketEntry::Present(entry, _) => Some(f(entry.value())),
992 BucketEntry::Pending(entry, _) => Some(f(entry.value())),
993 _ => None,
994 }
995 }
996
997 fn update_on_reping(&mut self, record: NodeRecord, mut last_enr_seq: Option<u64>) {
1004 if record.id == self.local_node_record.id {
1005 return
1006 }
1007
1008 if !self.config.enable_eip868 {
1010 last_enr_seq = None;
1011 }
1012
1013 let key = kad_key(record.id);
1014 let old_enr = match self.kbuckets.entry(&key) {
1015 kbucket::Entry::Present(mut entry, _) => {
1016 entry.value_mut().update_with_enr(last_enr_seq)
1017 }
1018 kbucket::Entry::Pending(mut entry, _) => {
1019 entry.value_mut().update_with_enr(last_enr_seq)
1020 }
1021 _ => return,
1022 };
1023
1024 match (last_enr_seq, old_enr) {
1026 (Some(new), Some(old)) if new > old => {
1027 self.send_enr_request(record);
1028 }
1029 (Some(_), None) => {
1030 self.send_enr_request(record);
1032 }
1033 _ => {}
1034 };
1035 }
1036
1037 fn update_on_pong(&mut self, record: NodeRecord, mut last_enr_seq: Option<u64>) {
1039 if record.id == *self.local_peer_id() {
1040 return
1041 }
1042
1043 if !self.config.enable_eip868 {
1045 last_enr_seq = None;
1046 }
1047
1048 let has_enr_seq = last_enr_seq.is_some();
1051
1052 let key = kad_key(record.id);
1053 match self.kbuckets.entry(&key) {
1054 kbucket::Entry::Present(mut entry, old_status) => {
1055 entry.value_mut().establish_proof();
1057 entry.value_mut().update_with_enr(last_enr_seq);
1058
1059 if !old_status.is_connected() {
1060 let _ = entry.update(ConnectionState::Connected, Some(old_status.direction));
1061 trace!(target: "discv4", ?record, "added after successful endpoint proof");
1062 self.notify(DiscoveryUpdate::Added(record));
1063
1064 if has_enr_seq {
1065 self.send_enr_request(record);
1067 }
1068 }
1069 }
1070 kbucket::Entry::Pending(mut entry, mut status) => {
1071 entry.value_mut().establish_proof();
1073 entry.value_mut().update_with_enr(last_enr_seq);
1074
1075 if !status.is_connected() {
1076 status.state = ConnectionState::Connected;
1077 let _ = entry.update(status);
1078 trace!(target: "discv4", ?record, "added after successful endpoint proof");
1079 self.notify(DiscoveryUpdate::Added(record));
1080
1081 if has_enr_seq {
1082 self.send_enr_request(record);
1084 }
1085 }
1086 }
1087 _ => {}
1088 };
1089 }
1090
1091 pub fn add_all_nodes(&mut self, records: impl IntoIterator<Item = NodeRecord>) {
1095 for record in records {
1096 self.add_node(record);
1097 }
1098 }
1099
1100 pub fn add_node(&mut self, record: NodeRecord) -> bool {
1106 let key = kad_key(record.id);
1107 match self.kbuckets.entry(&key) {
1108 kbucket::Entry::Absent(entry) => {
1109 let node = NodeEntry::new(record);
1110 match entry.insert(
1111 node,
1112 NodeStatus {
1113 direction: ConnectionDirection::Outgoing,
1114 state: ConnectionState::Disconnected,
1115 },
1116 ) {
1117 BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1118 trace!(target: "discv4", ?record, "inserted new record");
1119 }
1120 _ => return false,
1121 }
1122 }
1123 _ => return false,
1124 }
1125
1126 self.try_ping(record, PingReason::InitialInsert);
1128 true
1129 }
1130
1131 pub(crate) fn send_packet(&self, msg: Message, to: SocketAddr) -> B256 {
1133 let (payload, hash) = msg.encode(&self.secret_key);
1134 trace!(target: "discv4", r#type=?msg.msg_type(), ?to, ?hash, "sending packet");
1135 let _ = self.egress.try_send((payload, to)).map_err(|err| {
1136 debug!(
1137 target: "discv4",
1138 %err,
1139 "dropped outgoing packet",
1140 );
1141 });
1142 hash
1143 }
1144
1145 fn find_node_packet(&mut self, target: PeerId) -> (Bytes, B256) {
1147 let expire = self.find_node_expiration();
1148 let cache_ttl = self.config.request_timeout / 4;
1149 CachedFindNode::get_or_sign(
1150 &mut self.cached_find_node,
1151 target,
1152 cache_ttl,
1153 &self.secret_key,
1154 expire,
1155 )
1156 }
1157
1158 fn on_ping(&mut self, ping: Ping, remote_addr: SocketAddr, remote_id: PeerId, hash: B256) {
1160 if self.is_expired(ping.expire) {
1161 return
1163 }
1164
1165 let record = NodeRecord {
1167 address: remote_addr.ip(),
1168 udp_port: remote_addr.port(),
1169 tcp_port: ping.from.tcp_port,
1170 id: remote_id,
1171 }
1172 .into_ipv4_mapped();
1173
1174 let key = kad_key(record.id);
1175
1176 let mut is_new_insert = false;
1183 let mut needs_bond = false;
1184 let mut is_proven = false;
1185
1186 let old_enr = match self.kbuckets.entry(&key) {
1187 kbucket::Entry::Present(mut entry, _) => {
1188 if entry.value().is_expired() {
1189 needs_bond = true;
1192 } else {
1193 is_proven = entry.value().has_endpoint_proof;
1194 }
1195 entry.value_mut().update_with_enr(ping.enr_sq)
1196 }
1197 kbucket::Entry::Pending(mut entry, _) => {
1198 if entry.value().is_expired() {
1199 needs_bond = true;
1202 } else {
1203 is_proven = entry.value().has_endpoint_proof;
1204 }
1205 entry.value_mut().update_with_enr(ping.enr_sq)
1206 }
1207 kbucket::Entry::Absent(entry) => {
1208 let mut node = NodeEntry::new(record);
1209 node.last_enr_seq = ping.enr_sq;
1210
1211 match entry.insert(
1212 node,
1213 NodeStatus {
1214 direction: ConnectionDirection::Incoming,
1215 state: ConnectionState::Disconnected,
1217 },
1218 ) {
1219 BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1220 is_new_insert = true;
1222 }
1223 BucketInsertResult::Full => {
1224 trace!(target: "discv4", ?record, "discovered new record but bucket is full");
1228 self.notify(DiscoveryUpdate::DiscoveredAtCapacity(record));
1229 needs_bond = true;
1230 }
1231 BucketInsertResult::TooManyIncoming | BucketInsertResult::NodeExists => {
1232 needs_bond = true;
1233 }
1235 BucketInsertResult::FailedFilter => return,
1236 }
1237
1238 None
1239 }
1240 kbucket::Entry::SelfEntry => return,
1241 };
1242
1243 let pong = Message::Pong(Pong {
1246 to: record.into(),
1248 echo: hash,
1249 expire: ping.expire,
1250 enr_sq: self.enr_seq(),
1251 });
1252 self.send_packet(pong, remote_addr);
1253
1254 if is_new_insert {
1256 self.try_ping(record, PingReason::InitialInsert);
1257 } else if needs_bond {
1258 self.try_ping(record, PingReason::EstablishBond);
1259 } else {
1260 if is_proven && let Some((_, ctx)) = self.pending_lookup.remove(&record.id) {
1264 if self.pending_find_nodes.contains_key(&record.id) {
1265 ctx.unmark_queried(record.id);
1268 } else {
1269 self.find_node(&record, ctx);
1272 }
1273 }
1274
1275 match (ping.enr_sq, old_enr) {
1277 (Some(new), Some(old)) if new > old => {
1278 self.send_enr_request(record);
1279 }
1280 (Some(_), None) => {
1281 self.send_enr_request(record);
1282 }
1283 _ => {}
1284 };
1285 }
1286 }
1287
1288 fn try_ping(&mut self, node: NodeRecord, reason: PingReason) {
1290 if node.id == *self.local_peer_id() {
1291 return
1293 }
1294
1295 if self.pending_pings.contains_key(&node.id) ||
1296 self.pending_find_nodes.contains_key(&node.id)
1297 {
1298 return
1299 }
1300
1301 if self.queued_pings.iter().any(|(n, _)| n.id == node.id) {
1302 return
1303 }
1304
1305 if self.pending_pings.len() < MAX_NODES_PING {
1306 self.send_ping(node, reason);
1307 } else if self.queued_pings.len() < MAX_QUEUED_PINGS {
1308 self.queued_pings.push_back((node, reason));
1309 }
1310 }
1311
1312 pub(crate) fn send_ping(&mut self, node: NodeRecord, reason: PingReason) -> B256 {
1316 let remote_addr = node.udp_addr();
1317 let id = node.id;
1318 let ping = Ping {
1319 from: self.local_node_record.into(),
1320 to: node.into(),
1321 expire: self.ping_expiration(),
1322 enr_sq: self.enr_seq(),
1323 };
1324 trace!(target: "discv4", ?ping, "sending ping");
1325 let echo_hash = self.send_packet(Message::Ping(ping), remote_addr);
1326
1327 self.pending_pings
1328 .insert(id, PingRequest { sent_at: Instant::now(), node, echo_hash, reason });
1329 echo_hash
1330 }
1331
1332 pub(crate) fn send_enr_request(&mut self, node: NodeRecord) {
1336 if !self.config.enable_eip868 {
1337 return
1338 }
1339 let remote_addr = node.udp_addr();
1340 let enr_request = EnrRequest { expire: self.enr_request_expiration() };
1341
1342 trace!(target: "discv4", ?enr_request, "sending enr request");
1343 let echo_hash = self.send_packet(Message::EnrRequest(enr_request), remote_addr);
1344
1345 self.pending_enr_requests
1346 .insert(node.id, EnrRequestState { sent_at: Instant::now(), echo_hash });
1347 }
1348
1349 fn on_pong(&mut self, pong: Pong, remote_addr: SocketAddr, remote_id: PeerId) {
1351 if self.is_expired(pong.expire) {
1352 return
1353 }
1354
1355 let PingRequest { node, reason, .. } = match self.pending_pings.entry(remote_id) {
1356 Entry::Occupied(entry) => {
1357 {
1358 let request = entry.get();
1359 if request.echo_hash != pong.echo {
1360 trace!(target: "discv4", from=?remote_addr, expected=?request.echo_hash, echo_hash=?pong.echo,"Got unexpected Pong");
1361 return
1362 }
1363 }
1364 entry.remove()
1365 }
1366 Entry::Vacant(_) => return,
1367 };
1368
1369 self.received_pongs.on_pong(remote_id, remote_addr.ip());
1371
1372 match reason {
1373 PingReason::InitialInsert => {
1374 self.update_on_pong(node, pong.enr_sq);
1375 if self.pending_lookup_reset && self.config.bootstrap_nodes.contains(&node) {
1378 self.pending_lookup_reset = false;
1379 self.lookup_interval.reset();
1380 }
1381 }
1382 PingReason::EstablishBond => {
1383 self.update_on_pong(node, pong.enr_sq);
1385 }
1386 PingReason::RePing => {
1387 self.update_on_reping(node, pong.enr_sq);
1388 }
1389 PingReason::Lookup(node, ctx) => {
1390 self.update_on_pong(node, pong.enr_sq);
1391 self.pending_lookup.insert(node.id, (Instant::now(), ctx));
1396 }
1397 }
1398 }
1399
1400 fn on_find_node(&mut self, msg: FindNode, remote_addr: SocketAddr, node_id: PeerId) {
1402 if self.is_expired(msg.expire) {
1403 return
1405 }
1406 if node_id == *self.local_peer_id() {
1407 return
1409 }
1410
1411 if self.has_bond(node_id, remote_addr.ip()) {
1412 self.respond_closest(msg.id, remote_addr)
1413 }
1414 }
1415
1416 fn on_enr_response(&mut self, msg: EnrResponse, remote_addr: SocketAddr, id: PeerId) {
1418 trace!(target: "discv4", ?remote_addr, ?msg, "received ENR response");
1419 let Entry::Occupied(request) = self.pending_enr_requests.entry(id) else { return };
1420
1421 if request.get().echo_hash != msg.request_hash || id != pk2id(&msg.enr.public_key()) {
1424 return
1425 }
1426 request.remove();
1427
1428 let key = kad_key(id);
1429 let fork_id = msg.eth_fork_id();
1430 let (record, old_fork_id) = match self.kbuckets.entry(&key) {
1431 kbucket::Entry::Present(mut entry, _) => {
1432 let id = entry.value_mut().update_with_fork_id(fork_id);
1433 (entry.value().record, id)
1434 }
1435 kbucket::Entry::Pending(mut entry, _) => {
1436 let id = entry.value_mut().update_with_fork_id(fork_id);
1437 (entry.value().record, id)
1438 }
1439 _ => return,
1440 };
1441 match (fork_id, old_fork_id) {
1442 (Some(new), Some(old)) if new != old => {
1443 self.notify(DiscoveryUpdate::EnrForkId(record, new))
1444 }
1445 (Some(new), None) => self.notify(DiscoveryUpdate::EnrForkId(record, new)),
1446 _ => {}
1447 }
1448 }
1449
1450 fn on_enr_request(
1452 &self,
1453 msg: EnrRequest,
1454 remote_addr: SocketAddr,
1455 id: PeerId,
1456 request_hash: B256,
1457 ) {
1458 if !self.config.enable_eip868 || self.is_expired(msg.expire) {
1459 return
1460 }
1461
1462 if self.has_bond(id, remote_addr.ip()) {
1463 self.send_packet(
1464 Message::EnrResponse(EnrResponse {
1465 request_hash,
1466 enr: self.local_eip_868_enr.clone(),
1467 }),
1468 remote_addr,
1469 );
1470 }
1471 }
1472
1473 fn on_neighbours(&mut self, msg: Neighbours, remote_addr: SocketAddr, node_id: PeerId) {
1476 if self.is_expired(msg.expire) {
1477 return
1479 }
1480 let ctx = match self.pending_find_nodes.entry(node_id) {
1482 Entry::Occupied(mut entry) => {
1483 {
1484 let request = entry.get_mut();
1485 request.answered = true;
1487 let total = request.response_count + msg.nodes.len();
1488
1489 if total <= MAX_NODES_PER_BUCKET {
1491 request.response_count = total;
1492 } else {
1493 trace!(target: "discv4", total, from=?remote_addr, "Received neighbors packet entries exceeds max nodes per bucket");
1494 return
1495 }
1496 };
1497
1498 if entry.get().response_count == MAX_NODES_PER_BUCKET {
1499 let ctx = entry.remove().lookup_context;
1501 ctx.mark_responded(node_id);
1502 ctx
1503 } else {
1504 entry.get().lookup_context.clone()
1505 }
1506 }
1507 Entry::Vacant(_) => {
1508 trace!(target: "discv4", from=?remote_addr, "Received unsolicited Neighbours");
1510 return
1511 }
1512 };
1513
1514 trace!(target: "discv4",
1516 target=format!("{:#?}", node_id),
1517 peers_count=msg.nodes.len(),
1518 peers=format!("[{:#}]", msg.nodes.iter()
1519 .map(|node_rec| node_rec.id
1520 ).format(", ")),
1521 "Received peers from Neighbours packet"
1522 );
1523
1524 for node in msg.nodes.into_iter().map(NodeRecord::into_ipv4_mapped) {
1527 if self.config.ban_list.is_banned(&node.id, &node.address) {
1529 trace!(target: "discv4", peer_id=?node.id, ip=?node.address, "ignoring banned record");
1530 continue
1531 }
1532
1533 ctx.add_node(node);
1534 }
1535
1536 let closest =
1538 ctx.filter_closest(ALPHA, |node| !self.pending_find_nodes.contains_key(&node.id));
1539
1540 for closest in closest {
1541 let key = kad_key(closest.id);
1542 match self.kbuckets.entry(&key) {
1543 BucketEntry::Absent(entry) => {
1544 ctx.mark_queried(closest.id);
1550 let node = NodeEntry::new(closest);
1551 match entry.insert(
1552 node,
1553 NodeStatus {
1554 direction: ConnectionDirection::Outgoing,
1555 state: ConnectionState::Disconnected,
1556 },
1557 ) {
1558 BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1559 self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1561 }
1562 BucketInsertResult::Full => {
1563 self.notify(DiscoveryUpdate::DiscoveredAtCapacity(closest))
1565 }
1566 _ => {}
1567 }
1568 }
1569 BucketEntry::SelfEntry => {
1570 }
1572 BucketEntry::Present(entry, _) => {
1573 if entry.value().has_endpoint_proof {
1574 if entry
1575 .value()
1576 .exceeds_find_node_failures(self.config.max_find_node_failures)
1577 {
1578 self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1579 } else {
1580 self.find_node(&closest, ctx.clone());
1581 }
1582 }
1583 }
1584 BucketEntry::Pending(entry, _) => {
1585 if entry.value().has_endpoint_proof {
1586 if entry
1587 .value()
1588 .exceeds_find_node_failures(self.config.max_find_node_failures)
1589 {
1590 self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1591 } else {
1592 self.find_node(&closest, ctx.clone());
1593 }
1594 }
1595 }
1596 }
1597 }
1598 }
1599
1600 fn respond_closest(&mut self, target: PeerId, to: SocketAddr) {
1602 let key = kad_key(target);
1603 let expire = self.send_neighbours_expiration();
1604
1605 let enforce_eip868 = self.config.enable_eip868 && self.config.enforce_eip868_neighbours;
1606
1607 let closest_nodes = self
1610 .kbuckets
1611 .closest_values(&key)
1612 .filter(|entry| !enforce_eip868 || entry.value.fork_id.is_some())
1613 .take(MAX_NODES_PER_BUCKET)
1614 .collect::<Vec<_>>();
1615
1616 if closest_nodes.is_empty() {
1617 let msg = Message::Neighbours(Neighbours { nodes: Vec::new(), expire });
1619 self.send_packet(msg, to);
1620 return;
1621 }
1622
1623 for nodes in closest_nodes.chunks(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS) {
1624 let nodes = nodes.iter().map(|node| node.value.record).collect::<Vec<NodeRecord>>();
1625 trace!(target: "discv4", len = nodes.len(), to=?to,"Sent neighbours packet");
1626 let msg = Message::Neighbours(Neighbours { nodes, expire });
1627 self.send_packet(msg, to);
1628 }
1629 }
1630
1631 fn evict_expired_requests(&mut self, now: Instant) {
1632 let mut failed_enr_requests = Vec::new();
1633 self.pending_enr_requests.retain(|node_id, enr_request| {
1634 if now.duration_since(enr_request.sent_at) < self.config.enr_expiration {
1635 return true
1636 }
1637 failed_enr_requests.push(*node_id);
1638 false
1639 });
1640
1641 for node_id in failed_enr_requests {
1643 match self.kbuckets.entry(&kad_key(node_id)) {
1644 kbucket::Entry::Present(mut entry, _) => entry.value_mut().last_enr_seq = None,
1645 kbucket::Entry::Pending(mut entry, _) => entry.value_mut().last_enr_seq = None,
1646 _ => {}
1647 }
1648 }
1649
1650 let mut failed_pings = Vec::new();
1651 self.pending_pings.retain(|node_id, ping_request| {
1652 if now.duration_since(ping_request.sent_at) > self.config.ping_expiration {
1653 failed_pings.push(*node_id);
1654 return false
1655 }
1656 true
1657 });
1658
1659 if !failed_pings.is_empty() {
1660 trace!(target: "discv4", num=%failed_pings.len(), "evicting nodes due to failed pong");
1662 for node_id in failed_pings {
1663 self.remove_node(node_id);
1664 }
1665 }
1666
1667 let mut failed_lookups = Vec::new();
1668 self.pending_lookup.retain(|node_id, (lookup_sent_at, _)| {
1669 if now.duration_since(*lookup_sent_at) > self.config.request_timeout {
1670 failed_lookups.push(*node_id);
1671 return false
1672 }
1673 true
1674 });
1675
1676 if !failed_lookups.is_empty() {
1677 trace!(target: "discv4", num=%failed_lookups.len(), "evicting nodes due to failed lookup");
1679 for node_id in failed_lookups {
1680 self.remove_node(node_id);
1681 }
1682 }
1683
1684 self.evict_failed_find_nodes(now);
1685 }
1686
1687 fn evict_failed_find_nodes(&mut self, now: Instant) {
1689 let mut failed_find_nodes = Vec::new();
1690 self.pending_find_nodes.retain(|node_id, find_node_request| {
1691 if now.duration_since(find_node_request.sent_at) > self.config.neighbours_expiration {
1692 if !find_node_request.answered {
1693 failed_find_nodes.push(*node_id);
1696 }
1697 return false
1698 }
1699 true
1700 });
1701
1702 if failed_find_nodes.is_empty() {
1703 return
1704 }
1705
1706 trace!(target: "discv4", num=%failed_find_nodes.len(), "processing failed find nodes");
1707
1708 for node_id in failed_find_nodes {
1709 let key = kad_key(node_id);
1710 let failures = match self.kbuckets.entry(&key) {
1711 kbucket::Entry::Present(mut entry, _) => {
1712 entry.value_mut().inc_failed_request();
1713 entry.value().find_node_failures
1714 }
1715 kbucket::Entry::Pending(mut entry, _) => {
1716 entry.value_mut().inc_failed_request();
1717 entry.value().find_node_failures
1718 }
1719 _ => continue,
1720 };
1721
1722 if failures > self.config.max_find_node_failures {
1726 self.soft_remove_node(node_id);
1727 }
1728 }
1729 }
1730
1731 fn re_ping_oldest(&mut self) {
1736 let mut nodes = self
1737 .kbuckets
1738 .iter_ref()
1739 .filter(|entry| entry.node.value.is_expired())
1740 .map(|n| n.node.value)
1741 .collect::<Vec<_>>();
1742 nodes.sort_unstable_by_key(|a| a.last_seen);
1743 let to_ping = nodes.into_iter().map(|n| n.record).take(MAX_NODES_PING).collect::<Vec<_>>();
1744 for node in to_ping {
1745 self.try_ping(node, PingReason::RePing)
1746 }
1747 }
1748
1749 fn is_expired(&self, expiration: u64) -> bool {
1751 self.ensure_not_expired(expiration).is_err()
1752 }
1753
1754 fn ensure_not_expired(&self, timestamp: u64) -> Result<(), ()> {
1764 let _ = i64::try_from(timestamp).map_err(drop)?;
1766
1767 let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
1768 if self.config.enforce_expiration_timestamps && timestamp < now {
1769 trace!(target: "discv4", "Expired packet");
1770 return Err(())
1771 }
1772 Ok(())
1773 }
1774
1775 fn ping_buffered(&mut self) {
1777 while self.pending_pings.len() < MAX_NODES_PING {
1778 match self.queued_pings.pop_front() {
1779 Some((next, reason)) => self.try_ping(next, reason),
1780 None => break,
1781 }
1782 }
1783 }
1784
1785 fn ping_expiration(&self) -> u64 {
1786 (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.ping_expiration)
1787 .as_secs()
1788 }
1789
1790 fn find_node_expiration(&self) -> u64 {
1791 (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.request_timeout)
1792 .as_secs()
1793 }
1794
1795 fn enr_request_expiration(&self) -> u64 {
1796 (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.enr_expiration)
1797 .as_secs()
1798 }
1799
1800 fn send_neighbours_expiration(&self) -> u64 {
1801 (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.neighbours_expiration)
1802 .as_secs()
1803 }
1804
1805 pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Discv4Event> {
1811 loop {
1812 if let Some(event) = self.queued_events.pop_front() {
1814 return Poll::Ready(event)
1815 }
1816
1817 if self.config.enable_lookup {
1819 while self.lookup_interval.poll_tick(cx).is_ready() {
1820 let target = self.lookup_rotator.next(&self.local_node_record.id);
1821 self.lookup_with(target, None);
1822 }
1823 }
1824
1825 while self.ping_interval.poll_tick(cx).is_ready() {
1827 self.re_ping_oldest();
1828 }
1829
1830 if let Some(Poll::Ready(Some(ip))) =
1831 self.resolve_external_ip_interval.as_mut().map(|r| r.poll_tick(cx))
1832 {
1833 self.set_external_ip_addr(ip);
1834 }
1835
1836 while let Poll::Ready(Some(cmd)) = self.commands_rx.poll_recv(cx) {
1838 match cmd {
1839 Discv4Command::Add(enr) => {
1840 self.add_node(enr);
1841 }
1842 Discv4Command::AddBootNode(record) => {
1843 self.add_boot_node(record);
1844 }
1845 Discv4Command::Lookup { node_id, tx } => {
1846 let node_id = node_id.unwrap_or(self.local_node_record.id);
1847 self.lookup_with(node_id, tx);
1848 }
1849 Discv4Command::SetLookupInterval(duration) => {
1850 self.set_lookup_interval(duration);
1851 }
1852 Discv4Command::Updates(tx) => {
1853 let rx = self.update_stream();
1854 let _ = tx.send(rx);
1855 }
1856 Discv4Command::BanPeer(node_id) => self.ban_node(node_id),
1857 Discv4Command::Remove(node_id) => {
1858 self.remove_node(node_id);
1859 }
1860 Discv4Command::Ban(node_id, ip) => {
1861 self.ban_node(node_id);
1862 self.ban_ip(ip);
1863 }
1864 Discv4Command::BanIp(ip) => {
1865 self.ban_ip(ip);
1866 }
1867 Discv4Command::SetEIP868RLPPair { key, rlp } => {
1868 debug!(target: "discv4", key=%String::from_utf8_lossy(&key), "Update EIP-868 extension pair");
1869
1870 let _ = self.local_eip_868_enr.insert_raw_rlp(key, rlp, &self.secret_key);
1871 }
1872 Discv4Command::SetTcpPort(port) => {
1873 debug!(target: "discv4", %port, "Update tcp port");
1874 self.local_node_record.tcp_port = port;
1875 if self.local_node_record.address.is_ipv4() {
1876 let _ = self.local_eip_868_enr.set_tcp4(port, &self.secret_key);
1877 } else {
1878 let _ = self.local_eip_868_enr.set_tcp6(port, &self.secret_key);
1879 }
1880 }
1881
1882 Discv4Command::Terminated => {
1883 self.queued_events.push_back(Discv4Event::Terminated);
1885 }
1886 }
1887 }
1888
1889 let mut udp_message_budget = UDP_MESSAGE_POLL_LOOP_BUDGET;
1891
1892 while let Poll::Ready(Some(event)) = self.ingress.poll_recv(cx) {
1894 match event {
1895 IngressEvent::RecvError(err) => {
1896 debug!(target: "discv4", %err, "failed to read datagram");
1897 }
1898 IngressEvent::BadPacket(from, err, data) => {
1899 trace!(target: "discv4", ?from, %err, packet=?hex::encode(&data), "bad packet");
1900 }
1901 IngressEvent::Packet(remote_addr, Packet { msg, node_id, hash }) => {
1902 trace!(target: "discv4", r#type=?msg.msg_type(), from=?remote_addr,"received packet");
1903 let event = match msg {
1904 Message::Ping(ping) => {
1905 self.on_ping(ping, remote_addr, node_id, hash);
1906 Discv4Event::Ping
1907 }
1908 Message::Pong(pong) => {
1909 self.on_pong(pong, remote_addr, node_id);
1910 Discv4Event::Pong
1911 }
1912 Message::FindNode(msg) => {
1913 self.on_find_node(msg, remote_addr, node_id);
1914 Discv4Event::FindNode
1915 }
1916 Message::Neighbours(msg) => {
1917 self.on_neighbours(msg, remote_addr, node_id);
1918 Discv4Event::Neighbours
1919 }
1920 Message::EnrRequest(msg) => {
1921 self.on_enr_request(msg, remote_addr, node_id, hash);
1922 Discv4Event::EnrRequest
1923 }
1924 Message::EnrResponse(msg) => {
1925 self.on_enr_response(msg, remote_addr, node_id);
1926 Discv4Event::EnrResponse
1927 }
1928 };
1929
1930 self.queued_events.push_back(event);
1931 }
1932 }
1933
1934 udp_message_budget -= 1;
1935 if udp_message_budget < 0 {
1936 trace!(target: "discv4", budget=UDP_MESSAGE_POLL_LOOP_BUDGET, "exhausted message poll budget");
1937 if self.queued_events.is_empty() {
1938 cx.waker().wake_by_ref();
1941 }
1942 break
1943 }
1944 }
1945
1946 self.ping_buffered();
1948
1949 while self.evict_expired_requests_interval.poll_tick(cx).is_ready() {
1951 self.evict_expired_requests(Instant::now());
1952 }
1953
1954 while self.expire_interval.poll_tick(cx).is_ready() {
1956 self.received_pongs.evict_expired(Instant::now(), EXPIRE_DURATION);
1957 }
1958
1959 if self.queued_events.is_empty() {
1960 return Poll::Pending
1961 }
1962 }
1963 }
1964}
1965
1966impl Stream for Discv4Service {
1968 type Item = Discv4Event;
1969
1970 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1971 match ready!(self.get_mut().poll(cx)) {
1973 Discv4Event::Terminated => Poll::Ready(None),
1975 ev => Poll::Ready(Some(ev)),
1977 }
1978 }
1979}
1980
1981impl fmt::Debug for Discv4Service {
1982 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1983 f.debug_struct("Discv4Service")
1984 .field("local_address", &self.local_address)
1985 .field("local_peer_id", &self.local_peer_id())
1986 .field("local_node_record", &self.local_node_record)
1987 .field("queued_pings", &self.queued_pings)
1988 .field("pending_lookup", &self.pending_lookup)
1989 .field("pending_find_nodes", &self.pending_find_nodes)
1990 .field("lookup_interval", &self.lookup_interval)
1991 .finish_non_exhaustive()
1992 }
1993}
1994
1995#[derive(Debug, Eq, PartialEq)]
1999pub enum Discv4Event {
2000 Ping,
2002 Pong,
2004 FindNode,
2006 Neighbours,
2008 EnrRequest,
2010 EnrResponse,
2012 Terminated,
2014}
2015
2016pub(crate) async fn send_loop(udp: Arc<UdpSocket>, rx: EgressReceiver) {
2018 let mut stream = ReceiverStream::new(rx);
2019 while let Some((payload, to)) = stream.next().await {
2020 match udp.send_to(&payload, to).await {
2021 Ok(size) => {
2022 trace!(target: "discv4", ?to, ?size,"sent payload");
2023 }
2024 Err(err) => {
2025 debug!(target: "discv4", ?to, %err,"Failed to send datagram.");
2026 }
2027 }
2028 }
2029}
2030
2031const MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP: usize = 60usize;
2033
2034pub(crate) async fn receive_loop(udp: Arc<UdpSocket>, tx: IngressSender, local_id: PeerId) {
2039 let mut handler = IngressHandler::new(tx, local_id);
2040 let mut buf = [0; MAX_PACKET_SIZE];
2041 loop {
2042 let res = udp.recv_from(&mut buf).await;
2043 match res {
2044 Err(err) => {
2045 debug!(target: "discv4", %err, "Failed to read datagram.");
2046 handler.send(IngressEvent::RecvError(err)).await;
2047 }
2048 Ok((read, remote_addr)) => {
2049 handler.handle_packet(&buf[..read], remote_addr).await;
2050 }
2051 }
2052 }
2053}
2054
2055#[derive(Debug)]
2060pub struct IngressHandler {
2061 tx: IngressSender,
2062 local_id: PeerId,
2063 tick: usize,
2064 tick_interval: Duration,
2065 cache: ReceiveCache,
2066 last_tick: Instant,
2067}
2068
2069impl IngressHandler {
2070 fn new(tx: IngressSender, local_id: PeerId) -> Self {
2071 let tick = MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP / 2;
2072 Self {
2073 tx,
2074 local_id,
2075 tick,
2076 tick_interval: Duration::from_secs(tick as u64),
2077 cache: ReceiveCache::default(),
2078 last_tick: Instant::now(),
2079 }
2080 }
2081
2082 async fn send(&self, event: IngressEvent) {
2083 let _ = self.tx.send(event).await.map_err(|err| {
2084 debug!(target: "discv4", %err, "failed send incoming packet");
2085 });
2086 }
2087
2088 pub async fn handle_packet(&mut self, data: &[u8], src: SocketAddr) {
2091 if self.last_tick.elapsed() >= self.tick_interval {
2092 self.cache.tick_ips(self.tick);
2093 self.last_tick = Instant::now();
2094 }
2095
2096 if self.cache.inc_ip(src.ip()) > MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP {
2098 trace!(target: "discv4", ?src, "Too many incoming packets from IP.");
2099 return
2100 }
2101
2102 if data.len() >= MIN_PACKET_SIZE &&
2108 self.cache.contains_packet(B256::from_slice(&data[..32]))
2109 {
2110 trace!(target: "discv4", ?src, "Received duplicate packet.");
2111 return
2112 }
2113
2114 let event = match Message::decode(data) {
2115 Ok(packet) => {
2116 if packet.node_id == self.local_id {
2117 debug!(target: "discv4", ?src, "Received own packet.");
2118 return
2119 }
2120
2121 self.cache.insert_packet(packet.hash);
2124
2125 IngressEvent::Packet(src, packet)
2126 }
2127 Err(err) => {
2128 trace!(target: "discv4", %err, "Failed to decode packet");
2129 IngressEvent::BadPacket(src, err, data.to_vec())
2130 }
2131 };
2132
2133 self.send(event).await;
2134 }
2135}
2136
2137#[derive(Debug)]
2141struct ReceiveCache {
2142 ip_messages: HashMap<IpAddr, usize>,
2148 unique_packets: schnellru::LruMap<B256, ()>,
2150}
2151
2152impl ReceiveCache {
2153 fn tick_ips(&mut self, tick: usize) {
2157 self.ip_messages.retain(|_, count| {
2158 if let Some(reset) = count.checked_sub(tick) {
2159 *count = reset;
2160 true
2161 } else {
2162 false
2163 }
2164 });
2165 }
2166
2167 fn inc_ip(&mut self, ip: IpAddr) -> usize {
2169 let ctn = self.ip_messages.entry(ip).or_default();
2170 *ctn = ctn.saturating_add(1);
2171 *ctn
2172 }
2173
2174 fn contains_packet(&mut self, hash: B256) -> bool {
2178 self.unique_packets.get(&hash).is_some()
2179 }
2180
2181 fn insert_packet(&mut self, hash: B256) {
2183 self.unique_packets.insert(hash, ());
2184 }
2185}
2186
2187impl Default for ReceiveCache {
2188 fn default() -> Self {
2189 Self {
2190 ip_messages: Default::default(),
2191 unique_packets: schnellru::LruMap::new(schnellru::ByLength::new(32)),
2192 }
2193 }
2194}
2195
2196enum Discv4Command {
2198 Add(NodeRecord),
2199 AddBootNode(NodeRecord),
2200 SetTcpPort(u16),
2201 SetEIP868RLPPair { key: Vec<u8>, rlp: Bytes },
2202 Ban(PeerId, IpAddr),
2203 BanPeer(PeerId),
2204 BanIp(IpAddr),
2205 Remove(PeerId),
2206 Lookup { node_id: Option<PeerId>, tx: Option<NodeRecordSender> },
2207 SetLookupInterval(Duration),
2208 Updates(OneshotSender<ReceiverStream<DiscoveryUpdate>>),
2209 Terminated,
2210}
2211
2212#[derive(Debug)]
2214pub(crate) enum IngressEvent {
2215 RecvError(io::Error),
2217 BadPacket(SocketAddr, DecodePacketError, Vec<u8>),
2219 Packet(SocketAddr, Packet),
2221}
2222
2223#[derive(Debug)]
2225struct PingRequest {
2226 sent_at: Instant,
2228 node: NodeRecord,
2230 echo_hash: B256,
2232 reason: PingReason,
2234}
2235
2236#[derive(Debug)]
2240struct LookupTargetRotator {
2241 interval: usize,
2242 counter: usize,
2243}
2244
2245impl LookupTargetRotator {
2248 const fn local_only() -> Self {
2250 Self { interval: 1, counter: 0 }
2251 }
2252}
2253
2254impl Default for LookupTargetRotator {
2255 fn default() -> Self {
2256 Self {
2257 interval: 4,
2259 counter: 3,
2260 }
2261 }
2262}
2263
2264impl LookupTargetRotator {
2265 fn next(&mut self, local: &PeerId) -> PeerId {
2267 self.counter += 1;
2268 self.counter %= self.interval;
2269 if self.counter == 0 {
2270 return *local
2271 }
2272 PeerId::random()
2273 }
2274}
2275
2276#[derive(Clone, Debug)]
2281struct LookupContext {
2282 inner: Rc<LookupContextInner>,
2283}
2284
2285impl LookupContext {
2286 fn new(
2288 target: discv5::Key<NodeKey>,
2289 nearest_nodes: impl IntoIterator<Item = (Distance, NodeRecord)>,
2290 listener: Option<NodeRecordSender>,
2291 ) -> Self {
2292 let closest_nodes = nearest_nodes
2293 .into_iter()
2294 .map(|(distance, record)| {
2295 (distance, QueryNode { record, queried: false, responded: false })
2296 })
2297 .collect();
2298
2299 let inner = Rc::new(LookupContextInner {
2300 target,
2301 closest_nodes: RefCell::new(closest_nodes),
2302 listener,
2303 });
2304 Self { inner }
2305 }
2306
2307 fn target(&self) -> PeerId {
2309 self.inner.target.preimage().0
2310 }
2311
2312 fn closest(&self, num: usize) -> Vec<NodeRecord> {
2313 self.inner
2314 .closest_nodes
2315 .borrow()
2316 .iter()
2317 .filter(|(_, node)| !node.queried)
2318 .map(|(_, n)| n.record)
2319 .take(num)
2320 .collect()
2321 }
2322
2323 fn filter_closest<P>(&self, num: usize, filter: P) -> Vec<NodeRecord>
2325 where
2326 P: FnMut(&NodeRecord) -> bool,
2327 {
2328 self.inner
2329 .closest_nodes
2330 .borrow()
2331 .iter()
2332 .filter(|(_, node)| !node.queried)
2333 .map(|(_, n)| n.record)
2334 .filter(filter)
2335 .take(num)
2336 .collect()
2337 }
2338
2339 fn add_node(&self, record: NodeRecord) {
2341 let distance = self.inner.target.distance(&kad_key(record.id));
2342 let mut closest = self.inner.closest_nodes.borrow_mut();
2343 if let btree_map::Entry::Vacant(entry) = closest.entry(distance) {
2344 entry.insert(QueryNode { record, queried: false, responded: false });
2345 }
2346 }
2347
2348 fn set_queried(&self, id: PeerId, val: bool) {
2349 if let Some((_, node)) =
2350 self.inner.closest_nodes.borrow_mut().iter_mut().find(|(_, node)| node.record.id == id)
2351 {
2352 node.queried = val;
2353 }
2354 }
2355
2356 fn mark_queried(&self, id: PeerId) {
2358 self.set_queried(id, true)
2359 }
2360
2361 fn unmark_queried(&self, id: PeerId) {
2363 self.set_queried(id, false)
2364 }
2365
2366 fn mark_responded(&self, id: PeerId) {
2368 if let Some((_, node)) =
2369 self.inner.closest_nodes.borrow_mut().iter_mut().find(|(_, node)| node.record.id == id)
2370 {
2371 node.responded = true;
2372 }
2373 }
2374}
2375
2376unsafe impl Send for LookupContext {}
2383#[derive(Debug)]
2384struct LookupContextInner {
2385 target: discv5::Key<NodeKey>,
2387 closest_nodes: RefCell<BTreeMap<Distance, QueryNode>>,
2389 listener: Option<NodeRecordSender>,
2394}
2395
2396impl Drop for LookupContextInner {
2397 fn drop(&mut self) {
2398 if let Some(tx) = self.listener.take() {
2399 let nodes = self
2402 .closest_nodes
2403 .take()
2404 .into_values()
2405 .filter(|node| node.responded)
2406 .map(|node| node.record)
2407 .collect();
2408 let _ = tx.send(nodes);
2409 }
2410 }
2411}
2412
2413#[derive(Debug, Clone, Copy)]
2415struct QueryNode {
2416 record: NodeRecord,
2417 queried: bool,
2418 responded: bool,
2419}
2420
2421#[derive(Debug)]
2422struct FindNodeRequest {
2423 sent_at: Instant,
2425 response_count: usize,
2427 answered: bool,
2429 lookup_context: LookupContext,
2431}
2432
2433impl FindNodeRequest {
2436 fn new(resp: LookupContext) -> Self {
2437 Self { sent_at: Instant::now(), response_count: 0, answered: false, lookup_context: resp }
2438 }
2439}
2440
2441#[derive(Debug)]
2443struct CachedFindNode {
2444 target: PeerId,
2445 payload: Bytes,
2446 hash: B256,
2447 cached_at: Instant,
2448}
2449
2450impl CachedFindNode {
2451 fn get_or_sign(
2454 cache: &mut Option<Self>,
2455 target: PeerId,
2456 ttl: Duration,
2457 secret_key: &secp256k1::SecretKey,
2458 expire: u64,
2459 ) -> (Bytes, B256) {
2460 if let Some(c) = cache.as_ref() &&
2461 c.target == target &&
2462 c.cached_at.elapsed() < ttl
2463 {
2464 return (c.payload.clone(), c.hash);
2465 }
2466
2467 let msg = Message::FindNode(FindNode { id: target, expire });
2468 let (payload, hash) = msg.encode(secret_key);
2469
2470 *cache = Some(Self { target, payload: payload.clone(), hash, cached_at: Instant::now() });
2471
2472 (payload, hash)
2473 }
2474}
2475
2476#[derive(Debug)]
2477struct EnrRequestState {
2478 sent_at: Instant,
2480 echo_hash: B256,
2482}
2483
2484#[derive(Debug, Clone, Eq, PartialEq)]
2486struct NodeEntry {
2487 record: NodeRecord,
2489 last_seen: Instant,
2491 last_enr_seq: Option<u64>,
2493 fork_id: Option<ForkId>,
2495 find_node_failures: u8,
2497 has_endpoint_proof: bool,
2499}
2500
2501impl NodeEntry {
2504 fn new(record: NodeRecord) -> Self {
2506 Self {
2507 record,
2508 last_seen: Instant::now(),
2509 last_enr_seq: None,
2510 fork_id: None,
2511 find_node_failures: 0,
2512 has_endpoint_proof: false,
2513 }
2514 }
2515
2516 #[cfg(test)]
2517 fn new_proven(record: NodeRecord) -> Self {
2518 let mut node = Self::new(record);
2519 node.has_endpoint_proof = true;
2520 node
2521 }
2522
2523 const fn establish_proof(&mut self) {
2525 self.has_endpoint_proof = true;
2526 self.find_node_failures = 0;
2527 }
2528
2529 const fn exceeds_find_node_failures(&self, max_failures: u8) -> bool {
2531 self.find_node_failures >= max_failures
2532 }
2533
2534 fn update_with_enr(&mut self, last_enr_seq: Option<u64>) -> Option<u64> {
2536 self.update_now(|s| std::mem::replace(&mut s.last_enr_seq, last_enr_seq))
2537 }
2538
2539 const fn inc_failed_request(&mut self) {
2541 self.find_node_failures += 1;
2542 }
2543
2544 fn update_with_fork_id(&mut self, fork_id: Option<ForkId>) -> Option<ForkId> {
2546 self.update_now(|s| std::mem::replace(&mut s.fork_id, fork_id))
2547 }
2548
2549 fn update_now<F, R>(&mut self, f: F) -> R
2551 where
2552 F: FnOnce(&mut Self) -> R,
2553 {
2554 self.last_seen = Instant::now();
2555 f(self)
2556 }
2557}
2558
2559impl NodeEntry {
2562 fn is_expired(&self) -> bool {
2564 self.last_seen.elapsed() > (ENDPOINT_PROOF_EXPIRATION / 2)
2565 }
2566}
2567
2568#[derive(Debug)]
2570enum PingReason {
2571 InitialInsert,
2573 EstablishBond,
2575 RePing,
2577 Lookup(NodeRecord, LookupContext),
2579}
2580
2581#[derive(Debug, Clone)]
2583pub enum DiscoveryUpdate {
2584 Added(NodeRecord),
2586 DiscoveredAtCapacity(NodeRecord),
2588 EnrForkId(NodeRecord, ForkId),
2590 Removed(PeerId),
2592 Batch(Vec<Self>),
2594}
2595
2596#[cfg(test)]
2597mod tests {
2598 use super::*;
2599 use crate::test_utils::{create_discv4, create_discv4_with_config, rng_endpoint, rng_record};
2600 use alloy_primitives::hex;
2601 use alloy_rlp::{Decodable, Encodable};
2602 use rand_08::Rng;
2603 use reth_ethereum_forks::{EnrForkIdEntry, ForkHash};
2604 use reth_network_peers::mainnet_nodes;
2605 use secp256k1::SECP256K1;
2606 use std::future::poll_fn;
2607
2608 #[tokio::test]
2609 async fn test_duplicate_packet_rejected_without_decoding() {
2610 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
2611 let local_id = pk2id(&secret_key.public_key(SECP256K1));
2612 let (tx, mut rx) = mpsc::channel(16);
2613 let mut handler = IngressHandler::new(tx, local_id);
2614
2615 let remote_key = SecretKey::new(&mut rand_08::thread_rng());
2616 let msg = Message::Ping(Ping {
2617 from: rng_endpoint(&mut rand_08::thread_rng()),
2618 to: rng_endpoint(&mut rand_08::thread_rng()),
2619 expire: u64::MAX,
2620 enr_sq: None,
2621 });
2622 let (packet, hash) = msg.encode(&remote_key);
2623 let src = "10.0.0.1:30303".parse().unwrap();
2624
2625 handler.handle_packet(&packet, src).await;
2626 assert!(matches!(rx.try_recv(), Ok(IngressEvent::Packet(_, _))));
2627
2628 handler.handle_packet(&packet, src).await;
2630 assert!(rx.try_recv().is_err());
2631 assert!(handler.cache.contains_packet(hash));
2632 }
2633
2634 #[tokio::test]
2635 async fn test_undecodable_packet_does_not_poison_cache() {
2636 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
2637 let local_id = pk2id(&secret_key.public_key(SECP256K1));
2638 let (tx, mut rx) = mpsc::channel(16);
2639 let mut handler = IngressHandler::new(tx, local_id);
2640
2641 let remote_key = SecretKey::new(&mut rand_08::thread_rng());
2642 let msg = Message::Ping(Ping {
2643 from: rng_endpoint(&mut rand_08::thread_rng()),
2644 to: rng_endpoint(&mut rand_08::thread_rng()),
2645 expire: u64::MAX,
2646 enr_sq: None,
2647 });
2648 let (packet, hash) = msg.encode(&remote_key);
2649 let src = "10.0.0.1:30303".parse().unwrap();
2650
2651 let mut forged = packet.to_vec();
2654 let last = forged.len() - 1;
2655 forged[last] ^= 0xff;
2656 handler.handle_packet(&forged, src).await;
2657 assert!(matches!(rx.try_recv(), Ok(IngressEvent::BadPacket(..))));
2658 assert!(!handler.cache.contains_packet(hash));
2659
2660 handler.handle_packet(&packet, src).await;
2662 assert!(matches!(rx.try_recv(), Ok(IngressEvent::Packet(_, _))));
2663 }
2664
2665 #[tokio::test]
2666 async fn test_configured_enr_forkid_entry() {
2667 let fork: ForkId = ForkId { hash: ForkHash([220, 233, 108, 45]), next: 0u64 };
2668 let mut disc_conf = Discv4Config::default();
2669 disc_conf.add_eip868_pair("eth", EnrForkIdEntry::from(fork));
2670 let (_discv4, service) = create_discv4_with_config(disc_conf).await;
2671 let eth = service.local_eip_868_enr.get_raw_rlp(b"eth").unwrap();
2672 let fork_entry_id = EnrForkIdEntry::decode(&mut ð[..]).unwrap();
2673
2674 let raw: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2675 let decoded = EnrForkIdEntry::decode(&mut &raw[..]).unwrap();
2676 let expected = EnrForkIdEntry {
2677 fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2678 };
2679 assert_eq!(expected, fork_entry_id);
2680 assert_eq!(expected, decoded);
2681 }
2682
2683 #[test]
2684 fn test_enr_forkid_entry_decode() {
2685 let raw: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2686 let decoded = EnrForkIdEntry::decode(&mut &raw[..]).unwrap();
2687 let expected = EnrForkIdEntry {
2688 fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2689 };
2690 assert_eq!(expected, decoded);
2691 }
2692
2693 #[test]
2694 fn test_enr_forkid_entry_encode() {
2695 let original = EnrForkIdEntry {
2696 fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2697 };
2698 let expected: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2699 let mut encoded = Vec::with_capacity(expected.len());
2700 original.encode(&mut encoded);
2701 assert_eq!(&expected[..], encoded.as_slice());
2702 }
2703
2704 #[tokio::test]
2705 async fn test_request_enr_of_proven_node() {
2706 reth_tracing::init_test_tracing();
2707 let mut rng = rand_08::thread_rng();
2708 let (_discv4, mut service) = create_discv4().await;
2709
2710 let id = PeerId::random();
2711 let addr: SocketAddr = (Ipv4Addr::new(10, 0, 0, 1), DEFAULT_DISCOVERY_PORT).into();
2712 let mut entry = NodeEntry::new_proven(NodeRecord::new(addr, id));
2713 entry.last_enr_seq = Some(1);
2714 let _ = service.kbuckets.insert_or_update(
2715 &kad_key(id),
2716 entry,
2717 NodeStatus {
2718 direction: ConnectionDirection::Outgoing,
2719 state: ConnectionState::Connected,
2720 },
2721 );
2722 let ping = |service: &Discv4Service, rng: &mut rand_08::rngs::ThreadRng| Ping {
2723 from: rng_endpoint(rng),
2724 to: rng_endpoint(rng),
2725 expire: service.ping_expiration(),
2726 enr_sq: Some(2),
2727 };
2728
2729 service.on_ping(ping(&service, &mut rng), addr, id, B256::random());
2731 assert!(service.pending_enr_requests.contains_key(&id));
2732
2733 service.evict_expired_requests(Instant::now() + service.config.enr_expiration * 2);
2735 assert!(!service.pending_enr_requests.contains_key(&id));
2736 service.on_ping(ping(&service, &mut rng), addr, id, B256::random());
2737 assert!(service.pending_enr_requests.contains_key(&id));
2738 }
2739
2740 #[test]
2741 fn test_local_rotator() {
2742 let id = PeerId::random();
2743 let mut rotator = LookupTargetRotator::local_only();
2744 assert_eq!(rotator.next(&id), id);
2745 assert_eq!(rotator.next(&id), id);
2746 }
2747
2748 #[test]
2749 fn test_rotator() {
2750 let id = PeerId::random();
2751 let mut rotator = LookupTargetRotator::default();
2752 assert_eq!(rotator.next(&id), id);
2753 assert_ne!(rotator.next(&id), id);
2754 assert_ne!(rotator.next(&id), id);
2755 assert_ne!(rotator.next(&id), id);
2756 assert_eq!(rotator.next(&id), id);
2757 }
2758
2759 #[tokio::test]
2760 async fn test_pending_ping() {
2761 let (_, mut service) = create_discv4().await;
2762
2763 let local_addr = service.local_addr();
2764
2765 let mut num_inserted = 0;
2766 loop {
2767 let node = NodeRecord::new(local_addr, PeerId::random());
2768 if service.add_node(node) {
2769 num_inserted += 1;
2770 assert!(service.pending_pings.contains_key(&node.id));
2771 assert_eq!(service.pending_pings.len(), num_inserted);
2772 if num_inserted == MAX_NODES_PING {
2773 break
2774 }
2775 }
2776 }
2777
2778 num_inserted = 0;
2780 for _ in 0..MAX_NODES_PING {
2781 let node = NodeRecord::new(local_addr, PeerId::random());
2782 if service.add_node(node) {
2783 num_inserted += 1;
2784 assert!(!service.pending_pings.contains_key(&node.id));
2785 assert_eq!(service.pending_pings.len(), MAX_NODES_PING);
2786 assert_eq!(service.queued_pings.len(), num_inserted);
2787 }
2788 }
2789 }
2790
2791 #[tokio::test(flavor = "multi_thread")]
2793 #[ignore]
2794 async fn test_mainnet_lookup() {
2795 reth_tracing::init_test_tracing();
2796 let fork_id = ForkId { hash: ForkHash(hex!("743f3d89")), next: 16191202 };
2797
2798 let all_nodes = mainnet_nodes();
2799 let config = Discv4Config::builder()
2800 .add_boot_nodes(all_nodes)
2801 .lookup_interval(Duration::from_secs(1))
2802 .add_eip868_pair("eth", fork_id)
2803 .build();
2804 let (_discv4, mut service) = create_discv4_with_config(config).await;
2805
2806 let mut updates = service.update_stream();
2807
2808 let _handle = service.spawn();
2809
2810 let mut table = HashMap::new();
2811 while let Some(update) = updates.next().await {
2812 match update {
2813 DiscoveryUpdate::EnrForkId(record, fork_id) => {
2814 println!("{record:?}, {fork_id:?}");
2815 }
2816 DiscoveryUpdate::Added(record) => {
2817 table.insert(record.id, record);
2818 }
2819 DiscoveryUpdate::Removed(id) => {
2820 table.remove(&id);
2821 }
2822 _ => {}
2823 }
2824 println!("total peers {}", table.len());
2825 }
2826 }
2827
2828 #[tokio::test]
2829 async fn test_mapped_ipv4() {
2830 reth_tracing::init_test_tracing();
2831 let mut rng = rand_08::thread_rng();
2832 let config = Discv4Config::builder().build();
2833 let (_discv4, mut service) = create_discv4_with_config(config).await;
2834
2835 let v4: Ipv4Addr = "0.0.0.0".parse().unwrap();
2836 let v6 = v4.to_ipv6_mapped();
2837 let addr: SocketAddr = (v6, DEFAULT_DISCOVERY_PORT).into();
2838
2839 let ping = Ping {
2840 from: rng_endpoint(&mut rng),
2841 to: rng_endpoint(&mut rng),
2842 expire: service.ping_expiration(),
2843 enr_sq: Some(rng.r#gen()),
2844 };
2845
2846 let id = PeerId::random();
2847 service.on_ping(ping, addr, id, B256::random());
2848
2849 let key = kad_key(id);
2850 match service.kbuckets.entry(&key) {
2851 kbucket::Entry::Present(entry, _) => {
2852 let node_addr = entry.value().record.address;
2853 assert!(node_addr.is_ipv4());
2854 assert_eq!(node_addr, IpAddr::from(v4));
2855 }
2856 _ => unreachable!(),
2857 };
2858 }
2859
2860 #[tokio::test]
2861 async fn test_respect_ping_expiration() {
2862 reth_tracing::init_test_tracing();
2863 let mut rng = rand_08::thread_rng();
2864 let config = Discv4Config::builder().build();
2865 let (_discv4, mut service) = create_discv4_with_config(config).await;
2866
2867 let v4: Ipv4Addr = "0.0.0.0".parse().unwrap();
2868 let v6 = v4.to_ipv6_mapped();
2869 let addr: SocketAddr = (v6, DEFAULT_DISCOVERY_PORT).into();
2870
2871 let ping = Ping {
2872 from: rng_endpoint(&mut rng),
2873 to: rng_endpoint(&mut rng),
2874 expire: SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() - 1,
2875 enr_sq: Some(rng.r#gen()),
2876 };
2877
2878 let id = PeerId::random();
2879 service.on_ping(ping, addr, id, B256::random());
2880
2881 let key = kad_key(id);
2882 match service.kbuckets.entry(&key) {
2883 kbucket::Entry::Absent(_) => {}
2884 _ => unreachable!(),
2885 };
2886 }
2887
2888 #[tokio::test]
2889 async fn test_single_lookups() {
2890 reth_tracing::init_test_tracing();
2891
2892 let config = Discv4Config::builder().build();
2893 let (_discv4, mut service) = create_discv4_with_config(config.clone()).await;
2894
2895 let id = PeerId::random();
2896 let key = kad_key(id);
2897 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
2898
2899 let _ = service.kbuckets.insert_or_update(
2900 &key,
2901 NodeEntry::new_proven(record),
2902 NodeStatus {
2903 direction: ConnectionDirection::Incoming,
2904 state: ConnectionState::Connected,
2905 },
2906 );
2907
2908 service.lookup_self();
2909 assert_eq!(service.pending_find_nodes.len(), 1);
2910
2911 poll_fn(|cx| {
2912 let _ = service.poll(cx);
2913 assert_eq!(service.pending_find_nodes.len(), 1);
2914
2915 Poll::Ready(())
2916 })
2917 .await;
2918 }
2919
2920 #[tokio::test]
2921 async fn test_on_neighbours_recursive_lookup() {
2922 reth_tracing::init_test_tracing();
2923
2924 let config = Discv4Config::builder().build();
2925 let (_discv4, mut service) = create_discv4_with_config(config.clone()).await;
2926 let (_discv4, mut service2) = create_discv4_with_config(config).await;
2927
2928 let id = PeerId::random();
2929 let key = kad_key(id);
2930 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
2931
2932 let _ = service.kbuckets.insert_or_update(
2933 &key,
2934 NodeEntry::new_proven(record),
2935 NodeStatus {
2936 direction: ConnectionDirection::Incoming,
2937 state: ConnectionState::Connected,
2938 },
2939 );
2940 service.lookup_self();
2943 assert_eq!(service.pending_find_nodes.len(), 1);
2944
2945 poll_fn(|cx| {
2946 let _ = service.poll(cx);
2947 assert_eq!(service.pending_find_nodes.len(), 1);
2948
2949 Poll::Ready(())
2950 })
2951 .await;
2952
2953 let expiry = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs() +
2954 10000000000000;
2955 let msg = Neighbours { nodes: vec![service2.local_node_record], expire: expiry };
2956 service.on_neighbours(msg, record.tcp_addr(), id);
2957 let event = poll_fn(|cx| service2.poll(cx)).await;
2959 assert_eq!(event, Discv4Event::Ping);
2960 assert_eq!(service.pending_find_nodes.len(), 1);
2963 let event = poll_fn(|cx| service.poll(cx)).await;
2965 assert_eq!(event, Discv4Event::Pong);
2966 let event = poll_fn(|cx| service.poll(cx)).await;
2971 assert_eq!(event, Discv4Event::Ping);
2972 assert_eq!(service.pending_find_nodes.len(), 2);
2975 }
2976
2977 #[tokio::test]
2978 async fn test_no_local_in_closest() {
2979 reth_tracing::init_test_tracing();
2980
2981 let config = Discv4Config::builder().build();
2982 let (_discv4, mut service) = create_discv4_with_config(config).await;
2983
2984 let target_key = kad_key(PeerId::random());
2985
2986 let id = PeerId::random();
2987 let key = kad_key(id);
2988 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
2989
2990 let _ = service.kbuckets.insert_or_update(
2991 &key,
2992 NodeEntry::new(record),
2993 NodeStatus {
2994 direction: ConnectionDirection::Incoming,
2995 state: ConnectionState::Connected,
2996 },
2997 );
2998
2999 let closest = service
3000 .kbuckets
3001 .closest_values(&target_key)
3002 .map(|n| n.value.record)
3003 .take(MAX_NODES_PER_BUCKET)
3004 .collect::<Vec<_>>();
3005
3006 assert_eq!(closest.len(), 1);
3007 assert!(!closest.iter().any(|r| r.id == *service.local_peer_id()));
3008 }
3009
3010 #[tokio::test]
3011 async fn test_random_lookup() {
3012 reth_tracing::init_test_tracing();
3013
3014 let config = Discv4Config::builder().build();
3015 let (_discv4, mut service) = create_discv4_with_config(config).await;
3016
3017 let target = PeerId::random();
3018
3019 let id = PeerId::random();
3020 let key = kad_key(id);
3021 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3022
3023 let _ = service.kbuckets.insert_or_update(
3024 &key,
3025 NodeEntry::new_proven(record),
3026 NodeStatus {
3027 direction: ConnectionDirection::Incoming,
3028 state: ConnectionState::Connected,
3029 },
3030 );
3031
3032 service.lookup(target);
3033 assert_eq!(service.pending_find_nodes.len(), 1);
3034
3035 let ctx = service.pending_find_nodes.values().next().unwrap().lookup_context.clone();
3036
3037 assert_eq!(ctx.target(), target);
3038 assert_eq!(ctx.inner.closest_nodes.borrow().len(), 1);
3039
3040 ctx.add_node(record);
3041 assert_eq!(ctx.inner.closest_nodes.borrow().len(), 1);
3042 }
3043
3044 #[tokio::test]
3045 async fn test_reping_on_find_node_failures() {
3046 reth_tracing::init_test_tracing();
3047
3048 let config = Discv4Config::builder().build();
3049 let (_discv4, mut service) = create_discv4_with_config(config).await;
3050
3051 let target = PeerId::random();
3052
3053 let id = PeerId::random();
3054 let key = kad_key(id);
3055 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3056
3057 let mut entry = NodeEntry::new_proven(record);
3058 entry.find_node_failures = u8::MAX;
3059 let _ = service.kbuckets.insert_or_update(
3060 &key,
3061 entry,
3062 NodeStatus {
3063 direction: ConnectionDirection::Incoming,
3064 state: ConnectionState::Connected,
3065 },
3066 );
3067
3068 service.lookup(target);
3069 assert_eq!(service.pending_find_nodes.len(), 0);
3070 assert_eq!(service.pending_pings.len(), 1);
3071
3072 service.update_on_pong(record, None);
3073
3074 service
3075 .on_entry(record.id, |entry| {
3076 assert_eq!(entry.find_node_failures, 0);
3078 assert!(entry.has_endpoint_proof);
3079 })
3080 .unwrap();
3081 }
3082
3083 #[tokio::test]
3084 async fn test_service_commands() {
3085 reth_tracing::init_test_tracing();
3086
3087 let config = Discv4Config::builder().build();
3088 let (discv4, mut service) = create_discv4_with_config(config).await;
3089
3090 service.lookup_self();
3091
3092 let _handle = service.spawn();
3093 discv4.send_lookup_self();
3094 let _ = discv4.lookup_self().await;
3095 }
3096
3097 #[tokio::test]
3098 async fn test_requests_timeout() {
3099 reth_tracing::init_test_tracing();
3100 let fork_id = ForkId { hash: ForkHash(hex!("743f3d89")), next: 16191202 };
3101
3102 let config = Discv4Config::builder()
3103 .request_timeout(Duration::from_millis(200))
3104 .ping_expiration(Duration::from_millis(200))
3105 .lookup_neighbours_expiration(Duration::from_millis(200))
3106 .add_eip868_pair("eth", fork_id)
3107 .build();
3108 let (_disv4, mut service) = create_discv4_with_config(config).await;
3109
3110 let id = PeerId::random();
3111 let key = kad_key(id);
3112 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3113
3114 let _ = service.kbuckets.insert_or_update(
3115 &key,
3116 NodeEntry::new_proven(record),
3117 NodeStatus {
3118 direction: ConnectionDirection::Incoming,
3119 state: ConnectionState::Connected,
3120 },
3121 );
3122
3123 service.lookup_self();
3124 assert_eq!(service.pending_find_nodes.len(), 1);
3125
3126 let ctx = service.pending_find_nodes.values().next().unwrap().lookup_context.clone();
3127
3128 service.pending_lookup.insert(record.id, (Instant::now(), ctx));
3129
3130 assert_eq!(service.pending_lookup.len(), 1);
3131
3132 let ping = Ping {
3133 from: service.local_node_record.into(),
3134 to: record.into(),
3135 expire: service.ping_expiration(),
3136 enr_sq: service.enr_seq(),
3137 };
3138 let echo_hash = service.send_packet(Message::Ping(ping), record.udp_addr());
3139 let ping_request = PingRequest {
3140 sent_at: Instant::now(),
3141 node: record,
3142 echo_hash,
3143 reason: PingReason::InitialInsert,
3144 };
3145 service.pending_pings.insert(record.id, ping_request);
3146
3147 assert_eq!(service.pending_pings.len(), 1);
3148
3149 tokio::time::sleep(Duration::from_secs(1)).await;
3150
3151 poll_fn(|cx| {
3152 let _ = service.poll(cx);
3153
3154 assert_eq!(service.pending_find_nodes.len(), 0);
3155 assert_eq!(service.pending_lookup.len(), 0);
3156 assert_eq!(service.pending_pings.len(), 0);
3157
3158 Poll::Ready(())
3159 })
3160 .await;
3161 }
3162
3163 #[tokio::test(flavor = "multi_thread")]
3165 async fn test_check_wrong_to() {
3166 reth_tracing::init_test_tracing();
3167
3168 let config = Discv4Config::builder().external_ip_resolver(None).build();
3169 let (_discv4, mut service_1) = create_discv4_with_config(config.clone()).await;
3170 let (_discv4, mut service_2) = create_discv4_with_config(config).await;
3171
3172 let mut ping = Ping {
3174 from: service_1.local_node_record.into(),
3175 to: service_2.local_node_record.into(),
3176 expire: service_1.ping_expiration(),
3177 enr_sq: service_1.enr_seq(),
3178 };
3179 ping.to.address = "192.0.2.0".parse().unwrap();
3180
3181 let echo_hash = service_1.send_packet(Message::Ping(ping), service_2.local_addr());
3182 let ping_request = PingRequest {
3183 sent_at: Instant::now(),
3184 node: service_2.local_node_record,
3185 echo_hash,
3186 reason: PingReason::InitialInsert,
3187 };
3188 service_1.pending_pings.insert(*service_2.local_peer_id(), ping_request);
3189
3190 let event = poll_fn(|cx| service_2.poll(cx)).await;
3192 assert_eq!(event, Discv4Event::Ping);
3193
3194 let event = poll_fn(|cx| service_1.poll(cx)).await;
3196 assert_eq!(event, Discv4Event::Pong);
3197 let event = poll_fn(|cx| service_1.poll(cx)).await;
3199 assert_eq!(event, Discv4Event::Ping);
3200 }
3201
3202 #[tokio::test(flavor = "multi_thread")]
3203 async fn test_check_ping_pong() {
3204 reth_tracing::init_test_tracing();
3205
3206 let config = Discv4Config::builder().external_ip_resolver(None).build();
3207 let (_discv4, mut service_1) = create_discv4_with_config(config.clone()).await;
3208 let (_discv4, mut service_2) = create_discv4_with_config(config).await;
3209
3210 service_1.add_node(service_2.local_node_record);
3212
3213 let event = poll_fn(|cx| service_2.poll(cx)).await;
3215 assert_eq!(event, Discv4Event::Ping);
3216
3217 let key1 = kad_key(*service_1.local_peer_id());
3219 match service_2.kbuckets.entry(&key1) {
3220 kbucket::Entry::Present(_entry, status) => {
3221 assert!(!status.is_connected());
3222 }
3223 _ => unreachable!(),
3224 }
3225
3226 let event = poll_fn(|cx| service_1.poll(cx)).await;
3228 assert_eq!(event, Discv4Event::Pong);
3229
3230 let key2 = kad_key(*service_2.local_peer_id());
3232 match service_1.kbuckets.entry(&key2) {
3233 kbucket::Entry::Present(_entry, status) => {
3234 assert!(status.is_connected());
3235 }
3236 _ => unreachable!(),
3237 }
3238
3239 let event = poll_fn(|cx| service_1.poll(cx)).await;
3241 assert_eq!(event, Discv4Event::Ping);
3242
3243 tokio::time::timeout(Duration::from_secs(5), async {
3247 loop {
3248 let event = poll_fn(|cx| service_2.poll(cx)).await;
3249 match event {
3250 Discv4Event::Pong => break,
3251 Discv4Event::EnrRequest | Discv4Event::FindNode => {}
3252 ev => unreachable!("{ev:?}"),
3253 }
3254 }
3255 })
3256 .await
3257 .expect("timed out waiting for Pong from service_2");
3258
3259 match service_2.kbuckets.entry(&key1) {
3261 kbucket::Entry::Present(_entry, status) => {
3262 assert!(status.is_connected());
3263 }
3264 ev => unreachable!("{ev:?}"),
3265 }
3266 }
3267
3268 #[test]
3269 fn test_insert() {
3270 let local_node_record = rng_record(&mut rand_08::thread_rng());
3271 let mut kbuckets: KBucketsTable<NodeKey, NodeEntry> = KBucketsTable::new(
3272 NodeKey::from(&local_node_record).into(),
3273 Duration::from_secs(60),
3274 MAX_NODES_PER_BUCKET,
3275 None,
3276 None,
3277 );
3278
3279 let new_record = rng_record(&mut rand_08::thread_rng());
3280 let key = kad_key(new_record.id);
3281 match kbuckets.entry(&key) {
3282 kbucket::Entry::Absent(entry) => {
3283 let node = NodeEntry::new(new_record);
3284 let _ = entry.insert(
3285 node,
3286 NodeStatus {
3287 direction: ConnectionDirection::Outgoing,
3288 state: ConnectionState::Disconnected,
3289 },
3290 );
3291 }
3292 _ => {
3293 unreachable!()
3294 }
3295 };
3296 match kbuckets.entry(&key) {
3297 kbucket::Entry::Present(_, _) => {}
3298 _ => {
3299 unreachable!()
3300 }
3301 }
3302 }
3303
3304 #[tokio::test]
3305 async fn test_bootnode_not_in_update_stream() {
3306 reth_tracing::init_test_tracing();
3307 let (_, service_1) = create_discv4().await;
3308 let peerid_1 = *service_1.local_peer_id();
3309
3310 let config = Discv4Config::builder().add_boot_node(service_1.local_node_record).build();
3311 service_1.spawn();
3312
3313 let (_, mut service_2) = create_discv4_with_config(config).await;
3314
3315 let mut updates = service_2.update_stream();
3316
3317 service_2.spawn();
3318
3319 let mut bootnode_appeared = false;
3321 let timeout = tokio::time::sleep(Duration::from_secs(1));
3322 tokio::pin!(timeout);
3323
3324 loop {
3325 tokio::select! {
3326 Some(update) = updates.next() => {
3327 if let DiscoveryUpdate::Added(record) = update
3328 && record.id == peerid_1 {
3329 bootnode_appeared = true;
3330 break;
3331 }
3332 }
3333 _ = &mut timeout => break,
3334 }
3335 }
3336
3337 assert!(bootnode_appeared, "Bootnode should appear in update stream");
3339 }
3340
3341 fn insert_proven_node(service: &mut Discv4Service, record: NodeRecord) {
3342 let key = kad_key(record.id);
3343 let _ = service.kbuckets.insert_or_update(
3344 &key,
3345 NodeEntry::new_proven(record),
3346 NodeStatus {
3347 direction: ConnectionDirection::Incoming,
3348 state: ConnectionState::Connected,
3349 },
3350 );
3351 }
3352
3353 fn insert_initial_ping(service: &mut Discv4Service, record: NodeRecord) -> B256 {
3354 let echo_hash = B256::random();
3355 service.pending_pings.insert(
3356 record.id,
3357 PingRequest {
3358 sent_at: Instant::now(),
3359 node: record,
3360 echo_hash,
3361 reason: PingReason::InitialInsert,
3362 },
3363 );
3364 echo_hash
3365 }
3366
3367 fn make_pong(service: &Discv4Service, echo_hash: B256) -> Pong {
3368 Pong {
3369 to: rng_endpoint(&mut rand_08::thread_rng()),
3370 echo: echo_hash,
3371 expire: service.ping_expiration(),
3372 enr_sq: None,
3373 }
3374 }
3375
3376 #[tokio::test]
3377 async fn test_lookup_reset_on_first_bootnode_pong() {
3378 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3379 let config = Discv4Config::builder().add_boot_node(record).build();
3380 let (_discv4, mut service) = create_discv4_with_config(config).await;
3381
3382 assert!(service.pending_lookup_reset);
3384
3385 insert_proven_node(&mut service, record);
3387 let echo_hash = insert_initial_ping(&mut service, record);
3388
3389 service.on_pong(make_pong(&service, echo_hash), record.udp_addr(), record.id);
3391
3392 assert!(!service.pending_lookup_reset, "flag should be consumed");
3394 }
3395
3396 #[tokio::test]
3397 async fn test_lookup_reset_fires_only_once() {
3398 let records: Vec<_> = (0..2)
3399 .map(|_| NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random()))
3400 .collect();
3401 let config = Discv4Config::builder().add_boot_nodes(records.clone()).build();
3402 let (_discv4, mut service) = create_discv4_with_config(config).await;
3403
3404 for &r in &records {
3406 insert_proven_node(&mut service, r);
3407 }
3408 let hashes: Vec<_> =
3409 records.iter().map(|r| insert_initial_ping(&mut service, *r)).collect();
3410
3411 service.on_pong(make_pong(&service, hashes[0]), records[0].udp_addr(), records[0].id);
3413 assert!(!service.pending_lookup_reset);
3414
3415 service.on_pong(make_pong(&service, hashes[1]), records[1].udp_addr(), records[1].id);
3417 assert!(!service.pending_lookup_reset);
3418 }
3419
3420 #[tokio::test]
3421 async fn test_lookup_reset_not_triggered_by_non_bootnode() {
3422 let bootnode = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3423 let config = Discv4Config::builder().add_boot_node(bootnode).build();
3424 let (_discv4, mut service) = create_discv4_with_config(config).await;
3425
3426 assert!(service.pending_lookup_reset);
3427
3428 let stranger = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3430 insert_proven_node(&mut service, stranger);
3431 let echo_hash = insert_initial_ping(&mut service, stranger);
3432 service.on_pong(make_pong(&service, echo_hash), stranger.udp_addr(), stranger.id);
3433
3434 assert!(service.pending_lookup_reset, "flag should not be consumed by non-bootnode");
3435 }
3436
3437 #[tokio::test]
3438 async fn test_lookup_reset_disabled_when_lookup_disabled() {
3439 let config = Discv4Config::builder().enable_lookup(false).build();
3440 let (_discv4, service) = create_discv4_with_config(config).await;
3441
3442 assert!(!service.pending_lookup_reset);
3444 }
3445
3446 #[tokio::test]
3447 async fn test_enr_response_preserves_pending_request() {
3448 let fork_id = ForkId { hash: ForkHash([1, 2, 3, 4]), next: 0 };
3449 let mut config = Discv4Config::default();
3450 config.add_eip868_pair("eth", EnrForkIdEntry::from(fork_id));
3451 let (_remote, mut remote) = create_discv4_with_config(config).await;
3452 let (_discv4, mut service) = create_discv4().await;
3453 let record = remote.local_node_record;
3454 let id = record.id;
3455 let addr = record.udp_addr();
3456 insert_proven_node(&mut service, record);
3457 service.update_on_reping(record, remote.enr_seq());
3458 let old_enr = remote.local_eip_868_enr.clone();
3459
3460 let (_, old_hash) =
3462 Message::EnrRequest(EnrRequest { expire: service.enr_request_expiration() - 1 })
3463 .encode(&service.secret_key);
3464 service.pending_enr_requests.get_mut(&id).unwrap().echo_hash = old_hash;
3465
3466 remote.local_eip_868_enr.set_tcp4(30304, &remote.secret_key).unwrap();
3468 service.on_ping(
3469 Ping {
3470 from: record.into(),
3471 to: service.local_node_record.into(),
3472 expire: service.ping_expiration(),
3473 enr_sq: remote.enr_seq(),
3474 },
3475 addr,
3476 id,
3477 B256::random(),
3478 );
3479 let request_hash = service.pending_enr_requests[&id].echo_hash;
3480 assert_ne!(request_hash, old_hash);
3481
3482 service.on_enr_response(EnrResponse { request_hash: old_hash, enr: old_enr }, addr, id);
3483 assert_eq!(service.pending_enr_requests[&id].echo_hash, request_hash);
3484 assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(None));
3485
3486 service.on_enr_response(
3488 EnrResponse { request_hash, enr: service.local_eip_868_enr.clone() },
3489 addr,
3490 id,
3491 );
3492 assert_eq!(service.pending_enr_requests[&id].echo_hash, request_hash);
3493 assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(None));
3494
3495 let mut updates = service.update_stream().into_inner();
3496 service.on_enr_response(
3497 EnrResponse { request_hash, enr: remote.local_eip_868_enr.clone() },
3498 addr,
3499 id,
3500 );
3501 assert!(!service.pending_enr_requests.contains_key(&id));
3502 assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(Some(fork_id)));
3503 assert_matches::assert_matches!(
3504 updates.try_recv().unwrap(),
3505 DiscoveryUpdate::EnrForkId(found_record, found_fork_id) => {
3506 assert_eq!(found_record, record);
3507 assert_eq!(found_fork_id, fork_id);
3508 }
3509 );
3510 }
3511}