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, NodeEndpoint, 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(
242 local_address: SocketAddr,
243 local_node_record: NodeRecord,
244 secret_key: SecretKey,
245 config: Discv4Config,
246 ) -> io::Result<(Self, Discv4Service)> {
247 let socket = Arc::new(UdpSocket::bind(local_address).await?);
248 trace!(target: "discv4", local_addr=?socket.local_addr(), "opened UDP socket");
249 let (tx, rx) = mpsc::channel(config.udp_ingress_message_buffer);
250
251 Self::bind_with_socket(socket, None, Some(tx), rx, local_node_record, secret_key, config)
252 }
253
254 pub fn bind_shared(
263 socket: Arc<UdpSocket>,
264 secondary_socket: Option<Arc<UdpSocket>>,
265 local_node_record: NodeRecord,
266 secret_key: SecretKey,
267 config: Discv4Config,
268 ) -> io::Result<(Self, Discv4Service, IngressHandler)> {
269 let (tx, rx) = mpsc::channel(config.udp_ingress_message_buffer);
270 let local_id = local_node_record.id;
271 let (discv4, service) = Self::bind_with_socket(
272 socket,
273 secondary_socket,
274 None,
275 rx,
276 local_node_record,
277 secret_key,
278 config,
279 )?;
280
281 let handler = IngressHandler::new(tx, local_id);
282
283 Ok((discv4, service, handler))
284 }
285
286 fn bind_with_socket(
287 socket: Arc<UdpSocket>,
288 secondary_socket: Option<Arc<UdpSocket>>,
289 ingress_tx: Option<IngressSender>,
290 ingress_rx: IngressReceiver,
291 mut local_node_record: NodeRecord,
292 secret_key: SecretKey,
293 config: Discv4Config,
294 ) -> io::Result<(Self, Discv4Service)> {
295 let local_addr = socket.local_addr()?;
296 local_node_record.udp_port = local_addr.port();
297
298 let mut service = Discv4Service::new(
299 socket,
300 secondary_socket,
301 ingress_tx,
302 ingress_rx,
303 local_addr,
304 local_node_record,
305 secret_key,
306 config,
307 );
308
309 service.resolve_external_ip();
311
312 let discv4 = service.handle();
313 Ok((discv4, service))
314 }
315
316 pub const fn local_addr(&self) -> SocketAddr {
318 self.local_addr
319 }
320
321 pub fn node_record(&self) -> NodeRecord {
325 *self.node_record.lock()
326 }
327
328 pub fn external_ip(&self) -> IpAddr {
330 self.node_record.lock().address
331 }
332
333 pub fn set_lookup_interval(&self, duration: Duration) {
335 self.send_to_service(Discv4Command::SetLookupInterval(duration))
336 }
337
338 pub async fn lookup_self(&self) -> Result<Vec<NodeRecord>, Discv4Error> {
353 self.lookup_node(None).await
354 }
355
356 pub async fn lookup(&self, node_id: PeerId) -> Result<Vec<NodeRecord>, Discv4Error> {
360 self.lookup_node(Some(node_id)).await
361 }
362
363 pub async fn lookup_random(&self) -> Result<Vec<NodeRecord>, Discv4Error> {
365 let target = PeerId::random();
366 self.lookup_node(Some(target)).await
367 }
368
369 pub fn send_lookup(&self, node_id: PeerId) {
371 let cmd = Discv4Command::Lookup { node_id: Some(node_id), tx: None };
372 self.send_to_service(cmd);
373 }
374
375 async fn lookup_node(&self, node_id: Option<PeerId>) -> Result<Vec<NodeRecord>, Discv4Error> {
376 let (tx, rx) = oneshot::channel();
377 let cmd = Discv4Command::Lookup { node_id, tx: Some(tx) };
378 self.to_service.send(cmd)?;
379 Ok(rx.await?)
380 }
381
382 pub fn send_lookup_self(&self) {
384 let cmd = Discv4Command::Lookup { node_id: None, tx: None };
385 self.send_to_service(cmd);
386 }
387
388 pub fn remove_peer(&self, node_id: PeerId) {
390 let cmd = Discv4Command::Remove(node_id);
391 self.send_to_service(cmd);
392 }
393
394 pub fn add_node(&self, node_record: NodeRecord) {
396 let cmd = Discv4Command::Add(node_record);
397 self.send_to_service(cmd);
398 }
399
400 pub fn add_boot_node(&self, node_record: NodeRecord) {
405 let cmd = Discv4Command::AddBootNode(node_record);
406 self.send_to_service(cmd);
407 }
408
409 pub fn ban(&self, node_id: PeerId, ip: IpAddr) {
413 let cmd = Discv4Command::Ban(node_id, ip);
414 self.send_to_service(cmd);
415 }
416
417 pub fn ban_ip(&self, ip: IpAddr) {
421 let cmd = Discv4Command::BanIp(ip);
422 self.send_to_service(cmd);
423 }
424
425 pub fn ban_node(&self, node_id: PeerId) {
429 let cmd = Discv4Command::BanPeer(node_id);
430 self.send_to_service(cmd);
431 }
432
433 pub fn set_tcp_port(&self, port: u16) {
437 let cmd = Discv4Command::SetTcpPort(port);
438 self.send_to_service(cmd);
439 }
440
441 pub fn set_eip868_rlp_pair(&self, key: Vec<u8>, rlp: Bytes) {
447 let cmd = Discv4Command::SetEIP868RLPPair { key, rlp };
448 self.send_to_service(cmd);
449 }
450
451 pub fn set_eip868_rlp(&self, key: Vec<u8>, value: impl alloy_rlp::Encodable) {
455 self.set_eip868_rlp_pair(key, Bytes::from(alloy_rlp::encode(&value)))
456 }
457
458 #[inline]
459 fn send_to_service(&self, cmd: Discv4Command) {
460 let _ = self.to_service.send(cmd).map_err(|err| {
461 debug!(
462 target: "discv4",
463 %err,
464 "channel capacity reached, dropping command",
465 )
466 });
467 }
468
469 pub async fn update_stream(&self) -> Result<ReceiverStream<DiscoveryUpdate>, Discv4Error> {
471 let (tx, rx) = oneshot::channel();
472 let cmd = Discv4Command::Updates(tx);
473 self.to_service.send(cmd)?;
474 Ok(rx.await?)
475 }
476
477 pub fn terminate(&self) {
479 self.send_to_service(Discv4Command::Terminated);
480 }
481}
482
483#[must_use = "Stream does nothing unless polled"]
497pub struct Discv4Service {
498 local_address: SocketAddr,
500 local_eip_868_enr: Enr<SecretKey>,
502 local_node_record: NodeRecord,
504 shared_node_record: Arc<Mutex<NodeRecord>>,
506 secret_key: SecretKey,
508 _socket: Arc<UdpSocket>,
510 _tasks: JoinSet<()>,
514 kbuckets: KBucketsTable<NodeKey, NodeEntry>,
516 ingress: IngressReceiver,
520 egress: EgressSender,
524 queued_pings: VecDeque<(NodeRecord, PingReason)>,
531 pending_pings: HashMap<PeerId, PingRequest>,
533 pending_lookup: HashMap<PeerId, (Instant, LookupContext)>,
538 pending_find_nodes: HashMap<PeerId, FindNodeRequest>,
540 pending_enr_requests: HashMap<PeerId, EnrRequestState>,
542 to_service: mpsc::UnboundedSender<Discv4Command>,
544 commands_rx: mpsc::UnboundedReceiver<Discv4Command>,
546 update_listeners: Vec<mpsc::Sender<DiscoveryUpdate>>,
548 lookup_interval: Interval,
550 lookup_rotator: LookupTargetRotator,
552 pending_lookup_reset: bool,
554 evict_expired_requests_interval: Interval,
556 ping_interval: Interval,
558 resolve_external_ip_interval: Option<ResolveNatInterval>,
560 config: Discv4Config,
562 queued_events: VecDeque<Discv4Event>,
564 received_pongs: PongTable,
566 expire_interval: Interval,
568 cached_find_node: Option<CachedFindNode>,
570}
571
572impl Discv4Service {
573 #[expect(clippy::too_many_arguments)]
578 pub(crate) fn new(
579 socket: Arc<UdpSocket>,
580 secondary_socket: Option<Arc<UdpSocket>>,
581 ingress_tx: Option<IngressSender>,
582 ingress_rx: IngressReceiver,
583 local_address: SocketAddr,
584 local_node_record: NodeRecord,
585 secret_key: SecretKey,
586 mut config: Discv4Config,
587 ) -> Self {
588 config.secondary_advertised_ip = config.secondary_advertised_ip.filter(|ip| {
590 secondary_socket.is_some() &&
591 !ip.is_unspecified() &&
592 ip.is_ipv4() != local_node_record.address.is_ipv4()
593 });
594
595 let (egress_tx, egress_rx) = mpsc::channel(config.udp_egress_message_buffer);
596 let mut tasks = JoinSet::<()>::new();
597
598 if let Some(ingress_tx) = ingress_tx {
599 let udp = Arc::clone(&socket);
600 tasks.spawn(receive_loop(udp, ingress_tx, local_node_record.id));
601 }
602
603 let udp = Arc::clone(&socket);
604 tasks.spawn(send_loop(udp, secondary_socket, local_address.is_ipv4(), egress_rx));
605
606 let kbuckets = KBucketsTable::new(
607 NodeKey::from(&local_node_record).into(),
608 Duration::from_secs(60),
609 MAX_NODES_PER_BUCKET,
610 None,
611 None,
612 );
613
614 let self_lookup_interval = tokio::time::interval(config.lookup_interval);
615
616 let ping_interval = tokio::time::interval_at(
619 tokio::time::Instant::now() + config.ping_interval,
620 config.ping_interval,
621 );
622
623 let evict_expired_requests_interval = tokio::time::interval_at(
624 tokio::time::Instant::now() + config.request_timeout,
625 config.request_timeout,
626 );
627
628 let lookup_rotator = if config.enable_dht_random_walk {
629 LookupTargetRotator::default()
630 } else {
631 LookupTargetRotator::local_only()
632 };
633
634 let local_eip_868_enr = {
636 let mut builder = Enr::builder();
637 for ip in
638 std::iter::once(local_node_record.address).chain(config.secondary_advertised_ip)
639 {
640 if !ip.is_unspecified() {
642 builder.ip(ip);
643 }
644 if ip.is_ipv4() {
645 builder.udp4(local_node_record.udp_port);
646 builder.tcp4(local_node_record.tcp_port);
647 } else {
648 builder.udp6(local_node_record.udp_port);
649 builder.tcp6(local_node_record.tcp_port);
650 }
651 }
652
653 for (key, val) in &config.additional_eip868_rlp_pairs {
654 builder.add_value_rlp(key, val.clone());
655 }
656 builder.build(&secret_key).expect("v4 is set")
657 };
658
659 let (to_service, commands_rx) = mpsc::unbounded_channel();
660
661 let shared_node_record = Arc::new(Mutex::new(local_node_record));
662
663 Self {
664 local_address,
665 local_eip_868_enr,
666 local_node_record,
667 shared_node_record,
668 _socket: socket,
669 kbuckets,
670 secret_key,
671 _tasks: tasks,
672 ingress: ingress_rx,
673 egress: egress_tx,
674 queued_pings: VecDeque::with_capacity(MAX_QUEUED_PINGS),
675 pending_pings: Default::default(),
676 pending_lookup: Default::default(),
677 pending_find_nodes: Default::default(),
678 pending_enr_requests: Default::default(),
679 commands_rx,
680 to_service,
681 update_listeners: Vec::with_capacity(1),
682 lookup_interval: self_lookup_interval,
683 ping_interval,
684 evict_expired_requests_interval,
685 lookup_rotator,
686 pending_lookup_reset: config.enable_lookup,
687 resolve_external_ip_interval: config.resolve_external_ip_interval(),
688 config,
689 queued_events: Default::default(),
690 received_pongs: Default::default(),
691 expire_interval: tokio::time::interval(EXPIRE_DURATION),
692 cached_find_node: None,
693 }
694 }
695
696 pub fn handle(&self) -> Discv4 {
698 Discv4 {
699 local_addr: self.local_address,
700 to_service: self.to_service.clone(),
701 node_record: self.shared_node_record.clone(),
702 }
703 }
704
705 fn enr_seq(&self) -> Option<u64> {
707 self.config.enable_eip868.then(|| self.local_eip_868_enr.seq())
708 }
709
710 pub fn set_lookup_interval(&mut self, duration: Duration) {
712 self.lookup_interval = tokio::time::interval(duration);
713 }
714
715 fn resolve_external_ip(&mut self) {
719 if let Some(r) = &self.resolve_external_ip_interval &&
720 let Some(external_ip) =
721 r.resolver().clone().as_external_ip(self.local_node_record.udp_port)
722 {
723 self.set_external_ip_addr(external_ip);
724 }
725 }
726
727 pub fn set_external_ip_addr(&mut self, external_ip: IpAddr) {
730 if self.config.secondary_advertised_ip.is_some() &&
733 external_ip.is_ipv4() != self.local_node_record.address.is_ipv4()
734 {
735 debug!(target: "discv4", ?external_ip, "Ignoring opposite-family external ip");
736 return
737 }
738 if self.local_node_record.address != external_ip {
739 debug!(target: "discv4", ?external_ip, "Updating external ip");
740 self.local_node_record.address = external_ip;
741 let _ = self.local_eip_868_enr.set_ip(external_ip, &self.secret_key);
742 let mut lock = self.shared_node_record.lock();
743 *lock = self.local_node_record;
744 debug!(target: "discv4", enr=?self.local_eip_868_enr, "Updated local ENR");
745 }
746 }
747
748 pub const fn local_peer_id(&self) -> &PeerId {
750 &self.local_node_record.id
751 }
752
753 pub const fn local_addr(&self) -> SocketAddr {
755 self.local_address
756 }
757
758 pub const fn local_enr(&self) -> NodeRecord {
762 self.local_node_record
763 }
764
765 #[cfg(test)]
767 pub const fn local_enr_mut(&mut self) -> &mut NodeRecord {
768 &mut self.local_node_record
769 }
770
771 pub fn contains_node(&self, id: PeerId) -> bool {
773 let key = kad_key(id);
774 self.kbuckets.get_index(&key).is_some()
775 }
776
777 pub fn bootstrap(&mut self) {
789 for record in self.config.bootstrap_nodes.clone() {
790 debug!(target: "discv4", ?record, "pinging boot node");
791 let key = kad_key(record.id);
792 let entry = NodeEntry::new(record);
793
794 match self.kbuckets.insert_or_update(
796 &key,
797 entry,
798 NodeStatus {
799 state: ConnectionState::Disconnected,
800 direction: ConnectionDirection::Outgoing,
801 },
802 ) {
803 InsertResult::Failed(_) => {}
804 _ => {
805 self.try_ping(record, PingReason::InitialInsert);
806 }
807 }
808 }
809 }
810
811 pub fn add_boot_node(&mut self, record: NodeRecord) -> bool {
816 self.config.bootstrap_nodes.insert(record);
817 self.add_node(record)
818 }
819
820 pub fn spawn(mut self) -> JoinHandle<()> {
824 tokio::task::spawn(async move {
825 self.bootstrap();
826
827 while let Some(event) = self.next().await {
828 trace!(target: "discv4", ?event, "processed");
829 }
830 trace!(target: "discv4", "service terminated");
831 })
832 }
833
834 pub fn update_stream(&mut self) -> ReceiverStream<DiscoveryUpdate> {
836 let (tx, rx) = mpsc::channel(512);
837 self.update_listeners.push(tx);
838 ReceiverStream::new(rx)
839 }
840
841 pub fn lookup_self(&mut self) {
843 self.lookup(self.local_node_record.id)
844 }
845
846 pub fn lookup(&mut self, target: PeerId) {
856 self.lookup_with(target, None)
857 }
858
859 fn lookup_with(&mut self, target: PeerId, tx: Option<NodeRecordSender>) {
869 trace!(target: "discv4", ?target, "Starting lookup");
870 let target_key = kad_key(target);
871
872 let ctx = LookupContext::new(
875 target_key.clone(),
876 self.kbuckets
877 .closest_values(&target_key)
878 .filter(|node| {
879 node.value.has_endpoint_proof &&
880 !self.pending_find_nodes.contains_key(&node.key.preimage().0)
881 })
882 .take(MAX_NODES_PER_BUCKET)
883 .map(|n| (target_key.distance(&n.key), n.value.record)),
884 tx,
885 );
886
887 let closest = ctx.closest(ALPHA);
889
890 if closest.is_empty() && self.pending_find_nodes.is_empty() {
891 self.bootstrap();
896 return
897 }
898
899 trace!(target: "discv4", ?target, num = closest.len(), "Start lookup closest nodes");
900
901 for node in closest {
902 self.find_node_checked(&node, ctx.clone());
906 }
907 }
908
909 fn find_node(&mut self, node: &NodeRecord, ctx: LookupContext) {
913 trace!(target: "discv4", ?node, lookup=?ctx.target(), "Sending FindNode");
914 ctx.mark_queried(node.id);
915 let (payload, hash) = self.find_node_packet(ctx.target());
916 let to = node.udp_addr();
917 trace!(target: "discv4", ?to, ?hash, "sending FindNode packet");
918 let _ = self.egress.try_send((payload, to)).map_err(|err| {
919 debug!(target: "discv4", %err, "dropped outgoing packet");
920 });
921 self.pending_find_nodes.insert(node.id, FindNodeRequest::new(ctx, to.ip()));
922 }
923
924 fn find_node_checked(&mut self, node: &NodeRecord, ctx: LookupContext) {
929 let max_failures = self.config.max_find_node_failures;
930 let needs_ping = self
931 .on_entry(node.id, |entry| entry.exceeds_find_node_failures(max_failures))
932 .unwrap_or(true);
933 if needs_ping {
934 self.try_ping(*node, PingReason::Lookup(*node, ctx))
935 } else {
936 self.find_node(node, ctx)
937 }
938 }
939
940 fn notify(&mut self, update: DiscoveryUpdate) {
944 self.update_listeners.retain_mut(|listener| match listener.try_send(update.clone()) {
945 Ok(()) => true,
946 Err(err) => match err {
947 TrySendError::Full(_) => true,
948 TrySendError::Closed(_) => false,
949 },
950 });
951 }
952
953 pub fn ban_ip(&mut self, ip: IpAddr) {
955 self.config.ban_list.ban_ip(ip);
956 }
957
958 pub fn ban_node(&mut self, node_id: PeerId) {
960 self.remove_node(node_id);
961 self.config.ban_list.ban_peer(node_id);
962 }
963
964 pub fn ban_ip_until(&mut self, ip: IpAddr, until: Instant) {
966 self.config.ban_list.ban_ip_until(ip, until);
967 }
968
969 pub fn ban_node_until(&mut self, node_id: PeerId, until: Instant) {
971 self.remove_node(node_id);
972 self.config.ban_list.ban_peer_until(node_id, until);
973 }
974
975 pub fn remove_node(&mut self, node_id: PeerId) -> bool {
980 let key = kad_key(node_id);
981 self.remove_key(node_id, key)
982 }
983
984 pub fn soft_remove_node(&mut self, node_id: PeerId) -> bool {
989 let key = kad_key(node_id);
990 let Some(bucket) = self.kbuckets.get_bucket(&key) else { return false };
991 if bucket.num_entries() < MAX_NODES_PER_BUCKET / 2 {
992 return false
994 }
995 self.remove_key(node_id, key)
996 }
997
998 fn remove_key(&mut self, node_id: PeerId, key: discv5::Key<NodeKey>) -> bool {
999 let removed = self.kbuckets.remove(&key);
1000 if removed {
1001 trace!(target: "discv4", ?node_id, "removed node");
1002 self.notify(DiscoveryUpdate::Removed(node_id));
1003 }
1004 removed
1005 }
1006
1007 pub fn num_connected(&self) -> usize {
1009 self.kbuckets.buckets_iter().fold(0, |count, bucket| count + bucket.num_connected())
1010 }
1011
1012 fn has_bond(&self, remote_id: PeerId, remote_ip: IpAddr) -> bool {
1014 if let Some(timestamp) = self.received_pongs.last_pong(remote_id, remote_ip) &&
1015 timestamp.elapsed() < self.config.bond_expiration
1016 {
1017 return true
1018 }
1019 false
1020 }
1021
1022 fn on_entry<F, R>(&mut self, peer_id: PeerId, f: F) -> Option<R>
1024 where
1025 F: FnOnce(&NodeEntry) -> R,
1026 {
1027 let key = kad_key(peer_id);
1028 match self.kbuckets.entry(&key) {
1029 BucketEntry::Present(entry, _) => Some(f(entry.value())),
1030 BucketEntry::Pending(entry, _) => Some(f(entry.value())),
1031 _ => None,
1032 }
1033 }
1034
1035 fn update_on_reping(&mut self, record: NodeRecord, mut last_enr_seq: Option<u64>) {
1042 if record.id == self.local_node_record.id {
1043 return
1044 }
1045
1046 if !self.config.enable_eip868 {
1048 last_enr_seq = None;
1049 }
1050
1051 let key = kad_key(record.id);
1052 let old_enr = match self.kbuckets.entry(&key) {
1053 kbucket::Entry::Present(mut entry, _) => {
1054 entry.value_mut().update_with_enr(last_enr_seq)
1055 }
1056 kbucket::Entry::Pending(mut entry, _) => {
1057 entry.value_mut().update_with_enr(last_enr_seq)
1058 }
1059 _ => return,
1060 };
1061
1062 match (last_enr_seq, old_enr) {
1064 (Some(new), Some(old)) if new > old => {
1065 self.send_enr_request(record);
1066 }
1067 (Some(_), None) => {
1068 self.send_enr_request(record);
1070 }
1071 _ => {}
1072 };
1073 }
1074
1075 fn update_on_pong(&mut self, record: NodeRecord, mut last_enr_seq: Option<u64>) {
1077 if record.id == *self.local_peer_id() {
1078 return
1079 }
1080
1081 if !self.config.enable_eip868 {
1083 last_enr_seq = None;
1084 }
1085
1086 let has_enr_seq = last_enr_seq.is_some();
1089
1090 let key = kad_key(record.id);
1091 match self.kbuckets.entry(&key) {
1092 kbucket::Entry::Present(mut entry, old_status) => {
1093 entry.value_mut().establish_proof();
1095 entry.value_mut().update_with_enr(last_enr_seq);
1096
1097 if !old_status.is_connected() {
1098 let _ = entry.update(ConnectionState::Connected, Some(old_status.direction));
1099 trace!(target: "discv4", ?record, "added after successful endpoint proof");
1100 self.notify(DiscoveryUpdate::Added(record));
1101
1102 if has_enr_seq {
1103 self.send_enr_request(record);
1105 }
1106 }
1107 }
1108 kbucket::Entry::Pending(mut entry, mut status) => {
1109 entry.value_mut().establish_proof();
1111 entry.value_mut().update_with_enr(last_enr_seq);
1112
1113 if !status.is_connected() {
1114 status.state = ConnectionState::Connected;
1115 let _ = entry.update(status);
1116 trace!(target: "discv4", ?record, "added after successful endpoint proof");
1117 self.notify(DiscoveryUpdate::Added(record));
1118
1119 if has_enr_seq {
1120 self.send_enr_request(record);
1122 }
1123 }
1124 }
1125 _ => {}
1126 };
1127 }
1128
1129 pub fn add_all_nodes(&mut self, records: impl IntoIterator<Item = NodeRecord>) {
1133 for record in records {
1134 self.add_node(record);
1135 }
1136 }
1137
1138 pub fn add_node(&mut self, record: NodeRecord) -> bool {
1144 let key = kad_key(record.id);
1145 match self.kbuckets.entry(&key) {
1146 kbucket::Entry::Absent(entry) => {
1147 let node = NodeEntry::new(record);
1148 match entry.insert(
1149 node,
1150 NodeStatus {
1151 direction: ConnectionDirection::Outgoing,
1152 state: ConnectionState::Disconnected,
1153 },
1154 ) {
1155 BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1156 trace!(target: "discv4", ?record, "inserted new record");
1157 }
1158 _ => return false,
1159 }
1160 }
1161 _ => return false,
1162 }
1163
1164 self.try_ping(record, PingReason::InitialInsert);
1166 true
1167 }
1168
1169 pub(crate) fn send_packet(&self, msg: Message, to: SocketAddr) -> B256 {
1171 let (payload, hash) = msg.encode(&self.secret_key);
1172 trace!(target: "discv4", r#type=?msg.msg_type(), ?to, ?hash, "sending packet");
1173 let _ = self.egress.try_send((payload, to)).map_err(|err| {
1174 debug!(
1175 target: "discv4",
1176 %err,
1177 "dropped outgoing packet",
1178 );
1179 });
1180 hash
1181 }
1182
1183 fn find_node_packet(&mut self, target: PeerId) -> (Bytes, B256) {
1185 let expire = self.find_node_expiration();
1186 let cache_ttl = self.config.request_timeout / 4;
1187 CachedFindNode::get_or_sign(
1188 &mut self.cached_find_node,
1189 target,
1190 cache_ttl,
1191 &self.secret_key,
1192 expire,
1193 )
1194 }
1195
1196 fn on_ping(&mut self, ping: Ping, remote_addr: SocketAddr, remote_id: PeerId, hash: B256) {
1198 if self.is_expired(ping.expire) {
1199 return
1201 }
1202
1203 let record = NodeRecord {
1205 address: remote_addr.ip(),
1206 udp_port: remote_addr.port(),
1207 tcp_port: ping.from.tcp_port,
1208 id: remote_id,
1209 }
1210 .into_ipv4_mapped();
1211
1212 let key = kad_key(record.id);
1213
1214 let mut is_new_insert = false;
1219 let mut needs_bond = !self.has_bond(record.id, record.address);
1220 let mut is_proven = false;
1221
1222 let old_enr = match self.kbuckets.entry(&key) {
1223 kbucket::Entry::Present(mut entry, _) => {
1224 if entry.value().is_expired() {
1225 needs_bond = true;
1228 } else {
1229 is_proven = entry.value().has_endpoint_proof;
1230 }
1231 entry.value_mut().update_with_enr(ping.enr_sq)
1232 }
1233 kbucket::Entry::Pending(mut entry, _) => {
1234 if entry.value().is_expired() {
1235 needs_bond = true;
1238 } else {
1239 is_proven = entry.value().has_endpoint_proof;
1240 }
1241 entry.value_mut().update_with_enr(ping.enr_sq)
1242 }
1243 kbucket::Entry::Absent(entry) => {
1244 let mut node = NodeEntry::new(record);
1245 node.last_enr_seq = ping.enr_sq;
1246
1247 match entry.insert(
1248 node,
1249 NodeStatus {
1250 direction: ConnectionDirection::Incoming,
1251 state: ConnectionState::Disconnected,
1253 },
1254 ) {
1255 BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1256 is_new_insert = true;
1258 }
1259 BucketInsertResult::Full => {
1260 trace!(target: "discv4", ?record, "discovered new record but bucket is full");
1264 self.notify(DiscoveryUpdate::DiscoveredAtCapacity(record));
1265 needs_bond = true;
1266 }
1267 BucketInsertResult::TooManyIncoming | BucketInsertResult::NodeExists => {
1268 needs_bond = true;
1269 }
1271 BucketInsertResult::FailedFilter => return,
1272 }
1273
1274 None
1275 }
1276 kbucket::Entry::SelfEntry => return,
1277 };
1278
1279 let pong = Message::Pong(Pong {
1282 to: record.into(),
1284 echo: hash,
1285 expire: ping.expire,
1286 enr_sq: self.enr_seq(),
1287 });
1288 self.send_packet(pong, remote_addr);
1289
1290 if is_new_insert {
1292 self.try_ping(record, PingReason::InitialInsert);
1293 } else if needs_bond {
1294 self.try_ping(record, PingReason::EstablishBond);
1295 } else {
1296 if is_proven && let Some((_, ctx)) = self.pending_lookup.remove(&record.id) {
1300 if self.pending_find_nodes.contains_key(&record.id) {
1301 ctx.unmark_queried(record.id);
1304 } else {
1305 self.find_node(&record, ctx);
1308 }
1309 }
1310
1311 match (ping.enr_sq, old_enr) {
1313 (Some(new), Some(old)) if new > old => {
1314 self.send_enr_request(record);
1315 }
1316 (Some(_), None) => {
1317 self.send_enr_request(record);
1318 }
1319 _ => {}
1320 };
1321 }
1322 }
1323
1324 fn try_ping(&mut self, node: NodeRecord, reason: PingReason) {
1326 if node.id == *self.local_peer_id() {
1327 return
1329 }
1330
1331 if self.pending_pings.get(&node.id).is_some_and(|ping| ping.node.address == node.address) {
1332 return
1333 }
1334
1335 if self.queued_pings.iter().any(|(n, _)| n.id == node.id && n.address == node.address) {
1336 return
1337 }
1338
1339 if self.pending_pings.contains_key(&node.id) ||
1340 self.pending_find_nodes.contains_key(&node.id)
1341 {
1342 if matches!(reason, PingReason::EstablishBond) &&
1343 self.queued_pings.len() < MAX_QUEUED_PINGS
1344 {
1345 self.queued_pings.push_back((node, reason));
1346 }
1347 return
1348 }
1349
1350 if self.pending_pings.len() < MAX_NODES_PING {
1351 self.send_ping(node, reason);
1352 } else if self.queued_pings.len() < MAX_QUEUED_PINGS {
1353 self.queued_pings.push_back((node, reason));
1354 }
1355 }
1356
1357 pub(crate) fn send_ping(&mut self, node: NodeRecord, reason: PingReason) -> B256 {
1361 let remote_addr = node.udp_addr();
1362 let id = node.id;
1363 let mut from: NodeEndpoint = self.local_node_record.into();
1364 from.tcp_port = if node.address.is_ipv4() {
1365 self.local_eip_868_enr.tcp4()
1366 } else {
1367 self.local_eip_868_enr.tcp6()
1368 }
1369 .unwrap_or_default();
1370 let ping =
1371 Ping { from, to: node.into(), expire: self.ping_expiration(), enr_sq: self.enr_seq() };
1372 trace!(target: "discv4", ?ping, "sending ping");
1373 let echo_hash = self.send_packet(Message::Ping(ping), remote_addr);
1374
1375 self.pending_pings
1376 .insert(id, PingRequest { sent_at: Instant::now(), node, echo_hash, reason });
1377 echo_hash
1378 }
1379
1380 pub(crate) fn send_enr_request(&mut self, node: NodeRecord) {
1384 if !self.config.enable_eip868 {
1385 return
1386 }
1387 let remote_addr = node.udp_addr();
1388 let enr_request = EnrRequest { expire: self.enr_request_expiration() };
1389
1390 trace!(target: "discv4", ?enr_request, "sending enr request");
1391 let echo_hash = self.send_packet(Message::EnrRequest(enr_request), remote_addr);
1392
1393 self.pending_enr_requests
1394 .insert(node.id, EnrRequestState { sent_at: Instant::now(), echo_hash });
1395 }
1396
1397 fn on_pong(&mut self, pong: Pong, remote_addr: SocketAddr, remote_id: PeerId) {
1399 if self.is_expired(pong.expire) {
1400 return
1401 }
1402
1403 let PingRequest { node, reason, .. } = match self.pending_pings.entry(remote_id) {
1404 Entry::Occupied(entry) => {
1405 {
1406 let request = entry.get();
1407 if request.echo_hash != pong.echo {
1408 trace!(target: "discv4", from=?remote_addr, expected=?request.echo_hash, echo_hash=?pong.echo,"Got unexpected Pong");
1409 return
1410 }
1411 if request.node.udp_addr().ip().to_canonical() !=
1413 remote_addr.ip().to_canonical()
1414 {
1415 trace!(target: "discv4", from=?remote_addr, expected=?request.node.udp_addr(), "Ignoring Pong from unexpected IP");
1416 return
1417 }
1418 }
1419 entry.remove()
1420 }
1421 Entry::Vacant(_) => return,
1422 };
1423
1424 self.received_pongs.on_pong(remote_id, remote_addr.ip());
1426
1427 match reason {
1428 PingReason::InitialInsert => {
1429 self.update_on_pong(node, pong.enr_sq);
1430 if self.pending_lookup_reset && self.config.bootstrap_nodes.contains(&node) {
1433 self.pending_lookup_reset = false;
1434 self.lookup_interval.reset();
1435 }
1436 }
1437 PingReason::EstablishBond => {
1438 self.update_on_pong(node, pong.enr_sq);
1440 }
1441 PingReason::RePing => {
1442 self.update_on_reping(node, pong.enr_sq);
1443 }
1444 PingReason::Lookup(node, ctx) => {
1445 self.update_on_pong(node, pong.enr_sq);
1446 self.pending_lookup.insert(node.id, (Instant::now(), ctx));
1451 }
1452 }
1453 }
1454
1455 fn on_find_node(&mut self, msg: FindNode, remote_addr: SocketAddr, node_id: PeerId) {
1457 if self.is_expired(msg.expire) {
1458 return
1460 }
1461 if node_id == *self.local_peer_id() {
1462 return
1464 }
1465
1466 if self.has_bond(node_id, remote_addr.ip()) {
1467 self.respond_closest(msg.id, remote_addr)
1468 }
1469 }
1470
1471 fn on_enr_response(&mut self, msg: EnrResponse, remote_addr: SocketAddr, id: PeerId) {
1473 trace!(target: "discv4", ?remote_addr, ?msg, "received ENR response");
1474 let Entry::Occupied(request) = self.pending_enr_requests.entry(id) else { return };
1475
1476 if request.get().echo_hash != msg.request_hash || id != pk2id(&msg.enr.public_key()) {
1479 return
1480 }
1481 request.remove();
1482
1483 let key = kad_key(id);
1484 let fork_id = msg.eth_fork_id();
1485 let (record, old_fork_id) = match self.kbuckets.entry(&key) {
1486 kbucket::Entry::Present(mut entry, _) => {
1487 let id = entry.value_mut().update_with_fork_id(fork_id);
1488 (entry.value().record, id)
1489 }
1490 kbucket::Entry::Pending(mut entry, _) => {
1491 let id = entry.value_mut().update_with_fork_id(fork_id);
1492 (entry.value().record, id)
1493 }
1494 _ => return,
1495 };
1496 match (fork_id, old_fork_id) {
1497 (Some(new), Some(old)) if new != old => {
1498 self.notify(DiscoveryUpdate::EnrForkId(record, new))
1499 }
1500 (Some(new), None) => self.notify(DiscoveryUpdate::EnrForkId(record, new)),
1501 _ => {}
1502 }
1503 }
1504
1505 fn on_enr_request(
1507 &self,
1508 msg: EnrRequest,
1509 remote_addr: SocketAddr,
1510 id: PeerId,
1511 request_hash: B256,
1512 ) {
1513 if !self.config.enable_eip868 || self.is_expired(msg.expire) {
1514 return
1515 }
1516
1517 if self.has_bond(id, remote_addr.ip()) {
1518 self.send_packet(
1519 Message::EnrResponse(EnrResponse {
1520 request_hash,
1521 enr: self.local_eip_868_enr.clone(),
1522 }),
1523 remote_addr,
1524 );
1525 }
1526 }
1527
1528 fn on_neighbours(&mut self, msg: Neighbours, remote_addr: SocketAddr, node_id: PeerId) {
1531 if self.is_expired(msg.expire) {
1532 return
1534 }
1535 let ctx = match self.pending_find_nodes.entry(node_id) {
1537 Entry::Occupied(mut entry) => {
1538 {
1539 let request = entry.get_mut();
1540 if !request.is_from_destination(remote_addr.ip()) {
1544 trace!(target: "discv4", from=?remote_addr, expected=?request.destination, "Ignoring Neighbours from unexpected IP");
1545 return
1546 }
1547 request.answered = true;
1549 let total = request.response_count + msg.nodes.len();
1550
1551 if total <= MAX_NODES_PER_BUCKET {
1553 request.response_count = total;
1554 } else {
1555 trace!(target: "discv4", total, from=?remote_addr, "Received neighbors packet entries exceeds max nodes per bucket");
1556 return
1557 }
1558 };
1559
1560 if entry.get().response_count == MAX_NODES_PER_BUCKET {
1561 let ctx = entry.remove().lookup_context;
1563 ctx.mark_responded(node_id);
1564 ctx
1565 } else {
1566 entry.get().lookup_context.clone()
1567 }
1568 }
1569 Entry::Vacant(_) => {
1570 trace!(target: "discv4", from=?remote_addr, "Received unsolicited Neighbours");
1572 return
1573 }
1574 };
1575
1576 trace!(target: "discv4",
1578 target=format!("{:#?}", node_id),
1579 peers_count=msg.nodes.len(),
1580 peers=format!("[{:#}]", msg.nodes.iter()
1581 .map(|node_rec| node_rec.id
1582 ).format(", ")),
1583 "Received peers from Neighbours packet"
1584 );
1585
1586 for node in msg.nodes.into_iter().map(NodeRecord::into_ipv4_mapped) {
1589 if self.config.ban_list.is_banned(&node.id, &node.address) {
1591 trace!(target: "discv4", peer_id=?node.id, ip=?node.address, "ignoring banned record");
1592 continue
1593 }
1594
1595 ctx.add_node(node);
1596 }
1597
1598 let closest =
1600 ctx.filter_closest(ALPHA, |node| !self.pending_find_nodes.contains_key(&node.id));
1601
1602 for closest in closest {
1603 let key = kad_key(closest.id);
1604 match self.kbuckets.entry(&key) {
1605 BucketEntry::Absent(entry) => {
1606 ctx.mark_queried(closest.id);
1612 let node = NodeEntry::new(closest);
1613 match entry.insert(
1614 node,
1615 NodeStatus {
1616 direction: ConnectionDirection::Outgoing,
1617 state: ConnectionState::Disconnected,
1618 },
1619 ) {
1620 BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1621 self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1623 }
1624 BucketInsertResult::Full => {
1625 self.notify(DiscoveryUpdate::DiscoveredAtCapacity(closest))
1627 }
1628 _ => {}
1629 }
1630 }
1631 BucketEntry::SelfEntry => {
1632 }
1634 BucketEntry::Present(entry, _) => {
1635 if entry.value().has_endpoint_proof {
1636 if entry
1637 .value()
1638 .exceeds_find_node_failures(self.config.max_find_node_failures)
1639 {
1640 self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1641 } else {
1642 self.find_node(&closest, ctx.clone());
1643 }
1644 }
1645 }
1646 BucketEntry::Pending(entry, _) => {
1647 if entry.value().has_endpoint_proof {
1648 if entry
1649 .value()
1650 .exceeds_find_node_failures(self.config.max_find_node_failures)
1651 {
1652 self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1653 } else {
1654 self.find_node(&closest, ctx.clone());
1655 }
1656 }
1657 }
1658 }
1659 }
1660 }
1661
1662 fn respond_closest(&mut self, target: PeerId, to: SocketAddr) {
1664 let key = kad_key(target);
1665 let expire = self.send_neighbours_expiration();
1666
1667 let enforce_eip868 = self.config.enable_eip868 && self.config.enforce_eip868_neighbours;
1668
1669 let closest_nodes = self
1672 .kbuckets
1673 .closest_values(&key)
1674 .filter(|entry| entry.value.has_endpoint_proof)
1675 .filter(|entry| !enforce_eip868 || entry.value.fork_id.is_some())
1676 .take(MAX_NODES_PER_BUCKET)
1677 .collect::<Vec<_>>();
1678
1679 if closest_nodes.is_empty() {
1680 let msg = Message::Neighbours(Neighbours { nodes: Vec::new(), expire });
1682 self.send_packet(msg, to);
1683 return;
1684 }
1685
1686 for nodes in closest_nodes.chunks(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS) {
1687 let nodes = nodes.iter().map(|node| node.value.record).collect::<Vec<NodeRecord>>();
1688 trace!(target: "discv4", len = nodes.len(), to=?to,"Sent neighbours packet");
1689 let msg = Message::Neighbours(Neighbours { nodes, expire });
1690 self.send_packet(msg, to);
1691 }
1692 }
1693
1694 fn evict_expired_requests(&mut self, now: Instant) {
1695 let mut failed_enr_requests = Vec::new();
1696 self.pending_enr_requests.retain(|node_id, enr_request| {
1697 if now.duration_since(enr_request.sent_at) < self.config.enr_expiration {
1698 return true
1699 }
1700 failed_enr_requests.push(*node_id);
1701 false
1702 });
1703
1704 for node_id in failed_enr_requests {
1706 match self.kbuckets.entry(&kad_key(node_id)) {
1707 kbucket::Entry::Present(mut entry, _) => entry.value_mut().last_enr_seq = None,
1708 kbucket::Entry::Pending(mut entry, _) => entry.value_mut().last_enr_seq = None,
1709 _ => {}
1710 }
1711 }
1712
1713 let mut failed_pings = Vec::new();
1714 self.pending_pings.retain(|node_id, ping_request| {
1715 if now.duration_since(ping_request.sent_at) > self.config.ping_expiration {
1716 failed_pings.push((*node_id, ping_request.node.udp_addr()));
1717 return false
1718 }
1719 true
1720 });
1721
1722 if !failed_pings.is_empty() {
1723 trace!(target: "discv4", num=%failed_pings.len(), "evicting nodes due to failed pong");
1725 for (node_id, pinged) in failed_pings {
1726 let proven_elsewhere = self
1729 .on_entry(node_id, |entry| entry.has_endpoint_proof.then_some(entry.record))
1730 .flatten()
1731 .filter(|record| record.udp_addr() != pinged);
1732 if let Some(record) = proven_elsewhere {
1733 self.try_ping(record, PingReason::RePing);
1734 } else {
1735 self.remove_node(node_id);
1736 }
1737 }
1738 }
1739
1740 let mut failed_lookups = Vec::new();
1741 self.pending_lookup.retain(|node_id, (lookup_sent_at, _)| {
1742 if now.duration_since(*lookup_sent_at) > self.config.request_timeout {
1743 failed_lookups.push(*node_id);
1744 return false
1745 }
1746 true
1747 });
1748
1749 if !failed_lookups.is_empty() {
1750 trace!(target: "discv4", num=%failed_lookups.len(), "evicting nodes due to failed lookup");
1752 for node_id in failed_lookups {
1753 self.remove_node(node_id);
1754 }
1755 }
1756
1757 self.evict_failed_find_nodes(now);
1758 }
1759
1760 fn evict_failed_find_nodes(&mut self, now: Instant) {
1762 let mut failed_find_nodes = Vec::new();
1763 self.pending_find_nodes.retain(|node_id, find_node_request| {
1764 if now.duration_since(find_node_request.sent_at) > self.config.neighbours_expiration {
1765 if !find_node_request.answered {
1766 failed_find_nodes.push((*node_id, find_node_request.destination));
1769 }
1770 return false
1771 }
1772 true
1773 });
1774
1775 if failed_find_nodes.is_empty() {
1776 return
1777 }
1778
1779 trace!(target: "discv4", num=%failed_find_nodes.len(), "processing failed find nodes");
1780
1781 for (node_id, destination) in failed_find_nodes {
1782 let key = kad_key(node_id);
1783 let failures = match self.kbuckets.entry(&key) {
1784 kbucket::Entry::Present(mut entry, _)
1785 if entry.value().record.address.to_canonical() == destination =>
1786 {
1787 entry.value_mut().inc_failed_request();
1788 entry.value().find_node_failures
1789 }
1790 kbucket::Entry::Pending(mut entry, _)
1791 if entry.value().record.address.to_canonical() == destination =>
1792 {
1793 entry.value_mut().inc_failed_request();
1794 entry.value().find_node_failures
1795 }
1796 _ => continue,
1798 };
1799
1800 if failures > self.config.max_find_node_failures {
1804 self.soft_remove_node(node_id);
1805 }
1806 }
1807 }
1808
1809 fn re_ping_oldest(&mut self) {
1814 let mut nodes = self
1815 .kbuckets
1816 .iter_ref()
1817 .filter(|entry| entry.node.value.is_expired())
1818 .map(|n| n.node.value)
1819 .collect::<Vec<_>>();
1820 nodes.sort_unstable_by_key(|a| a.last_seen);
1821 let to_ping = nodes.into_iter().map(|n| n.record).take(MAX_NODES_PING).collect::<Vec<_>>();
1822 for node in to_ping {
1823 self.try_ping(node, PingReason::RePing)
1824 }
1825 }
1826
1827 fn is_expired(&self, expiration: u64) -> bool {
1829 self.ensure_not_expired(expiration).is_err()
1830 }
1831
1832 fn ensure_not_expired(&self, timestamp: u64) -> Result<(), ()> {
1842 let _ = i64::try_from(timestamp).map_err(drop)?;
1844
1845 let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
1846 if self.config.enforce_expiration_timestamps && timestamp < now {
1847 trace!(target: "discv4", "Expired packet");
1848 return Err(())
1849 }
1850 Ok(())
1851 }
1852
1853 fn ping_buffered(&mut self) {
1855 for _ in 0..self.queued_pings.len() {
1856 if self.pending_pings.len() >= MAX_NODES_PING {
1857 break
1858 }
1859 match self.queued_pings.pop_front() {
1860 Some((next, reason)) => self.try_ping(next, reason),
1861 None => break,
1862 }
1863 }
1864 }
1865
1866 fn ping_expiration(&self) -> u64 {
1867 (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.ping_expiration)
1868 .as_secs()
1869 }
1870
1871 fn find_node_expiration(&self) -> u64 {
1872 (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.request_timeout)
1873 .as_secs()
1874 }
1875
1876 fn enr_request_expiration(&self) -> u64 {
1877 (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.enr_expiration)
1878 .as_secs()
1879 }
1880
1881 fn send_neighbours_expiration(&self) -> u64 {
1882 (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.neighbours_expiration)
1883 .as_secs()
1884 }
1885
1886 pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Discv4Event> {
1892 loop {
1893 if let Some(event) = self.queued_events.pop_front() {
1895 return Poll::Ready(event)
1896 }
1897
1898 if self.config.enable_lookup {
1900 while self.lookup_interval.poll_tick(cx).is_ready() {
1901 let target = self.lookup_rotator.next(&self.local_node_record.id);
1902 self.lookup_with(target, None);
1903 }
1904 }
1905
1906 while self.ping_interval.poll_tick(cx).is_ready() {
1908 self.re_ping_oldest();
1909 }
1910
1911 if let Some(Poll::Ready(Some(ip))) =
1912 self.resolve_external_ip_interval.as_mut().map(|r| r.poll_tick(cx))
1913 {
1914 self.set_external_ip_addr(ip);
1915 }
1916
1917 while let Poll::Ready(Some(cmd)) = self.commands_rx.poll_recv(cx) {
1919 match cmd {
1920 Discv4Command::Add(enr) => {
1921 self.add_node(enr);
1922 }
1923 Discv4Command::AddBootNode(record) => {
1924 self.add_boot_node(record);
1925 }
1926 Discv4Command::Lookup { node_id, tx } => {
1927 let node_id = node_id.unwrap_or(self.local_node_record.id);
1928 self.lookup_with(node_id, tx);
1929 }
1930 Discv4Command::SetLookupInterval(duration) => {
1931 self.set_lookup_interval(duration);
1932 }
1933 Discv4Command::Updates(tx) => {
1934 let rx = self.update_stream();
1935 let _ = tx.send(rx);
1936 }
1937 Discv4Command::BanPeer(node_id) => self.ban_node(node_id),
1938 Discv4Command::Remove(node_id) => {
1939 self.remove_node(node_id);
1940 }
1941 Discv4Command::Ban(node_id, ip) => {
1942 self.ban_node(node_id);
1943 self.ban_ip(ip);
1944 }
1945 Discv4Command::BanIp(ip) => {
1946 self.ban_ip(ip);
1947 }
1948 Discv4Command::SetEIP868RLPPair { key, rlp } => {
1949 debug!(target: "discv4", key=%String::from_utf8_lossy(&key), "Update EIP-868 extension pair");
1950
1951 let _ = self.local_eip_868_enr.insert_raw_rlp(key, rlp, &self.secret_key);
1952 }
1953 Discv4Command::SetTcpPort(port) => {
1954 debug!(target: "discv4", %port, "Update tcp port");
1955 self.local_node_record.tcp_port = port;
1956 if self.local_node_record.address.is_ipv4() {
1957 let _ = self.local_eip_868_enr.set_tcp4(port, &self.secret_key);
1958 } else {
1959 let _ = self.local_eip_868_enr.set_tcp6(port, &self.secret_key);
1960 }
1961 if let Some(ip) = self.config.secondary_advertised_ip {
1962 if ip.is_ipv4() {
1963 let _ = self.local_eip_868_enr.set_tcp4(port, &self.secret_key);
1964 } else {
1965 let _ = self.local_eip_868_enr.set_tcp6(port, &self.secret_key);
1966 }
1967 }
1968 }
1969
1970 Discv4Command::Terminated => {
1971 self.queued_events.push_back(Discv4Event::Terminated);
1973 }
1974 }
1975 }
1976
1977 let mut udp_message_budget = UDP_MESSAGE_POLL_LOOP_BUDGET;
1979
1980 while let Poll::Ready(Some(event)) = self.ingress.poll_recv(cx) {
1982 match event {
1983 IngressEvent::RecvError(err) => {
1984 debug!(target: "discv4", %err, "failed to read datagram");
1985 }
1986 IngressEvent::BadPacket(from, err, data) => {
1987 trace!(target: "discv4", ?from, %err, packet=?hex::encode(&data), "bad packet");
1988 }
1989 IngressEvent::Packet(remote_addr, Packet { msg, node_id, hash }) => {
1990 trace!(target: "discv4", r#type=?msg.msg_type(), from=?remote_addr,"received packet");
1991 let event = match msg {
1992 Message::Ping(ping) => {
1993 self.on_ping(ping, remote_addr, node_id, hash);
1994 Discv4Event::Ping
1995 }
1996 Message::Pong(pong) => {
1997 self.on_pong(pong, remote_addr, node_id);
1998 Discv4Event::Pong
1999 }
2000 Message::FindNode(msg) => {
2001 self.on_find_node(msg, remote_addr, node_id);
2002 Discv4Event::FindNode
2003 }
2004 Message::Neighbours(msg) => {
2005 self.on_neighbours(msg, remote_addr, node_id);
2006 Discv4Event::Neighbours
2007 }
2008 Message::EnrRequest(msg) => {
2009 self.on_enr_request(msg, remote_addr, node_id, hash);
2010 Discv4Event::EnrRequest
2011 }
2012 Message::EnrResponse(msg) => {
2013 self.on_enr_response(msg, remote_addr, node_id);
2014 Discv4Event::EnrResponse
2015 }
2016 };
2017
2018 self.queued_events.push_back(event);
2019 }
2020 }
2021
2022 udp_message_budget -= 1;
2023 if udp_message_budget < 0 {
2024 trace!(target: "discv4", budget=UDP_MESSAGE_POLL_LOOP_BUDGET, "exhausted message poll budget");
2025 if self.queued_events.is_empty() {
2026 cx.waker().wake_by_ref();
2029 }
2030 break
2031 }
2032 }
2033
2034 self.ping_buffered();
2036
2037 while self.evict_expired_requests_interval.poll_tick(cx).is_ready() {
2039 self.evict_expired_requests(Instant::now());
2040 }
2041
2042 while self.expire_interval.poll_tick(cx).is_ready() {
2044 self.received_pongs.evict_expired(Instant::now(), EXPIRE_DURATION);
2045 }
2046
2047 if self.queued_events.is_empty() {
2048 return Poll::Pending
2049 }
2050 }
2051 }
2052}
2053
2054impl Stream for Discv4Service {
2056 type Item = Discv4Event;
2057
2058 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2059 match ready!(self.get_mut().poll(cx)) {
2061 Discv4Event::Terminated => Poll::Ready(None),
2063 ev => Poll::Ready(Some(ev)),
2065 }
2066 }
2067}
2068
2069impl fmt::Debug for Discv4Service {
2070 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2071 f.debug_struct("Discv4Service")
2072 .field("local_address", &self.local_address)
2073 .field("local_peer_id", &self.local_peer_id())
2074 .field("local_node_record", &self.local_node_record)
2075 .field("queued_pings", &self.queued_pings)
2076 .field("pending_lookup", &self.pending_lookup)
2077 .field("pending_find_nodes", &self.pending_find_nodes)
2078 .field("lookup_interval", &self.lookup_interval)
2079 .finish_non_exhaustive()
2080 }
2081}
2082
2083#[derive(Debug, Eq, PartialEq)]
2087pub enum Discv4Event {
2088 Ping,
2090 Pong,
2092 FindNode,
2094 Neighbours,
2096 EnrRequest,
2098 EnrResponse,
2100 Terminated,
2102}
2103
2104pub(crate) async fn send_loop(
2106 udp: Arc<UdpSocket>,
2107 secondary_udp: Option<Arc<UdpSocket>>,
2108 primary_is_ipv4: bool,
2109 rx: EgressReceiver,
2110) {
2111 let mut stream = ReceiverStream::new(rx);
2112 while let Some((payload, to)) = stream.next().await {
2113 let socket = if to.is_ipv4() == primary_is_ipv4 {
2116 &udp
2117 } else {
2118 secondary_udp.as_ref().unwrap_or(&udp)
2119 };
2120 match socket.send_to(&payload, to).await {
2121 Ok(size) => {
2122 trace!(target: "discv4", ?to, ?size,"sent payload");
2123 }
2124 Err(err) => {
2125 debug!(target: "discv4", ?to, %err,"Failed to send datagram.");
2126 }
2127 }
2128 }
2129}
2130
2131const MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP: usize = 60usize;
2133
2134pub(crate) async fn receive_loop(udp: Arc<UdpSocket>, tx: IngressSender, local_id: PeerId) {
2139 let mut handler = IngressHandler::new(tx, local_id);
2140 let mut buf = [0; MAX_PACKET_SIZE];
2141 loop {
2142 let res = udp.recv_from(&mut buf).await;
2143 match res {
2144 Err(err) => {
2145 debug!(target: "discv4", %err, "Failed to read datagram.");
2146 handler.send(IngressEvent::RecvError(err)).await;
2147 }
2148 Ok((read, remote_addr)) => {
2149 handler.handle_packet(&buf[..read], remote_addr).await;
2150 }
2151 }
2152 }
2153}
2154
2155#[derive(Debug)]
2160pub struct IngressHandler {
2161 tx: IngressSender,
2162 local_id: PeerId,
2163 tick: usize,
2164 tick_interval: Duration,
2165 cache: ReceiveCache,
2166 last_tick: Instant,
2167}
2168
2169impl IngressHandler {
2170 fn new(tx: IngressSender, local_id: PeerId) -> Self {
2171 let tick = MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP / 2;
2172 Self {
2173 tx,
2174 local_id,
2175 tick,
2176 tick_interval: Duration::from_secs(tick as u64),
2177 cache: ReceiveCache::default(),
2178 last_tick: Instant::now(),
2179 }
2180 }
2181
2182 async fn send(&self, event: IngressEvent) {
2183 let _ = self.tx.send(event).await.map_err(|err| {
2184 debug!(target: "discv4", %err, "failed send incoming packet");
2185 });
2186 }
2187
2188 pub async fn handle_packet(&mut self, data: &[u8], src: SocketAddr) {
2191 if self.last_tick.elapsed() >= self.tick_interval {
2192 self.cache.tick_ips(self.tick);
2193 self.last_tick = Instant::now();
2194 }
2195
2196 if self.cache.inc_ip(src.ip()) > MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP {
2198 trace!(target: "discv4", ?src, "Too many incoming packets from IP.");
2199 return
2200 }
2201
2202 if data.len() >= MIN_PACKET_SIZE &&
2209 self.cache.contains_packet(B256::from_slice(&data[..32]), src.ip())
2210 {
2211 trace!(target: "discv4", ?src, "Received duplicate packet.");
2212 return
2213 }
2214
2215 let event = match Message::decode(data) {
2216 Ok(packet) => {
2217 if packet.node_id == self.local_id {
2218 debug!(target: "discv4", ?src, "Received own packet.");
2219 return
2220 }
2221
2222 self.cache.insert_packet(packet.hash, src.ip());
2225
2226 IngressEvent::Packet(src, packet)
2227 }
2228 Err(err) => {
2229 trace!(target: "discv4", %err, "Failed to decode packet");
2230 IngressEvent::BadPacket(src, err, data.to_vec())
2231 }
2232 };
2233
2234 self.send(event).await;
2235 }
2236}
2237
2238#[derive(Debug)]
2242struct ReceiveCache {
2243 ip_messages: HashMap<IpAddr, usize>,
2249 unique_packets: reth_lru::LruMap<(B256, IpAddr), ()>,
2253}
2254
2255impl ReceiveCache {
2256 fn tick_ips(&mut self, tick: usize) {
2260 self.ip_messages.retain(|_, count| {
2261 if let Some(reset) = count.checked_sub(tick) {
2262 *count = reset;
2263 true
2264 } else {
2265 false
2266 }
2267 });
2268 }
2269
2270 fn inc_ip(&mut self, ip: IpAddr) -> usize {
2272 let ctn = self.ip_messages.entry(ip).or_default();
2273 *ctn = ctn.saturating_add(1);
2274 *ctn
2275 }
2276
2277 fn contains_packet(&mut self, hash: B256, ip: IpAddr) -> bool {
2281 self.unique_packets.get(&(hash, ip.to_canonical())).is_some()
2282 }
2283
2284 fn insert_packet(&mut self, hash: B256, ip: IpAddr) {
2286 self.unique_packets.insert((hash, ip.to_canonical()), ());
2287 }
2288}
2289
2290impl Default for ReceiveCache {
2291 fn default() -> Self {
2292 Self { ip_messages: Default::default(), unique_packets: reth_lru::LruMap::new(32) }
2293 }
2294}
2295
2296enum Discv4Command {
2298 Add(NodeRecord),
2299 AddBootNode(NodeRecord),
2300 SetTcpPort(u16),
2301 SetEIP868RLPPair { key: Vec<u8>, rlp: Bytes },
2302 Ban(PeerId, IpAddr),
2303 BanPeer(PeerId),
2304 BanIp(IpAddr),
2305 Remove(PeerId),
2306 Lookup { node_id: Option<PeerId>, tx: Option<NodeRecordSender> },
2307 SetLookupInterval(Duration),
2308 Updates(OneshotSender<ReceiverStream<DiscoveryUpdate>>),
2309 Terminated,
2310}
2311
2312#[derive(Debug)]
2314pub(crate) enum IngressEvent {
2315 RecvError(io::Error),
2317 BadPacket(SocketAddr, DecodePacketError, Vec<u8>),
2319 Packet(SocketAddr, Packet),
2321}
2322
2323#[derive(Debug)]
2325struct PingRequest {
2326 sent_at: Instant,
2328 node: NodeRecord,
2330 echo_hash: B256,
2332 reason: PingReason,
2334}
2335
2336#[derive(Debug)]
2340struct LookupTargetRotator {
2341 interval: usize,
2342 counter: usize,
2343}
2344
2345impl LookupTargetRotator {
2348 const fn local_only() -> Self {
2350 Self { interval: 1, counter: 0 }
2351 }
2352}
2353
2354impl Default for LookupTargetRotator {
2355 fn default() -> Self {
2356 Self {
2357 interval: 4,
2359 counter: 3,
2360 }
2361 }
2362}
2363
2364impl LookupTargetRotator {
2365 fn next(&mut self, local: &PeerId) -> PeerId {
2367 self.counter += 1;
2368 self.counter %= self.interval;
2369 if self.counter == 0 {
2370 return *local
2371 }
2372 PeerId::random()
2373 }
2374}
2375
2376#[derive(Clone, Debug)]
2381struct LookupContext {
2382 inner: Rc<LookupContextInner>,
2383}
2384
2385impl LookupContext {
2386 fn new(
2388 target: discv5::Key<NodeKey>,
2389 nearest_nodes: impl IntoIterator<Item = (Distance, NodeRecord)>,
2390 listener: Option<NodeRecordSender>,
2391 ) -> Self {
2392 let closest_nodes = nearest_nodes
2393 .into_iter()
2394 .map(|(distance, record)| {
2395 (distance, QueryNode { record, queried: false, responded: false })
2396 })
2397 .collect();
2398
2399 let inner = Rc::new(LookupContextInner {
2400 target,
2401 closest_nodes: RefCell::new(closest_nodes),
2402 listener,
2403 });
2404 Self { inner }
2405 }
2406
2407 fn target(&self) -> PeerId {
2409 self.inner.target.preimage().0
2410 }
2411
2412 fn closest(&self, num: usize) -> Vec<NodeRecord> {
2413 self.inner
2414 .closest_nodes
2415 .borrow()
2416 .iter()
2417 .filter(|(_, node)| !node.queried)
2418 .map(|(_, n)| n.record)
2419 .take(num)
2420 .collect()
2421 }
2422
2423 fn filter_closest<P>(&self, num: usize, filter: P) -> Vec<NodeRecord>
2425 where
2426 P: FnMut(&NodeRecord) -> bool,
2427 {
2428 self.inner
2429 .closest_nodes
2430 .borrow()
2431 .iter()
2432 .filter(|(_, node)| !node.queried)
2433 .map(|(_, n)| n.record)
2434 .filter(filter)
2435 .take(num)
2436 .collect()
2437 }
2438
2439 fn add_node(&self, record: NodeRecord) {
2441 let distance = self.inner.target.distance(&kad_key(record.id));
2442 let mut closest = self.inner.closest_nodes.borrow_mut();
2443 if let btree_map::Entry::Vacant(entry) = closest.entry(distance) {
2444 entry.insert(QueryNode { record, queried: false, responded: false });
2445 }
2446 }
2447
2448 fn set_queried(&self, id: PeerId, val: bool) {
2449 if let Some((_, node)) =
2450 self.inner.closest_nodes.borrow_mut().iter_mut().find(|(_, node)| node.record.id == id)
2451 {
2452 node.queried = val;
2453 }
2454 }
2455
2456 fn mark_queried(&self, id: PeerId) {
2458 self.set_queried(id, true)
2459 }
2460
2461 fn unmark_queried(&self, id: PeerId) {
2463 self.set_queried(id, false)
2464 }
2465
2466 fn mark_responded(&self, id: PeerId) {
2468 if let Some((_, node)) =
2469 self.inner.closest_nodes.borrow_mut().iter_mut().find(|(_, node)| node.record.id == id)
2470 {
2471 node.responded = true;
2472 }
2473 }
2474}
2475
2476unsafe impl Send for LookupContext {}
2483#[derive(Debug)]
2484struct LookupContextInner {
2485 target: discv5::Key<NodeKey>,
2487 closest_nodes: RefCell<BTreeMap<Distance, QueryNode>>,
2489 listener: Option<NodeRecordSender>,
2494}
2495
2496impl Drop for LookupContextInner {
2497 fn drop(&mut self) {
2498 if let Some(tx) = self.listener.take() {
2499 let nodes = self
2502 .closest_nodes
2503 .take()
2504 .into_values()
2505 .filter(|node| node.responded)
2506 .map(|node| node.record)
2507 .collect();
2508 let _ = tx.send(nodes);
2509 }
2510 }
2511}
2512
2513#[derive(Debug, Clone, Copy)]
2515struct QueryNode {
2516 record: NodeRecord,
2517 queried: bool,
2518 responded: bool,
2519}
2520
2521#[derive(Debug)]
2522struct FindNodeRequest {
2523 sent_at: Instant,
2525 response_count: usize,
2527 answered: bool,
2529 lookup_context: LookupContext,
2531 destination: IpAddr,
2533}
2534
2535impl FindNodeRequest {
2538 fn new(resp: LookupContext, destination: IpAddr) -> Self {
2539 Self {
2540 sent_at: Instant::now(),
2541 response_count: 0,
2542 answered: false,
2543 lookup_context: resp,
2544 destination: destination.to_canonical(),
2545 }
2546 }
2547
2548 fn is_from_destination(&self, ip: IpAddr) -> bool {
2553 self.destination == ip.to_canonical()
2554 }
2555}
2556
2557#[derive(Debug)]
2559struct CachedFindNode {
2560 target: PeerId,
2561 payload: Bytes,
2562 hash: B256,
2563 cached_at: Instant,
2564}
2565
2566impl CachedFindNode {
2567 fn get_or_sign(
2570 cache: &mut Option<Self>,
2571 target: PeerId,
2572 ttl: Duration,
2573 secret_key: &secp256k1::SecretKey,
2574 expire: u64,
2575 ) -> (Bytes, B256) {
2576 if let Some(c) = cache.as_ref() &&
2577 c.target == target &&
2578 c.cached_at.elapsed() < ttl
2579 {
2580 return (c.payload.clone(), c.hash);
2581 }
2582
2583 let msg = Message::FindNode(FindNode { id: target, expire });
2584 let (payload, hash) = msg.encode(secret_key);
2585
2586 *cache = Some(Self { target, payload: payload.clone(), hash, cached_at: Instant::now() });
2587
2588 (payload, hash)
2589 }
2590}
2591
2592#[derive(Debug)]
2593struct EnrRequestState {
2594 sent_at: Instant,
2596 echo_hash: B256,
2598}
2599
2600#[derive(Debug, Clone, Eq, PartialEq)]
2602struct NodeEntry {
2603 record: NodeRecord,
2605 last_seen: Instant,
2607 last_enr_seq: Option<u64>,
2609 fork_id: Option<ForkId>,
2611 find_node_failures: u8,
2613 has_endpoint_proof: bool,
2615}
2616
2617impl NodeEntry {
2620 fn new(record: NodeRecord) -> Self {
2622 Self {
2623 record,
2624 last_seen: Instant::now(),
2625 last_enr_seq: None,
2626 fork_id: None,
2627 find_node_failures: 0,
2628 has_endpoint_proof: false,
2629 }
2630 }
2631
2632 #[cfg(test)]
2633 fn new_proven(record: NodeRecord) -> Self {
2634 let mut node = Self::new(record);
2635 node.has_endpoint_proof = true;
2636 node
2637 }
2638
2639 const fn establish_proof(&mut self) {
2641 self.has_endpoint_proof = true;
2642 self.find_node_failures = 0;
2643 }
2644
2645 const fn exceeds_find_node_failures(&self, max_failures: u8) -> bool {
2647 self.find_node_failures >= max_failures
2648 }
2649
2650 fn update_with_enr(&mut self, last_enr_seq: Option<u64>) -> Option<u64> {
2652 self.update_now(|s| std::mem::replace(&mut s.last_enr_seq, last_enr_seq))
2653 }
2654
2655 const fn inc_failed_request(&mut self) {
2657 self.find_node_failures += 1;
2658 }
2659
2660 fn update_with_fork_id(&mut self, fork_id: Option<ForkId>) -> Option<ForkId> {
2662 self.update_now(|s| std::mem::replace(&mut s.fork_id, fork_id))
2663 }
2664
2665 fn update_now<F, R>(&mut self, f: F) -> R
2667 where
2668 F: FnOnce(&mut Self) -> R,
2669 {
2670 self.last_seen = Instant::now();
2671 f(self)
2672 }
2673}
2674
2675impl NodeEntry {
2678 fn is_expired(&self) -> bool {
2680 self.last_seen.elapsed() > (ENDPOINT_PROOF_EXPIRATION / 2)
2681 }
2682}
2683
2684#[derive(Debug)]
2686enum PingReason {
2687 InitialInsert,
2689 EstablishBond,
2691 RePing,
2693 Lookup(NodeRecord, LookupContext),
2695}
2696
2697#[derive(Debug, Clone)]
2699pub enum DiscoveryUpdate {
2700 Added(NodeRecord),
2702 DiscoveredAtCapacity(NodeRecord),
2704 EnrForkId(NodeRecord, ForkId),
2706 Removed(PeerId),
2708 Batch(Vec<Self>),
2710}
2711
2712#[cfg(test)]
2713mod tests {
2714 use super::*;
2715 use crate::test_utils::{create_discv4, create_discv4_with_config, rng_endpoint, rng_record};
2716 use alloy_primitives::hex;
2717 use alloy_rlp::{Decodable, Encodable};
2718 use rand::Rng;
2719 use reth_ethereum_forks::{EnrForkIdEntry, ForkHash};
2720 use reth_network_peers::mainnet_nodes;
2721 use std::{future::poll_fn, net::Ipv6Addr};
2722
2723 #[tokio::test]
2724 async fn test_duplicate_packet_rejected_without_decoding() {
2725 let secret_key = SecretKey::new(&mut rand::rng());
2726 let local_id = pk2id(&secret_key.public_key());
2727 let (tx, mut rx) = mpsc::channel(16);
2728 let mut handler = IngressHandler::new(tx, local_id);
2729
2730 let remote_key = SecretKey::new(&mut rand::rng());
2731 let msg = Message::Ping(Ping {
2732 from: rng_endpoint(&mut rand::rng()),
2733 to: rng_endpoint(&mut rand::rng()),
2734 expire: u64::MAX,
2735 enr_sq: None,
2736 });
2737 let (packet, hash) = msg.encode(&remote_key);
2738 let src = SocketAddr::from((Ipv4Addr::new(10, 0, 0, 1).to_ipv6_mapped(), 30303));
2739
2740 handler.handle_packet(&packet, src).await;
2741 assert!(matches!(rx.try_recv(), Ok(IngressEvent::Packet(_, _))));
2742
2743 for duplicate_src in [src, SocketAddr::from(([10, 0, 0, 1], 40404))] {
2745 handler.handle_packet(&packet, duplicate_src).await;
2746 assert_matches::assert_matches!(rx.try_recv(), Err(mpsc::error::TryRecvError::Empty));
2747 }
2748 assert!(handler.cache.contains_packet(hash, src.ip()));
2749 }
2750
2751 #[tokio::test]
2752 async fn test_undecodable_packet_does_not_poison_cache() {
2753 let secret_key = SecretKey::new(&mut rand::rng());
2754 let local_id = pk2id(&secret_key.public_key());
2755 let (tx, mut rx) = mpsc::channel(16);
2756 let mut handler = IngressHandler::new(tx, local_id);
2757
2758 let remote_key = SecretKey::new(&mut rand::rng());
2759 let msg = Message::Ping(Ping {
2760 from: rng_endpoint(&mut rand::rng()),
2761 to: rng_endpoint(&mut rand::rng()),
2762 expire: u64::MAX,
2763 enr_sq: None,
2764 });
2765 let (packet, hash) = msg.encode(&remote_key);
2766 let src = "10.0.0.1:30303".parse().unwrap();
2767
2768 let mut forged = packet.to_vec();
2771 let last = forged.len() - 1;
2772 forged[last] ^= 0xff;
2773 handler.handle_packet(&forged, src).await;
2774 assert!(matches!(rx.try_recv(), Ok(IngressEvent::BadPacket(..))));
2775 assert!(!handler.cache.contains_packet(hash, src.ip()));
2776
2777 handler.handle_packet(&packet, src).await;
2779 assert!(matches!(rx.try_recv(), Ok(IngressEvent::Packet(_, _))));
2780 }
2781
2782 #[tokio::test]
2783 async fn test_configured_enr_forkid_entry() {
2784 let fork: ForkId = ForkId { hash: ForkHash([220, 233, 108, 45]), next: 0u64 };
2785 let mut disc_conf = Discv4Config::default();
2786 disc_conf.add_eip868_pair("eth", EnrForkIdEntry::from(fork));
2787 let (_discv4, service) = create_discv4_with_config(disc_conf).await;
2788 let eth = service.local_eip_868_enr.get_raw_rlp(b"eth").unwrap();
2789 let fork_entry_id = EnrForkIdEntry::decode(&mut ð[..]).unwrap();
2790
2791 let raw: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2792 let decoded = EnrForkIdEntry::decode(&mut &raw[..]).unwrap();
2793 let expected = EnrForkIdEntry {
2794 fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2795 };
2796 assert_eq!(expected, fork_entry_id);
2797 assert_eq!(expected, decoded);
2798 }
2799
2800 #[tokio::test]
2801 async fn test_secondary_advertised_ip_in_enr() {
2802 let (Ok(secondary), Ok(secondary2)) =
2803 (UdpSocket::bind("[::1]:0").await, UdpSocket::bind("[::1]:0").await)
2804 else {
2805 return;
2807 };
2808 let mut rng = rand::rng();
2809 let (secret_key, pk) = secp256k1::generate_keypair(&mut rng);
2810 let local_enr = NodeRecord {
2811 address: Ipv4Addr::LOCALHOST.into(),
2812 tcp_port: 30303,
2813 udp_port: 30303,
2814 id: pk2id(&pk),
2815 };
2816
2817 let ip6: Ipv6Addr = "2001:db8::1".parse().unwrap();
2818 let config = Discv4Config::builder().secondary_advertised_ip(ip6.into()).build();
2819 let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
2820 let (_discv4, service, _ingress) =
2821 Discv4::bind_shared(socket, Some(Arc::new(secondary)), local_enr, secret_key, config)
2822 .unwrap();
2823
2824 let enr = &service.local_eip_868_enr;
2825 assert_eq!(enr.ip6(), Some(ip6));
2826 assert_eq!(enr.udp6(), Some(service.local_node_record.udp_port));
2827 assert_eq!(enr.tcp6(), Some(service.local_node_record.tcp_port));
2828
2829 let config =
2831 Discv4Config::builder().secondary_advertised_ip(Ipv6Addr::UNSPECIFIED.into()).build();
2832 let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
2833 let (_discv4, service, _ingress) =
2834 Discv4::bind_shared(socket, Some(Arc::new(secondary2)), local_enr, secret_key, config)
2835 .unwrap();
2836
2837 let enr = &service.local_eip_868_enr;
2838 assert_eq!(enr.ip6(), None);
2839 assert_eq!(enr.udp6(), None);
2840 assert_eq!(enr.tcp6(), None);
2841 }
2842
2843 #[tokio::test]
2844 async fn test_secondary_advertised_ip_requires_secondary_socket() {
2845 let ip6: Ipv6Addr = "2001:db8::1".parse().unwrap();
2846 let config = Discv4Config::builder().secondary_advertised_ip(ip6.into()).build();
2847 let (_discv4, service) = create_discv4_with_config(config).await;
2848
2849 let enr = &service.local_eip_868_enr;
2850 assert_eq!(enr.ip6(), None);
2851 assert_eq!(enr.udp6(), None);
2852 assert_eq!(enr.tcp6(), None);
2853 }
2854
2855 #[tokio::test]
2856 async fn test_external_ip_completes_unspecified_enr_endpoint() {
2857 let mut rng = rand::rng();
2858 let socket: SocketAddr = "127.0.0.1:0".parse().unwrap();
2859 let (secret_key, pk) = secp256k1::generate_keypair(&mut rng);
2860 let local_enr = NodeRecord {
2861 address: Ipv4Addr::UNSPECIFIED.into(),
2862 tcp_port: 30303,
2863 udp_port: 30303,
2864 id: pk2id(&pk),
2865 };
2866 let (_discv4, mut service) =
2867 Discv4::bind(socket, local_enr, secret_key, Default::default()).await.unwrap();
2868
2869 let udp_port = service.local_node_record.udp_port;
2870 let enr = &service.local_eip_868_enr;
2871 assert_eq!(enr.ip4(), None);
2872 assert_eq!(enr.udp4(), Some(udp_port));
2873 assert_eq!(enr.tcp4(), Some(30303));
2874
2875 let external_ip: Ipv4Addr = "1.2.3.4".parse().unwrap();
2876 service.set_external_ip_addr(external_ip.into());
2877
2878 let enr = &service.local_eip_868_enr;
2879 assert_eq!(enr.ip4(), Some(external_ip));
2880 assert_eq!(enr.udp4(), Some(udp_port));
2881 assert_eq!(enr.tcp4(), Some(30303));
2882 }
2883
2884 #[tokio::test]
2885 async fn test_send_loop_routes_by_destination_family() {
2886 let (Ok(secondary), Ok(sink6)) =
2887 (UdpSocket::bind("[::1]:0").await, UdpSocket::bind("[::1]:0").await)
2888 else {
2889 return;
2891 };
2892 let primary = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
2893 let sink4 = UdpSocket::bind("127.0.0.1:0").await.unwrap();
2894
2895 let (tx, rx) = mpsc::channel(8);
2896 tokio::spawn(send_loop(primary, Some(Arc::new(secondary)), true, rx));
2897
2898 tx.send((Bytes::from_static(b"v4"), sink4.local_addr().unwrap())).await.unwrap();
2899 tx.send((Bytes::from_static(b"v6"), sink6.local_addr().unwrap())).await.unwrap();
2900
2901 let mut buf = [0u8; 2];
2902 let (read, _) = tokio::time::timeout(Duration::from_secs(5), sink4.recv_from(&mut buf))
2903 .await
2904 .unwrap()
2905 .unwrap();
2906 assert_eq!(&buf[..read], b"v4");
2907 let (read, _) = tokio::time::timeout(Duration::from_secs(5), sink6.recv_from(&mut buf))
2908 .await
2909 .unwrap()
2910 .unwrap();
2911 assert_eq!(&buf[..read], b"v6");
2912 }
2913
2914 #[test]
2915 fn test_enr_forkid_entry_decode() {
2916 let raw: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2917 let decoded = EnrForkIdEntry::decode(&mut &raw[..]).unwrap();
2918 let expected = EnrForkIdEntry {
2919 fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2920 };
2921 assert_eq!(expected, decoded);
2922 }
2923
2924 #[test]
2925 fn test_enr_forkid_entry_encode() {
2926 let original = EnrForkIdEntry {
2927 fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2928 };
2929 let expected: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2930 let mut encoded = Vec::with_capacity(expected.len());
2931 original.encode(&mut encoded);
2932 assert_eq!(&expected[..], encoded.as_slice());
2933 }
2934
2935 #[tokio::test]
2936 async fn test_request_enr_of_proven_node() {
2937 reth_tracing::init_test_tracing();
2938 let mut rng = rand::rng();
2939 let (_discv4, mut service) = create_discv4().await;
2940
2941 let id = PeerId::random();
2942 let addr: SocketAddr = (Ipv4Addr::new(10, 0, 0, 1), DEFAULT_DISCOVERY_PORT).into();
2943 let mut entry = NodeEntry::new_proven(NodeRecord::new(addr, id));
2944 entry.last_enr_seq = Some(1);
2945 let _ = service.kbuckets.insert_or_update(
2946 &kad_key(id),
2947 entry,
2948 NodeStatus {
2949 direction: ConnectionDirection::Outgoing,
2950 state: ConnectionState::Connected,
2951 },
2952 );
2953 service.received_pongs.on_pong(id, addr.ip());
2954 let ping = |service: &Discv4Service, rng: &mut rand::rngs::ThreadRng| Ping {
2955 from: rng_endpoint(rng),
2956 to: rng_endpoint(rng),
2957 expire: service.ping_expiration(),
2958 enr_sq: Some(2),
2959 };
2960
2961 service.on_ping(ping(&service, &mut rng), addr, id, B256::random());
2963 assert!(service.pending_enr_requests.contains_key(&id));
2964
2965 service.evict_expired_requests(Instant::now() + service.config.enr_expiration * 2);
2967 assert!(!service.pending_enr_requests.contains_key(&id));
2968 service.on_ping(ping(&service, &mut rng), addr, id, B256::random());
2969 assert!(service.pending_enr_requests.contains_key(&id));
2970 }
2971
2972 #[test]
2973 fn test_local_rotator() {
2974 let id = PeerId::random();
2975 let mut rotator = LookupTargetRotator::local_only();
2976 assert_eq!(rotator.next(&id), id);
2977 assert_eq!(rotator.next(&id), id);
2978 }
2979
2980 #[test]
2981 fn test_rotator() {
2982 let id = PeerId::random();
2983 let mut rotator = LookupTargetRotator::default();
2984 assert_eq!(rotator.next(&id), id);
2985 assert_ne!(rotator.next(&id), id);
2986 assert_ne!(rotator.next(&id), id);
2987 assert_ne!(rotator.next(&id), id);
2988 assert_eq!(rotator.next(&id), id);
2989 }
2990
2991 #[tokio::test]
2992 async fn test_ping_establishes_bond_per_ip() {
2993 for (primary, secondary) in
2994 [("127.0.0.1:30303", "[::1]:30303"), ("[::1]:30303", "127.0.0.1:30303")]
2995 {
2996 for initial_ping_pending in [false, true] {
2997 let (_, mut service) = create_discv4().await;
2998 let (egress, mut packets) = mpsc::channel(16);
2999 service.egress = egress;
3000 let primary = NodeRecord::new(primary.parse().unwrap(), PeerId::random());
3001 let secondary = NodeRecord::new(secondary.parse().unwrap(), primary.id);
3002 service.add_node(primary);
3003 let pong = Pong {
3004 to: service.local_node_record.into(),
3005 echo: service.pending_pings[&primary.id].echo_hash,
3006 expire: service.ping_expiration(),
3007 enr_sq: None,
3008 };
3009 if !initial_ping_pending {
3010 service.on_pong(pong.clone(), primary.udp_addr(), primary.id);
3011 }
3012
3013 let ping = Ping {
3014 from: secondary.into(),
3015 to: service.local_node_record.into(),
3016 expire: service.ping_expiration(),
3017 enr_sq: None,
3018 };
3019 service.on_ping(ping.clone(), secondary.udp_addr(), secondary.id, B256::random());
3020 if initial_ping_pending {
3021 assert_eq!(service.queued_pings.len(), 1);
3022 service.ping_buffered();
3023 assert_eq!(service.queued_pings.len(), 1);
3024 service.on_pong(pong, primary.udp_addr(), primary.id);
3025 service.ping_buffered();
3026 }
3027
3028 assert!(service.has_bond(primary.id, primary.address));
3029 assert!(!service.has_bond(secondary.id, secondary.address));
3030 let request = &service.pending_pings[&secondary.id];
3031 assert_eq!(request.node.address, secondary.address);
3032 assert!(matches!(request.reason, PingReason::EstablishBond));
3033 let pong = Pong {
3034 to: service.local_node_record.into(),
3035 echo: request.echo_hash,
3036 expire: service.ping_expiration(),
3037 enr_sq: None,
3038 };
3039 while packets.try_recv().is_ok() {}
3040 let find_node =
3041 FindNode { id: PeerId::random(), expire: service.ping_expiration() };
3042 service.on_find_node(find_node, secondary.udp_addr(), secondary.id);
3043 assert!(packets.try_recv().is_err());
3044
3045 service.on_pong(pong, secondary.udp_addr(), secondary.id);
3046 assert!(service.has_bond(primary.id, primary.address));
3047 assert!(service.has_bond(secondary.id, secondary.address));
3048 service.on_find_node(find_node, secondary.udp_addr(), secondary.id);
3049 let (packet, to) = packets.try_recv().unwrap();
3050 assert_eq!(to, secondary.udp_addr());
3051 assert!(matches!(Message::decode(&packet).unwrap().msg, Message::Neighbours(_)));
3052
3053 service.on_ping(ping, secondary.udp_addr(), secondary.id, B256::random());
3054 assert!(service.pending_pings.is_empty());
3055 assert!(service.queued_pings.is_empty());
3056 }
3057 }
3058 }
3059
3060 #[tokio::test]
3061 async fn test_neighbours_exclude_unproven_nodes() {
3062 let (_, mut service) = create_discv4().await;
3063 let (egress, mut packets) = mpsc::channel(16);
3064 service.egress = egress;
3065 let proven = NodeRecord::new("127.0.0.1:30303".parse().unwrap(), PeerId::random());
3066 insert_proven_node(&mut service, proven);
3067
3068 let unproven = NodeRecord::new("127.0.0.2:30303".parse().unwrap(), PeerId::random());
3070 let ping = Ping {
3071 from: unproven.into(),
3072 to: service.local_node_record.into(),
3073 expire: service.ping_expiration(),
3074 enr_sq: None,
3075 };
3076 service.on_ping(ping, unproven.udp_addr(), unproven.id, B256::random());
3077 assert_eq!(service.on_entry(unproven.id, |entry| entry.has_endpoint_proof), Some(false));
3078 while packets.try_recv().is_ok() {}
3079
3080 let find_node = FindNode { id: unproven.id, expire: service.ping_expiration() };
3081 service.on_find_node(find_node, proven.udp_addr(), proven.id);
3082 let (packet, to) = packets.try_recv().unwrap();
3083 assert_eq!(to, proven.udp_addr());
3084 let Message::Neighbours(neighbours) = Message::decode(&packet).unwrap().msg else {
3085 panic!("expected a neighbours packet")
3086 };
3087 assert_eq!(neighbours.nodes, vec![proven]);
3088 }
3089
3090 #[tokio::test]
3091 async fn test_ping_advertises_tcp_port_per_family() {
3092 let (_, mut service) = create_discv4().await;
3093 let (egress, mut packets) = mpsc::channel(16);
3094 service.egress = egress;
3095 service.local_node_record.tcp_port = 30303;
3096 service.local_eip_868_enr.set_tcp4(30303, &service.secret_key).unwrap();
3097
3098 for (tcp6, expected_tcp6) in [(None, 0), (Some(30304), 30304), (Some(0), 0)] {
3099 if let Some(port) = tcp6 {
3100 service.local_eip_868_enr.set_tcp6(port, &service.secret_key).unwrap();
3101 }
3102 for (addr, expected) in [("127.0.0.1:30303", 30303), ("[::1]:30303", expected_tcp6)] {
3103 let record = NodeRecord::new(addr.parse().unwrap(), PeerId::random());
3104 service.send_ping(record, PingReason::InitialInsert);
3105 let (packet, to) = packets.try_recv().unwrap();
3106 assert_eq!(to, record.udp_addr());
3107 let Message::Ping(ping) = Message::decode(&packet).unwrap().msg else {
3108 panic!("expected ping");
3109 };
3110 assert_eq!(ping.from.tcp_port, expected);
3111 }
3112 }
3113 }
3114
3115 #[tokio::test]
3116 async fn test_pending_ping() {
3117 let (_, mut service) = create_discv4().await;
3118
3119 let local_addr = service.local_addr();
3120
3121 let mut num_inserted = 0;
3122 loop {
3123 let node = NodeRecord::new(local_addr, PeerId::random());
3124 if service.add_node(node) {
3125 num_inserted += 1;
3126 assert!(service.pending_pings.contains_key(&node.id));
3127 assert_eq!(service.pending_pings.len(), num_inserted);
3128 if num_inserted == MAX_NODES_PING {
3129 break
3130 }
3131 }
3132 }
3133
3134 num_inserted = 0;
3136 for _ in 0..MAX_NODES_PING {
3137 let node = NodeRecord::new(local_addr, PeerId::random());
3138 if service.add_node(node) {
3139 num_inserted += 1;
3140 assert!(!service.pending_pings.contains_key(&node.id));
3141 assert_eq!(service.pending_pings.len(), MAX_NODES_PING);
3142 assert_eq!(service.queued_pings.len(), num_inserted);
3143 }
3144 }
3145 }
3146
3147 #[tokio::test(flavor = "multi_thread")]
3149 #[ignore]
3150 async fn test_mainnet_lookup() {
3151 reth_tracing::init_test_tracing();
3152 let fork_id = ForkId { hash: ForkHash(hex!("743f3d89")), next: 16191202 };
3153
3154 let all_nodes = mainnet_nodes();
3155 let config = Discv4Config::builder()
3156 .add_boot_nodes(all_nodes)
3157 .lookup_interval(Duration::from_secs(1))
3158 .add_eip868_pair("eth", fork_id)
3159 .build();
3160 let (_discv4, mut service) = create_discv4_with_config(config).await;
3161
3162 let mut updates = service.update_stream();
3163
3164 let _handle = service.spawn();
3165
3166 let mut table = HashMap::new();
3167 while let Some(update) = updates.next().await {
3168 match update {
3169 DiscoveryUpdate::EnrForkId(record, fork_id) => {
3170 println!("{record:?}, {fork_id:?}");
3171 }
3172 DiscoveryUpdate::Added(record) => {
3173 table.insert(record.id, record);
3174 }
3175 DiscoveryUpdate::Removed(id) => {
3176 table.remove(&id);
3177 }
3178 _ => {}
3179 }
3180 println!("total peers {}", table.len());
3181 }
3182 }
3183
3184 #[tokio::test]
3185 async fn test_mapped_ipv4() {
3186 reth_tracing::init_test_tracing();
3187 let mut rng = rand::rng();
3188 let config = Discv4Config::builder().build();
3189 let (_discv4, mut service) = create_discv4_with_config(config).await;
3190
3191 let v4: Ipv4Addr = "0.0.0.0".parse().unwrap();
3192 let v6 = v4.to_ipv6_mapped();
3193 let addr: SocketAddr = (v6, DEFAULT_DISCOVERY_PORT).into();
3194
3195 let ping = Ping {
3196 from: rng_endpoint(&mut rng),
3197 to: rng_endpoint(&mut rng),
3198 expire: service.ping_expiration(),
3199 enr_sq: Some(rng.random()),
3200 };
3201
3202 let id = PeerId::random();
3203 service.on_ping(ping, addr, id, B256::random());
3204
3205 let key = kad_key(id);
3206 match service.kbuckets.entry(&key) {
3207 kbucket::Entry::Present(entry, _) => {
3208 let node_addr = entry.value().record.address;
3209 assert!(node_addr.is_ipv4());
3210 assert_eq!(node_addr, IpAddr::from(v4));
3211 }
3212 _ => unreachable!(),
3213 };
3214 }
3215
3216 #[tokio::test]
3217 async fn test_respect_ping_expiration() {
3218 reth_tracing::init_test_tracing();
3219 let mut rng = rand::rng();
3220 let config = Discv4Config::builder().build();
3221 let (_discv4, mut service) = create_discv4_with_config(config).await;
3222
3223 let v4: Ipv4Addr = "0.0.0.0".parse().unwrap();
3224 let v6 = v4.to_ipv6_mapped();
3225 let addr: SocketAddr = (v6, DEFAULT_DISCOVERY_PORT).into();
3226
3227 let ping = Ping {
3228 from: rng_endpoint(&mut rng),
3229 to: rng_endpoint(&mut rng),
3230 expire: SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() - 1,
3231 enr_sq: Some(rng.random()),
3232 };
3233
3234 let id = PeerId::random();
3235 service.on_ping(ping, addr, id, B256::random());
3236
3237 let key = kad_key(id);
3238 match service.kbuckets.entry(&key) {
3239 kbucket::Entry::Absent(_) => {}
3240 _ => unreachable!(),
3241 };
3242 }
3243
3244 #[tokio::test]
3245 async fn test_single_lookups() {
3246 reth_tracing::init_test_tracing();
3247
3248 let config = Discv4Config::builder().build();
3249 let (_discv4, mut service) = create_discv4_with_config(config.clone()).await;
3250
3251 let id = PeerId::random();
3252 let key = kad_key(id);
3253 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3254
3255 let _ = service.kbuckets.insert_or_update(
3256 &key,
3257 NodeEntry::new_proven(record),
3258 NodeStatus {
3259 direction: ConnectionDirection::Incoming,
3260 state: ConnectionState::Connected,
3261 },
3262 );
3263
3264 service.lookup_self();
3265 assert_eq!(service.pending_find_nodes.len(), 1);
3266
3267 poll_fn(|cx| {
3268 let _ = service.poll(cx);
3269 assert_eq!(service.pending_find_nodes.len(), 1);
3270
3271 Poll::Ready(())
3272 })
3273 .await;
3274 }
3275
3276 #[tokio::test]
3277 async fn test_on_neighbours_recursive_lookup() {
3278 reth_tracing::init_test_tracing();
3279
3280 let config = Discv4Config::builder().build();
3281 let (_discv4, mut service) = create_discv4_with_config(config.clone()).await;
3282 let (_discv4, mut service2) = create_discv4_with_config(config).await;
3283
3284 let id = PeerId::random();
3285 let key = kad_key(id);
3286 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3287
3288 let _ = service.kbuckets.insert_or_update(
3289 &key,
3290 NodeEntry::new_proven(record),
3291 NodeStatus {
3292 direction: ConnectionDirection::Incoming,
3293 state: ConnectionState::Connected,
3294 },
3295 );
3296 service.lookup_self();
3299 assert_eq!(service.pending_find_nodes.len(), 1);
3300
3301 poll_fn(|cx| {
3302 let _ = service.poll(cx);
3303 assert_eq!(service.pending_find_nodes.len(), 1);
3304
3305 Poll::Ready(())
3306 })
3307 .await;
3308
3309 let expiry = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs() +
3310 10000000000000;
3311 let msg = Neighbours { nodes: vec![service2.local_node_record], expire: expiry };
3312 service.on_neighbours(msg, record.tcp_addr(), id);
3313 let event = poll_fn(|cx| service2.poll(cx)).await;
3315 assert_eq!(event, Discv4Event::Ping);
3316 assert_eq!(service.pending_find_nodes.len(), 1);
3319 let event = poll_fn(|cx| service.poll(cx)).await;
3321 assert_eq!(event, Discv4Event::Pong);
3322 let event = poll_fn(|cx| service.poll(cx)).await;
3327 assert_eq!(event, Discv4Event::Ping);
3328 assert_eq!(service.pending_find_nodes.len(), 2);
3331 }
3332
3333 #[tokio::test]
3334 async fn test_no_local_in_closest() {
3335 reth_tracing::init_test_tracing();
3336
3337 let config = Discv4Config::builder().build();
3338 let (_discv4, mut service) = create_discv4_with_config(config).await;
3339
3340 let target_key = kad_key(PeerId::random());
3341
3342 let id = PeerId::random();
3343 let key = kad_key(id);
3344 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3345
3346 let _ = service.kbuckets.insert_or_update(
3347 &key,
3348 NodeEntry::new(record),
3349 NodeStatus {
3350 direction: ConnectionDirection::Incoming,
3351 state: ConnectionState::Connected,
3352 },
3353 );
3354
3355 let closest = service
3356 .kbuckets
3357 .closest_values(&target_key)
3358 .map(|n| n.value.record)
3359 .take(MAX_NODES_PER_BUCKET)
3360 .collect::<Vec<_>>();
3361
3362 assert_eq!(closest.len(), 1);
3363 assert!(!closest.iter().any(|r| r.id == *service.local_peer_id()));
3364 }
3365
3366 #[tokio::test]
3367 async fn test_random_lookup() {
3368 reth_tracing::init_test_tracing();
3369
3370 let config = Discv4Config::builder().build();
3371 let (_discv4, mut service) = create_discv4_with_config(config).await;
3372
3373 let target = PeerId::random();
3374
3375 let id = PeerId::random();
3376 let key = kad_key(id);
3377 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3378
3379 let _ = service.kbuckets.insert_or_update(
3380 &key,
3381 NodeEntry::new_proven(record),
3382 NodeStatus {
3383 direction: ConnectionDirection::Incoming,
3384 state: ConnectionState::Connected,
3385 },
3386 );
3387
3388 service.lookup(target);
3389 assert_eq!(service.pending_find_nodes.len(), 1);
3390
3391 let ctx = service.pending_find_nodes.values().next().unwrap().lookup_context.clone();
3392
3393 assert_eq!(ctx.target(), target);
3394 assert_eq!(ctx.inner.closest_nodes.borrow().len(), 1);
3395
3396 ctx.add_node(record);
3397 assert_eq!(ctx.inner.closest_nodes.borrow().len(), 1);
3398 }
3399
3400 #[tokio::test]
3401 async fn test_reping_on_find_node_failures() {
3402 reth_tracing::init_test_tracing();
3403
3404 let config = Discv4Config::builder().build();
3405 let (_discv4, mut service) = create_discv4_with_config(config).await;
3406
3407 let target = PeerId::random();
3408
3409 let id = PeerId::random();
3410 let key = kad_key(id);
3411 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3412
3413 let mut entry = NodeEntry::new_proven(record);
3414 entry.find_node_failures = u8::MAX;
3415 let _ = service.kbuckets.insert_or_update(
3416 &key,
3417 entry,
3418 NodeStatus {
3419 direction: ConnectionDirection::Incoming,
3420 state: ConnectionState::Connected,
3421 },
3422 );
3423
3424 service.lookup(target);
3425 assert_eq!(service.pending_find_nodes.len(), 0);
3426 assert_eq!(service.pending_pings.len(), 1);
3427
3428 service.update_on_pong(record, None);
3429
3430 service
3431 .on_entry(record.id, |entry| {
3432 assert_eq!(entry.find_node_failures, 0);
3434 assert!(entry.has_endpoint_proof);
3435 })
3436 .unwrap();
3437 }
3438
3439 #[tokio::test]
3440 async fn test_service_commands() {
3441 reth_tracing::init_test_tracing();
3442
3443 let config = Discv4Config::builder().build();
3444 let (discv4, mut service) = create_discv4_with_config(config).await;
3445
3446 service.lookup_self();
3447
3448 let _handle = service.spawn();
3449 discv4.send_lookup_self();
3450 let _ = discv4.lookup_self().await;
3451 }
3452
3453 #[tokio::test]
3454 async fn test_requests_timeout() {
3455 reth_tracing::init_test_tracing();
3456 let fork_id = ForkId { hash: ForkHash(hex!("743f3d89")), next: 16191202 };
3457
3458 let config = Discv4Config::builder()
3459 .request_timeout(Duration::from_millis(200))
3460 .ping_expiration(Duration::from_millis(200))
3461 .lookup_neighbours_expiration(Duration::from_millis(200))
3462 .add_eip868_pair("eth", fork_id)
3463 .build();
3464 let (_disv4, mut service) = create_discv4_with_config(config).await;
3465
3466 let id = PeerId::random();
3467 let key = kad_key(id);
3468 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3469
3470 let _ = service.kbuckets.insert_or_update(
3471 &key,
3472 NodeEntry::new_proven(record),
3473 NodeStatus {
3474 direction: ConnectionDirection::Incoming,
3475 state: ConnectionState::Connected,
3476 },
3477 );
3478
3479 service.lookup_self();
3480 assert_eq!(service.pending_find_nodes.len(), 1);
3481
3482 let ctx = service.pending_find_nodes.values().next().unwrap().lookup_context.clone();
3483
3484 service.pending_lookup.insert(record.id, (Instant::now(), ctx));
3485
3486 assert_eq!(service.pending_lookup.len(), 1);
3487
3488 let ping = Ping {
3489 from: service.local_node_record.into(),
3490 to: record.into(),
3491 expire: service.ping_expiration(),
3492 enr_sq: service.enr_seq(),
3493 };
3494 let echo_hash = service.send_packet(Message::Ping(ping), record.udp_addr());
3495 let ping_request = PingRequest {
3496 sent_at: Instant::now(),
3497 node: record,
3498 echo_hash,
3499 reason: PingReason::InitialInsert,
3500 };
3501 service.pending_pings.insert(record.id, ping_request);
3502
3503 assert_eq!(service.pending_pings.len(), 1);
3504
3505 tokio::time::sleep(Duration::from_secs(1)).await;
3506
3507 poll_fn(|cx| {
3508 let _ = service.poll(cx);
3509
3510 assert_eq!(service.pending_find_nodes.len(), 0);
3511 assert_eq!(service.pending_lookup.len(), 0);
3512 assert_eq!(service.pending_pings.len(), 0);
3513
3514 Poll::Ready(())
3515 })
3516 .await;
3517 }
3518
3519 #[tokio::test(flavor = "multi_thread")]
3521 async fn test_check_wrong_to() {
3522 reth_tracing::init_test_tracing();
3523
3524 let config = Discv4Config::builder().external_ip_resolver(None).build();
3525 let (_discv4, mut service_1) = create_discv4_with_config(config.clone()).await;
3526 let (_discv4, mut service_2) = create_discv4_with_config(config).await;
3527
3528 let mut ping = Ping {
3530 from: service_1.local_node_record.into(),
3531 to: service_2.local_node_record.into(),
3532 expire: service_1.ping_expiration(),
3533 enr_sq: service_1.enr_seq(),
3534 };
3535 ping.to.address = "192.0.2.0".parse().unwrap();
3536
3537 let echo_hash = service_1.send_packet(Message::Ping(ping), service_2.local_addr());
3538 let ping_request = PingRequest {
3539 sent_at: Instant::now(),
3540 node: service_2.local_node_record,
3541 echo_hash,
3542 reason: PingReason::InitialInsert,
3543 };
3544 service_1.pending_pings.insert(*service_2.local_peer_id(), ping_request);
3545
3546 let event = poll_fn(|cx| service_2.poll(cx)).await;
3548 assert_eq!(event, Discv4Event::Ping);
3549
3550 let event = poll_fn(|cx| service_1.poll(cx)).await;
3552 assert_eq!(event, Discv4Event::Pong);
3553 let event = poll_fn(|cx| service_1.poll(cx)).await;
3555 assert_eq!(event, Discv4Event::Ping);
3556 }
3557
3558 #[tokio::test(flavor = "multi_thread")]
3559 async fn test_check_ping_pong() {
3560 reth_tracing::init_test_tracing();
3561
3562 let config = Discv4Config::builder().external_ip_resolver(None).build();
3563 let (_discv4, mut service_1) = create_discv4_with_config(config.clone()).await;
3564 let (_discv4, mut service_2) = create_discv4_with_config(config).await;
3565
3566 service_1.add_node(service_2.local_node_record);
3568
3569 let event = poll_fn(|cx| service_2.poll(cx)).await;
3571 assert_eq!(event, Discv4Event::Ping);
3572
3573 let key1 = kad_key(*service_1.local_peer_id());
3575 match service_2.kbuckets.entry(&key1) {
3576 kbucket::Entry::Present(_entry, status) => {
3577 assert!(!status.is_connected());
3578 }
3579 _ => unreachable!(),
3580 }
3581
3582 let event = poll_fn(|cx| service_1.poll(cx)).await;
3584 assert_eq!(event, Discv4Event::Pong);
3585
3586 let key2 = kad_key(*service_2.local_peer_id());
3588 match service_1.kbuckets.entry(&key2) {
3589 kbucket::Entry::Present(_entry, status) => {
3590 assert!(status.is_connected());
3591 }
3592 _ => unreachable!(),
3593 }
3594
3595 let event = poll_fn(|cx| service_1.poll(cx)).await;
3597 assert_eq!(event, Discv4Event::Ping);
3598
3599 tokio::time::timeout(Duration::from_secs(5), async {
3603 loop {
3604 let event = poll_fn(|cx| service_2.poll(cx)).await;
3605 match event {
3606 Discv4Event::Pong => break,
3607 Discv4Event::EnrRequest | Discv4Event::FindNode => {}
3608 ev => unreachable!("{ev:?}"),
3609 }
3610 }
3611 })
3612 .await
3613 .expect("timed out waiting for Pong from service_2");
3614
3615 match service_2.kbuckets.entry(&key1) {
3617 kbucket::Entry::Present(_entry, status) => {
3618 assert!(status.is_connected());
3619 }
3620 ev => unreachable!("{ev:?}"),
3621 }
3622 }
3623
3624 #[test]
3625 fn test_insert() {
3626 let local_node_record = rng_record(&mut rand::rng());
3627 let mut kbuckets: KBucketsTable<NodeKey, NodeEntry> = KBucketsTable::new(
3628 NodeKey::from(&local_node_record).into(),
3629 Duration::from_secs(60),
3630 MAX_NODES_PER_BUCKET,
3631 None,
3632 None,
3633 );
3634
3635 let new_record = rng_record(&mut rand::rng());
3636 let key = kad_key(new_record.id);
3637 match kbuckets.entry(&key) {
3638 kbucket::Entry::Absent(entry) => {
3639 let node = NodeEntry::new(new_record);
3640 let _ = entry.insert(
3641 node,
3642 NodeStatus {
3643 direction: ConnectionDirection::Outgoing,
3644 state: ConnectionState::Disconnected,
3645 },
3646 );
3647 }
3648 _ => {
3649 unreachable!()
3650 }
3651 };
3652 match kbuckets.entry(&key) {
3653 kbucket::Entry::Present(_, _) => {}
3654 _ => {
3655 unreachable!()
3656 }
3657 }
3658 }
3659
3660 #[tokio::test]
3661 async fn test_bootnode_in_update_stream() {
3662 reth_tracing::init_test_tracing();
3663 let (_, service_1) = create_discv4().await;
3664 let peerid_1 = *service_1.local_peer_id();
3665
3666 let config = Discv4Config::builder().add_boot_node(service_1.local_node_record).build();
3667 service_1.spawn();
3668
3669 let (_, mut service_2) = create_discv4_with_config(config).await;
3670
3671 let mut updates = service_2.update_stream();
3672
3673 service_2.spawn();
3674
3675 let mut bootnode_appeared = false;
3677 let timeout = tokio::time::sleep(Duration::from_secs(1));
3678 tokio::pin!(timeout);
3679
3680 loop {
3681 tokio::select! {
3682 Some(update) = updates.next() => {
3683 if let DiscoveryUpdate::Added(record) = update
3684 && record.id == peerid_1 {
3685 bootnode_appeared = true;
3686 break;
3687 }
3688 }
3689 _ = &mut timeout => break,
3690 }
3691 }
3692
3693 assert!(bootnode_appeared, "Bootnode should appear in update stream");
3694 }
3695
3696 fn insert_proven_node(service: &mut Discv4Service, record: NodeRecord) {
3697 service.received_pongs.on_pong(record.id, record.address);
3698 let key = kad_key(record.id);
3699 let _ = service.kbuckets.insert_or_update(
3700 &key,
3701 NodeEntry::new_proven(record),
3702 NodeStatus {
3703 direction: ConnectionDirection::Incoming,
3704 state: ConnectionState::Connected,
3705 },
3706 );
3707 }
3708
3709 fn insert_initial_ping(service: &mut Discv4Service, record: NodeRecord) -> B256 {
3710 let echo_hash = B256::random();
3711 service.pending_pings.insert(
3712 record.id,
3713 PingRequest {
3714 sent_at: Instant::now(),
3715 node: record,
3716 echo_hash,
3717 reason: PingReason::InitialInsert,
3718 },
3719 );
3720 echo_hash
3721 }
3722
3723 fn make_pong(service: &Discv4Service, echo_hash: B256) -> Pong {
3724 Pong {
3725 to: rng_endpoint(&mut rand::rng()),
3726 echo: echo_hash,
3727 expire: service.ping_expiration(),
3728 enr_sq: None,
3729 }
3730 }
3731
3732 #[tokio::test]
3733 async fn test_pong_from_unpinged_ip_does_not_establish_bond() {
3734 let (_discv4, mut service) = create_discv4().await;
3735 let record = NodeRecord::new("127.0.0.1:30303".parse().unwrap(), PeerId::random());
3736 let echo_hash = insert_initial_ping(&mut service, record);
3737 let unexpected_addr = "127.0.0.2:30303".parse().unwrap();
3738
3739 service.on_pong(make_pong(&service, echo_hash), unexpected_addr, record.id);
3740
3741 assert!(service.pending_pings.contains_key(&record.id));
3742 assert!(!service.has_bond(record.id, unexpected_addr.ip()));
3743
3744 let expected_addr = SocketAddr::from((Ipv4Addr::LOCALHOST.to_ipv6_mapped(), 30304));
3745 service.on_pong(make_pong(&service, echo_hash), expected_addr, record.id);
3746
3747 assert!(!service.pending_pings.contains_key(&record.id));
3748 assert!(service.has_bond(record.id, expected_addr.ip()));
3749 }
3750
3751 #[tokio::test]
3752 async fn test_failed_ping_to_other_endpoint_revalidates_proven_node() {
3753 let (_discv4, mut service) = create_discv4().await;
3754 let expired = Instant::now() + service.config.ping_expiration * 2;
3755
3756 for other in ["127.0.0.2:30303", "127.0.0.1:30304"] {
3758 let record = NodeRecord::new("127.0.0.1:30303".parse().unwrap(), PeerId::random());
3759 insert_proven_node(&mut service, record);
3760 let other = NodeRecord::new(other.parse().unwrap(), record.id);
3761 service.send_ping(other, PingReason::EstablishBond);
3762 service.evict_expired_requests(expired);
3763
3764 assert!(service.on_entry(record.id, |_| ()).is_some());
3765 let request = &service.pending_pings[&record.id];
3766 assert_eq!(request.node, record);
3767 assert!(matches!(request.reason, PingReason::RePing));
3768
3769 service.evict_expired_requests(expired);
3771 assert!(service.on_entry(record.id, |_| ()).is_none());
3772 }
3773 }
3774
3775 #[tokio::test]
3776 async fn test_failed_ping_to_other_endpoint_evicts_unproven_node() {
3777 let (_discv4, mut service) = create_discv4().await;
3778 let record = NodeRecord::new("127.0.0.1:30303".parse().unwrap(), PeerId::random());
3779 let other = NodeRecord::new("127.0.0.2:30303".parse().unwrap(), record.id);
3780
3781 service.send_ping(other, PingReason::EstablishBond);
3783 assert!(service.add_node(record));
3784 assert_eq!(service.pending_pings[&record.id].node, other);
3785
3786 service.evict_expired_requests(Instant::now() + service.config.ping_expiration * 2);
3787 assert!(service.on_entry(record.id, |_| ()).is_none());
3788 }
3789
3790 #[tokio::test]
3791 async fn test_find_node_timeout_counts_only_at_stored_address() {
3792 let (_discv4, mut service) = create_discv4().await;
3793 let record = NodeRecord::new("127.0.0.1:30303".parse().unwrap(), PeerId::random());
3794 insert_proven_node(&mut service, record);
3795 let other = NodeRecord::new("127.0.0.2:30303".parse().unwrap(), record.id);
3796 let ctx = LookupContext::new(kad_key(PeerId::random()), [], None);
3797 let expired = Instant::now() + service.config.neighbours_expiration * 2;
3798
3799 for (destination, failures) in [(other, 0), (record, 1)] {
3800 service.find_node(&destination, ctx.clone());
3801 service.evict_failed_find_nodes(expired);
3802 assert_eq!(
3803 service.on_entry(record.id, |entry| entry.find_node_failures),
3804 Some(failures)
3805 );
3806 }
3807 }
3808
3809 #[tokio::test]
3810 async fn test_lookup_reset_on_first_bootnode_pong() {
3811 let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3812 let config = Discv4Config::builder().add_boot_node(record).build();
3813 let (_discv4, mut service) = create_discv4_with_config(config).await;
3814
3815 assert!(service.pending_lookup_reset);
3817
3818 insert_proven_node(&mut service, record);
3820 let echo_hash = insert_initial_ping(&mut service, record);
3821
3822 service.on_pong(make_pong(&service, echo_hash), record.udp_addr(), record.id);
3824
3825 assert!(!service.pending_lookup_reset, "flag should be consumed");
3827 }
3828
3829 #[tokio::test]
3830 async fn test_lookup_reset_fires_only_once() {
3831 let records: Vec<_> = (0..2)
3832 .map(|_| NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random()))
3833 .collect();
3834 let config = Discv4Config::builder().add_boot_nodes(records.clone()).build();
3835 let (_discv4, mut service) = create_discv4_with_config(config).await;
3836
3837 for &r in &records {
3839 insert_proven_node(&mut service, r);
3840 }
3841 let hashes: Vec<_> =
3842 records.iter().map(|r| insert_initial_ping(&mut service, *r)).collect();
3843
3844 service.on_pong(make_pong(&service, hashes[0]), records[0].udp_addr(), records[0].id);
3846 assert!(!service.pending_lookup_reset);
3847
3848 service.on_pong(make_pong(&service, hashes[1]), records[1].udp_addr(), records[1].id);
3850 assert!(!service.pending_lookup_reset);
3851 }
3852
3853 #[tokio::test]
3854 async fn test_lookup_reset_not_triggered_by_non_bootnode() {
3855 let bootnode = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3856 let config = Discv4Config::builder().add_boot_node(bootnode).build();
3857 let (_discv4, mut service) = create_discv4_with_config(config).await;
3858
3859 assert!(service.pending_lookup_reset);
3860
3861 let stranger = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3863 insert_proven_node(&mut service, stranger);
3864 let echo_hash = insert_initial_ping(&mut service, stranger);
3865 service.on_pong(make_pong(&service, echo_hash), stranger.udp_addr(), stranger.id);
3866
3867 assert!(service.pending_lookup_reset, "flag should not be consumed by non-bootnode");
3868 }
3869
3870 #[tokio::test]
3871 async fn test_lookup_reset_disabled_when_lookup_disabled() {
3872 let config = Discv4Config::builder().enable_lookup(false).build();
3873 let (_discv4, service) = create_discv4_with_config(config).await;
3874
3875 assert!(!service.pending_lookup_reset);
3877 }
3878
3879 #[tokio::test]
3880 async fn test_enr_response_preserves_pending_request() {
3881 let fork_id = ForkId { hash: ForkHash([1, 2, 3, 4]), next: 0 };
3882 let mut config = Discv4Config::default();
3883 config.add_eip868_pair("eth", EnrForkIdEntry::from(fork_id));
3884 let (_remote, mut remote) = create_discv4_with_config(config).await;
3885 let (_discv4, mut service) = create_discv4().await;
3886 let record = remote.local_node_record;
3887 let id = record.id;
3888 let addr = record.udp_addr();
3889 insert_proven_node(&mut service, record);
3890 service.update_on_reping(record, remote.enr_seq());
3891 let old_enr = remote.local_eip_868_enr.clone();
3892
3893 let (_, old_hash) =
3895 Message::EnrRequest(EnrRequest { expire: service.enr_request_expiration() - 1 })
3896 .encode(&service.secret_key);
3897 service.pending_enr_requests.get_mut(&id).unwrap().echo_hash = old_hash;
3898
3899 remote.local_eip_868_enr.set_tcp4(30304, &remote.secret_key).unwrap();
3901 service.on_ping(
3902 Ping {
3903 from: record.into(),
3904 to: service.local_node_record.into(),
3905 expire: service.ping_expiration(),
3906 enr_sq: remote.enr_seq(),
3907 },
3908 addr,
3909 id,
3910 B256::random(),
3911 );
3912 let request_hash = service.pending_enr_requests[&id].echo_hash;
3913 assert_ne!(request_hash, old_hash);
3914
3915 service.on_enr_response(EnrResponse { request_hash: old_hash, enr: old_enr }, addr, id);
3916 assert_eq!(service.pending_enr_requests[&id].echo_hash, request_hash);
3917 assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(None));
3918
3919 service.on_enr_response(
3921 EnrResponse { request_hash, enr: service.local_eip_868_enr.clone() },
3922 addr,
3923 id,
3924 );
3925 assert_eq!(service.pending_enr_requests[&id].echo_hash, request_hash);
3926 assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(None));
3927
3928 let mut updates = service.update_stream().into_inner();
3929 service.on_enr_response(
3930 EnrResponse { request_hash, enr: remote.local_eip_868_enr.clone() },
3931 addr,
3932 id,
3933 );
3934 assert!(!service.pending_enr_requests.contains_key(&id));
3935 assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(Some(fork_id)));
3936 assert_matches::assert_matches!(
3937 updates.try_recv().unwrap(),
3938 DiscoveryUpdate::EnrForkId(found_record, found_fork_id) => {
3939 assert_eq!(found_record, record);
3940 assert_eq!(found_fork_id, fork_id);
3941 }
3942 );
3943 }
3944
3945 async fn pending_find_node_service(node_addr: SocketAddr) -> (Discv4Service, NodeRecord) {
3947 let (_discv4, mut service) = create_discv4().await;
3948
3949 let id = PeerId::random();
3950 let record = NodeRecord::new(node_addr, id);
3951 let _ = service.kbuckets.insert_or_update(
3952 &kad_key(id),
3953 NodeEntry::new_proven(record),
3954 NodeStatus {
3955 direction: ConnectionDirection::Incoming,
3956 state: ConnectionState::Connected,
3957 },
3958 );
3959
3960 service.lookup(PeerId::random());
3961 assert_eq!(service.pending_find_nodes.len(), 1);
3962 (service, record)
3963 }
3964
3965 #[tokio::test]
3966 async fn test_neighbours_from_unexpected_ip_ignored() {
3967 reth_tracing::init_test_tracing();
3968
3969 let node_addr = SocketAddr::from(([10, 0, 0, 1], 30303));
3970 let (mut service, record) = pending_find_node_service(node_addr).await;
3971 let ctx = service.pending_find_nodes[&record.id].lookup_context.clone();
3972 let lookup_nodes = ctx.inner.closest_nodes.borrow().len();
3973
3974 let expire = service.find_node_expiration() + 1000;
3975 let discovered =
3976 NodeRecord::new(SocketAddr::from(([10, 0, 0, 9], 30303)), PeerId::random());
3977 let msg = Neighbours { nodes: vec![discovered], expire };
3978
3979 for spoofed in [
3982 SocketAddr::from(([10, 0, 0, 2], 30303)),
3983 SocketAddr::from(([10, 0, 0, 2], 1)),
3984 SocketAddr::from((Ipv4Addr::new(10, 0, 0, 2).to_ipv6_mapped(), 30303)),
3985 SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], 30303)),
3986 ] {
3987 service.on_neighbours(msg.clone(), spoofed, record.id);
3988
3989 let request = &service.pending_find_nodes[&record.id];
3990 assert!(!request.answered);
3991 assert_eq!(request.response_count, 0);
3992 assert_eq!(ctx.inner.closest_nodes.borrow().len(), lookup_nodes);
3993 assert!(service.pending_pings.is_empty());
3994 }
3995
3996 service.on_neighbours(msg, SocketAddr::from(([10, 0, 0, 1], 40404)), record.id);
3998 let request = &service.pending_find_nodes[&record.id];
3999 assert!(request.answered);
4000 assert_eq!(request.response_count, 1);
4001 assert_eq!(ctx.inner.closest_nodes.borrow().len(), lookup_nodes + 1);
4002 }
4003
4004 #[tokio::test]
4005 async fn test_neighbours_ipv4_mapped_ip_accepted() {
4006 reth_tracing::init_test_tracing();
4007
4008 let node_addr = SocketAddr::from(([10, 0, 0, 1], 30303));
4010 let (mut service, record) = pending_find_node_service(node_addr).await;
4011
4012 let expire = service.find_node_expiration() + 1000;
4013 let msg = Neighbours { nodes: Vec::new(), expire };
4014 let mapped = SocketAddr::from((Ipv4Addr::new(10, 0, 0, 1).to_ipv6_mapped(), 30303));
4015 service.on_neighbours(msg, mapped, record.id);
4016
4017 let request = &service.pending_find_nodes[&record.id];
4018 assert!(request.answered);
4019 }
4020
4021 #[tokio::test]
4022 async fn test_neighbours_wrong_ip_replay_does_not_poison_cache() {
4023 let (_discv4, mut service) = create_discv4().await;
4024 let remote_key = SecretKey::new(&mut rand::rng());
4025 let node_addr = SocketAddr::from(([10, 0, 0, 1], 30303));
4026 let record = NodeRecord::from_secret_key(node_addr, &remote_key);
4027 insert_proven_node(&mut service, record);
4028 service.lookup(PeerId::random());
4029 assert_eq!(service.pending_find_nodes.len(), 1);
4030
4031 let ctx = service.pending_find_nodes[&record.id].lookup_context.clone();
4032 let lookup_nodes = ctx.inner.closest_nodes.borrow().len();
4033 let discovered =
4034 NodeRecord::new(SocketAddr::from(([10, 0, 0, 9], 30303)), PeerId::random());
4035 let msg = Message::Neighbours(Neighbours {
4036 nodes: vec![discovered],
4037 expire: service.find_node_expiration(),
4038 });
4039 let (packet, _) = msg.encode(&remote_key);
4040 let (tx, mut rx) = mpsc::channel(16);
4041 let mut handler = IngressHandler::new(tx, *service.local_peer_id());
4042
4043 handler.handle_packet(&packet, SocketAddr::from(([10, 0, 0, 2], 30303))).await;
4045 let IngressEvent::Packet(src, Packet { msg: Message::Neighbours(msg), node_id, .. }) =
4046 rx.try_recv().unwrap()
4047 else {
4048 panic!("expected a decoded Neighbours packet");
4049 };
4050 service.on_neighbours(msg, src, node_id);
4051 let request = &service.pending_find_nodes[&record.id];
4052 assert!(!request.answered);
4053 assert_eq!(request.response_count, 0);
4054 assert_eq!(ctx.inner.closest_nodes.borrow().len(), lookup_nodes);
4055 assert!(service.pending_pings.is_empty());
4056
4057 handler.handle_packet(&packet, node_addr).await;
4059 let IngressEvent::Packet(src, Packet { msg: Message::Neighbours(msg), node_id, .. }) =
4060 rx.try_recv().expect("genuine reply was suppressed by the wrong-IP replay")
4061 else {
4062 panic!("expected a decoded Neighbours packet");
4063 };
4064 service.on_neighbours(msg, src, node_id);
4065 let request = &service.pending_find_nodes[&record.id];
4066 assert!(request.answered);
4067 assert_eq!(request.response_count, 1);
4068 assert_eq!(ctx.inner.closest_nodes.borrow().len(), lookup_nodes + 1);
4069
4070 let mapped = SocketAddr::from((Ipv4Addr::new(10, 0, 0, 1).to_ipv6_mapped(), 40404));
4072 handler.handle_packet(&packet, mapped).await;
4073 assert_matches::assert_matches!(rx.try_recv(), Err(mpsc::error::TryRecvError::Empty));
4074 }
4075}