Skip to main content

reth_discv4/
lib.rs

1//! Discovery v4 implementation: <https://github.com/ethereum/devp2p/blob/master/discv4.md>
2//!
3//! Discv4 employs a kademlia-like routing table to store and manage discovered peers and topics.
4//! The protocol allows for external IP discovery in NAT environments through regular PING/PONG's
5//! with discovered nodes. Nodes return the external IP address that they have received and a simple
6//! majority is chosen as our external IP address. If an external IP address is updated, this is
7//! produced as an event to notify the swarm (if one is used for this behaviour).
8//!
9//! This implementation consists of a [`Discv4`] and [`Discv4Service`] pair. The service manages the
10//! state and drives the UDP socket. The (optional) [`Discv4`] serves as the frontend to interact
11//! with the service via a channel. Whenever the underlying table changes service produces a
12//! [`DiscoveryUpdate`] that listeners will receive.
13//!
14//! ## Feature Flags
15//!
16//! - `serde` (default): Enable serde support
17//! - `test-utils`: Export utilities for testing
18
19#![doc(
20    html_logo_url = "https://raw.githubusercontent.com/paradigmxyz/reth/main/assets/reth-docs.png",
21    html_favicon_url = "https://avatars0.githubusercontent.com/u/97369466?s=256",
22    issue_tracker_base_url = "https://github.com/paradigmxyz/reth/issues/"
23)]
24#![cfg_attr(not(test), warn(unused_crate_dependencies))]
25#![cfg_attr(docsrs, feature(doc_cfg))]
26
27use crate::{
28    error::{DecodePacketError, Discv4Error},
29    proto::{FindNode, Message, Neighbours, Packet, Ping, Pong},
30};
31use alloy_primitives::{bytes::Bytes, hex, B256};
32use discv5::{
33    kbucket,
34    kbucket::{
35        BucketInsertResult, Distance, Entry as BucketEntry, InsertResult, KBucketsTable,
36        NodeStatus, MAX_NODES_PER_BUCKET,
37    },
38    ConnectionDirection, ConnectionState,
39};
40use enr::Enr;
41use itertools::Itertools;
42use parking_lot::Mutex;
43use proto::{EnrRequest, EnrResponse};
44use reth_ethereum_forks::ForkId;
45use reth_network_peers::{pk2id, PeerId};
46use secp256k1::SecretKey;
47use std::{
48    cell::RefCell,
49    collections::{btree_map, hash_map::Entry, BTreeMap, HashMap, VecDeque},
50    fmt, io,
51    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
52    pin::Pin,
53    rc::Rc,
54    sync::Arc,
55    task::{ready, Context, Poll},
56    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
57};
58use tokio::{
59    net::UdpSocket,
60    sync::{mpsc, mpsc::error::TrySendError, oneshot, oneshot::Sender as OneshotSender},
61    task::{JoinHandle, JoinSet},
62    time::Interval,
63};
64use tokio_stream::{wrappers::ReceiverStream, Stream, StreamExt};
65use tracing::{debug, trace};
66
67pub mod error;
68pub mod proto;
69
70mod config;
71pub use config::{Discv4Config, Discv4ConfigBuilder};
72
73mod node;
74use node::{kad_key, NodeKey};
75
76mod table;
77
78// reexport NodeRecord primitive
79pub 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;
86/// reexport to get public ip.
87pub use reth_net_nat::{external_ip, NatResolver};
88
89/// The default address for discv4 via UDP
90///
91/// Note: the default TCP address is the same.
92pub const DEFAULT_DISCOVERY_ADDR: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED);
93
94/// The default port for discv4 via UDP
95///
96/// Note: the default TCP port is the same.
97pub const DEFAULT_DISCOVERY_PORT: u16 = 30303;
98
99/// The default address for discv4 via UDP: "0.0.0.0:30303"
100///
101/// Note: The default TCP address is the same.
102pub const DEFAULT_DISCOVERY_ADDRESS: SocketAddr =
103    SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, DEFAULT_DISCOVERY_PORT));
104
105/// The maximum size of any packet is 1280 bytes.
106const MAX_PACKET_SIZE: usize = 1280;
107
108/// Length of the UDP datagram packet-header: Hash(32b) + Signature(65b) + Packet Type(1b)
109const MIN_PACKET_SIZE: usize = 32 + 65 + 1;
110
111/// Concurrency factor for `FindNode` requests to pick `ALPHA` closest nodes, <https://github.com/ethereum/devp2p/blob/master/discv4.md#recursive-lookup>
112const ALPHA: usize = 3;
113
114/// Maximum number of nodes to ping at concurrently.
115///
116/// This corresponds to 2 full `Neighbours` responses with 16 _new_ nodes. This will apply some
117/// backpressure in recursive lookups.
118const MAX_NODES_PING: usize = 2 * MAX_NODES_PER_BUCKET;
119
120/// Maximum number of pings to keep queued.
121///
122/// If we are currently sending too many pings, any new pings will be queued. To prevent unbounded
123/// growth of the queue, the queue has a maximum capacity, after which any additional pings will be
124/// discarded.
125///
126/// This corresponds to 2 full `Neighbours` responses with 16 new nodes.
127const MAX_QUEUED_PINGS: usize = 2 * MAX_NODES_PER_BUCKET;
128
129/// The size of the datagram is limited [`MAX_PACKET_SIZE`], 16 nodes, as the discv4 specifies don't
130/// fit in one datagram. The safe number of nodes that always fit in a datagram is 12, with worst
131/// case all of them being IPv6 nodes. This is calculated by `(MAX_PACKET_SIZE - (header + expire +
132/// rlp overhead) / size(rlp(Node_IPv6))`
133/// Even in the best case where all nodes are IPv4, only 14 nodes fit into one packet.
134const SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS: usize = (MAX_PACKET_SIZE - 109) / 91;
135
136/// The timeout used to identify expired nodes, 24h
137///
138/// Mirrors geth's `bondExpiration` of 24h
139const ENDPOINT_PROOF_EXPIRATION: Duration = Duration::from_secs(24 * 60 * 60);
140
141/// Duration used to expire nodes from the routing table 1hr
142const EXPIRE_DURATION: Duration = Duration::from_secs(60 * 60);
143
144// Restricts how many udp messages can be processed in a single [Discv4Service::poll] call.
145//
146// This will act as a manual yield point when draining the socket messages where the most CPU
147// expensive part is handling outgoing messages: encoding and hashing the packet
148const 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/// The Discv4 frontend.
159///
160/// This is a cloneable type that communicates with the [`Discv4Service`] by sending commands over a
161/// shared channel.
162///
163/// See also [`Discv4::spawn`]
164#[derive(Debug, Clone)]
165pub struct Discv4 {
166    /// The address of the udp socket
167    local_addr: SocketAddr,
168    /// channel to send commands over to the service
169    to_service: mpsc::UnboundedSender<Discv4Command>,
170    /// Tracks the local node record.
171    ///
172    /// This includes the currently tracked external IP address of the node.
173    node_record: Arc<Mutex<NodeRecord>>,
174}
175
176impl Discv4 {
177    /// Same as [`Self::bind`] but also spawns the service onto a new task.
178    ///
179    /// See also: [`Discv4Service::spawn()`]
180    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    /// Returns a new instance with the given channel directly
194    ///
195    /// NOTE: this is only intended for test setups.
196    #[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    /// Binds a new `UdpSocket` and creates the service
212    ///
213    /// ```
214    /// use reth_discv4::{Discv4, Discv4Config};
215    /// use reth_network_peers::{pk2id, NodeRecord, PeerId};
216    /// use secp256k1::SECP256K1;
217    /// use std::{net::SocketAddr, str::FromStr};
218    /// # async fn t() -> std:: io::Result<()> {
219    ///
220    /// // generate a (random) keypair
221    /// let (secret_key, pk) = SECP256K1.generate_keypair(&mut rand_08::thread_rng());
222    /// let id = pk2id(&pk);
223    ///
224    /// let socket = SocketAddr::from_str("0.0.0.0:0").unwrap();
225    /// let local_enr =
226    ///     NodeRecord { address: socket.ip(), tcp_port: socket.port(), udp_port: socket.port(), id };
227    /// let config = Discv4Config::default();
228    ///
229    /// let (discv4, mut service) = Discv4::bind(socket, local_enr, secret_key, config).await.unwrap();
230    ///
231    /// // get an update strea
232    /// let updates = service.update_stream();
233    ///
234    /// let _handle = service.spawn();
235    ///
236    /// // lookup the local node in the DHT
237    /// let _discovered = discv4.lookup_self().await.unwrap();
238    ///
239    /// # Ok(())
240    /// # }
241    /// ```
242    pub async fn bind(
243        local_address: SocketAddr,
244        local_node_record: NodeRecord,
245        secret_key: SecretKey,
246        config: Discv4Config,
247    ) -> io::Result<(Self, Discv4Service)> {
248        let socket = Arc::new(UdpSocket::bind(local_address).await?);
249        trace!(target: "discv4", local_addr=?socket.local_addr(), "opened UDP socket");
250        let (tx, rx) = mpsc::channel(config.udp_ingress_message_buffer);
251
252        Self::bind_with_socket(socket, Some(tx), rx, local_node_record, secret_key, config)
253    }
254
255    /// Creates a new `Discv4` instance using a pre-bound shared socket. No receive loop is
256    /// spawned; instead returns an [`IngressHandler`] that should be used to forward raw packets
257    /// received by the socket owner (e.g. discv5 unrecognized frames).
258    pub fn bind_shared(
259        socket: Arc<UdpSocket>,
260        local_node_record: NodeRecord,
261        secret_key: SecretKey,
262        config: Discv4Config,
263    ) -> io::Result<(Self, Discv4Service, IngressHandler)> {
264        let (tx, rx) = mpsc::channel(config.udp_ingress_message_buffer);
265        let local_id = local_node_record.id;
266        let (discv4, service) =
267            Self::bind_with_socket(socket, None, rx, local_node_record, secret_key, config)?;
268
269        let handler = IngressHandler::new(tx, local_id);
270
271        Ok((discv4, service, handler))
272    }
273
274    fn bind_with_socket(
275        socket: Arc<UdpSocket>,
276        ingress_tx: Option<IngressSender>,
277        ingress_rx: IngressReceiver,
278        mut local_node_record: NodeRecord,
279        secret_key: SecretKey,
280        config: Discv4Config,
281    ) -> io::Result<(Self, Discv4Service)> {
282        let local_addr = socket.local_addr()?;
283        local_node_record.udp_port = local_addr.port();
284
285        let mut service = Discv4Service::new(
286            socket,
287            ingress_tx,
288            ingress_rx,
289            local_addr,
290            local_node_record,
291            secret_key,
292            config,
293        );
294
295        // resolve the external address immediately
296        service.resolve_external_ip();
297
298        let discv4 = service.handle();
299        Ok((discv4, service))
300    }
301
302    /// Returns the address of the UDP socket.
303    pub const fn local_addr(&self) -> SocketAddr {
304        self.local_addr
305    }
306
307    /// Returns the [`NodeRecord`] of the local node.
308    ///
309    /// This includes the currently tracked external IP address of the node.
310    pub fn node_record(&self) -> NodeRecord {
311        *self.node_record.lock()
312    }
313
314    /// Returns the currently tracked external IP of the node.
315    pub fn external_ip(&self) -> IpAddr {
316        self.node_record.lock().address
317    }
318
319    /// Sets the [Interval] used for periodically looking up targets over the network
320    pub fn set_lookup_interval(&self, duration: Duration) {
321        self.send_to_service(Discv4Command::SetLookupInterval(duration))
322    }
323
324    /// Starts a `FindNode` recursive lookup that locates the closest nodes to the given node id. See also: <https://github.com/ethereum/devp2p/blob/master/discv4.md#recursive-lookup>
325    ///
326    /// The lookup initiator starts by picking α closest nodes to the target it knows of. The
327    /// initiator then sends concurrent `FindNode` packets to those nodes. α is a system-wide
328    /// concurrency parameter, such as 3. In the recursive step, the initiator resends `FindNode` to
329    /// nodes it has learned about from previous queries. Of the k nodes the initiator has heard of
330    /// closest to the target, it picks α that it has not yet queried and resends `FindNode` to
331    /// them. Nodes that fail to respond quickly are removed from consideration until and unless
332    /// they do respond.
333    //
334    // If a round of FindNode queries fails to return a node any closer than the closest already
335    // seen, the initiator resends the find node to all of the k closest nodes it has not already
336    // queried. The lookup terminates when the initiator has queried and gotten responses from the k
337    // closest nodes it has seen.
338    pub async fn lookup_self(&self) -> Result<Vec<NodeRecord>, Discv4Error> {
339        self.lookup_node(None).await
340    }
341
342    /// Looks up the given node id.
343    ///
344    /// Returning the closest nodes to the given node id.
345    pub async fn lookup(&self, node_id: PeerId) -> Result<Vec<NodeRecord>, Discv4Error> {
346        self.lookup_node(Some(node_id)).await
347    }
348
349    /// Performs a random lookup for node records.
350    pub async fn lookup_random(&self) -> Result<Vec<NodeRecord>, Discv4Error> {
351        let target = PeerId::random();
352        self.lookup_node(Some(target)).await
353    }
354
355    /// Sends a message to the service to lookup the closest nodes
356    pub fn send_lookup(&self, node_id: PeerId) {
357        let cmd = Discv4Command::Lookup { node_id: Some(node_id), tx: None };
358        self.send_to_service(cmd);
359    }
360
361    async fn lookup_node(&self, node_id: Option<PeerId>) -> Result<Vec<NodeRecord>, Discv4Error> {
362        let (tx, rx) = oneshot::channel();
363        let cmd = Discv4Command::Lookup { node_id, tx: Some(tx) };
364        self.to_service.send(cmd)?;
365        Ok(rx.await?)
366    }
367
368    /// Triggers a new self lookup without expecting a response
369    pub fn send_lookup_self(&self) {
370        let cmd = Discv4Command::Lookup { node_id: None, tx: None };
371        self.send_to_service(cmd);
372    }
373
374    /// Removes the peer from the table, if it exists.
375    pub fn remove_peer(&self, node_id: PeerId) {
376        let cmd = Discv4Command::Remove(node_id);
377        self.send_to_service(cmd);
378    }
379
380    /// Adds the node to the table, if it is not already present.
381    pub fn add_node(&self, node_record: NodeRecord) {
382        let cmd = Discv4Command::Add(node_record);
383        self.send_to_service(cmd);
384    }
385
386    /// Adds the node as a bootnode.
387    ///
388    /// This registers the node in the configured bootstrap set and inserts it into the routing
389    /// table, pinging it to establish the endpoint proof, same as the nodes provided at startup.
390    pub fn add_boot_node(&self, node_record: NodeRecord) {
391        let cmd = Discv4Command::AddBootNode(node_record);
392        self.send_to_service(cmd);
393    }
394
395    /// Adds the peer and id to the ban list.
396    ///
397    /// This will prevent any future inclusion in the table
398    pub fn ban(&self, node_id: PeerId, ip: IpAddr) {
399        let cmd = Discv4Command::Ban(node_id, ip);
400        self.send_to_service(cmd);
401    }
402
403    /// Adds the ip to the ban list.
404    ///
405    /// This will prevent any future inclusion in the table
406    pub fn ban_ip(&self, ip: IpAddr) {
407        let cmd = Discv4Command::BanIp(ip);
408        self.send_to_service(cmd);
409    }
410
411    /// Adds the peer to the ban list.
412    ///
413    /// This will prevent any future inclusion in the table
414    pub fn ban_node(&self, node_id: PeerId) {
415        let cmd = Discv4Command::BanPeer(node_id);
416        self.send_to_service(cmd);
417    }
418
419    /// Sets the tcp port
420    ///
421    /// This will update our [`NodeRecord`]'s tcp port.
422    pub fn set_tcp_port(&self, port: u16) {
423        let cmd = Discv4Command::SetTcpPort(port);
424        self.send_to_service(cmd);
425    }
426
427    /// Sets the pair in the EIP-868 [`Enr`] of the node.
428    ///
429    /// If the key already exists, this will update it.
430    ///
431    /// CAUTION: The value **must** be rlp encoded
432    pub fn set_eip868_rlp_pair(&self, key: Vec<u8>, rlp: Bytes) {
433        let cmd = Discv4Command::SetEIP868RLPPair { key, rlp };
434        self.send_to_service(cmd);
435    }
436
437    /// Sets the pair in the EIP-868 [`Enr`] of the node.
438    ///
439    /// If the key already exists, this will update it.
440    pub fn set_eip868_rlp(&self, key: Vec<u8>, value: impl alloy_rlp::Encodable) {
441        self.set_eip868_rlp_pair(key, Bytes::from(alloy_rlp::encode(&value)))
442    }
443
444    #[inline]
445    fn send_to_service(&self, cmd: Discv4Command) {
446        let _ = self.to_service.send(cmd).map_err(|err| {
447            debug!(
448                target: "discv4",
449                %err,
450                "channel capacity reached, dropping command",
451            )
452        });
453    }
454
455    /// Returns the receiver half of new listener channel that streams [`DiscoveryUpdate`]s.
456    pub async fn update_stream(&self) -> Result<ReceiverStream<DiscoveryUpdate>, Discv4Error> {
457        let (tx, rx) = oneshot::channel();
458        let cmd = Discv4Command::Updates(tx);
459        self.to_service.send(cmd)?;
460        Ok(rx.await?)
461    }
462
463    /// Terminates the spawned [`Discv4Service`].
464    pub fn terminate(&self) {
465        self.send_to_service(Discv4Command::Terminated);
466    }
467}
468
469/// Manages discv4 peer discovery over UDP.
470///
471/// This is a [Stream] to handles incoming and outgoing discv4 messages and emits updates via:
472/// [`Discv4Service::update_stream`].
473///
474/// This type maintains the discv Kademlia routing table and is responsible for performing lookups.
475///
476/// ## Lookups
477///
478/// See also [Recursive Lookups](https://github.com/ethereum/devp2p/blob/master/discv4.md#recursive-lookup).
479/// Lookups are either triggered periodically or performaned on demand: [`Discv4::lookup`]
480/// Newly discovered nodes are emitted as [`DiscoveryUpdate::Added`] event to all subscribers:
481/// [`Discv4Service::update_stream`].
482#[must_use = "Stream does nothing unless polled"]
483pub struct Discv4Service {
484    /// Local address of the UDP socket.
485    local_address: SocketAddr,
486    /// The local ENR for EIP-868 <https://eips.ethereum.org/EIPS/eip-868>
487    local_eip_868_enr: Enr<SecretKey>,
488    /// Local ENR of the server.
489    local_node_record: NodeRecord,
490    /// Keeps track of the node record of the local node.
491    shared_node_record: Arc<Mutex<NodeRecord>>,
492    /// The secret key used to sign payloads
493    secret_key: SecretKey,
494    /// The UDP socket for sending and receiving messages.
495    _socket: Arc<UdpSocket>,
496    /// The spawned UDP tasks.
497    ///
498    /// Note: If dropped, the spawned send+receive tasks are aborted.
499    _tasks: JoinSet<()>,
500    /// The routing table.
501    kbuckets: KBucketsTable<NodeKey, NodeEntry>,
502    /// Receiver for incoming messages
503    ///
504    /// Receives incoming messages from the UDP task.
505    ingress: IngressReceiver,
506    /// Sender for sending outgoing messages
507    ///
508    /// Sends outgoing messages to the UDP task.
509    egress: EgressSender,
510    /// Buffered pending pings to apply backpressure.
511    ///
512    /// Lookups behave like bursts of requests: Endpoint proof followed by `FindNode` request. [Recursive lookups](https://github.com/ethereum/devp2p/blob/master/discv4.md#recursive-lookup) can trigger multiple followup Pings+FindNode requests.
513    /// A cap on concurrent `Ping` prevents escalation where: A large number of new nodes
514    /// discovered via `FindNode` in a recursive lookup triggers a large number of `Ping`s, and
515    /// followup `FindNode` requests.... Buffering them effectively prevents high `Ping` peaks.
516    queued_pings: VecDeque<(NodeRecord, PingReason)>,
517    /// Currently active pings to specific nodes.
518    pending_pings: HashMap<PeerId, PingRequest>,
519    /// Currently active endpoint proof verification lookups to specific nodes.
520    ///
521    /// Entries here means we've proven the peer's endpoint but haven't completed our end of the
522    /// endpoint proof
523    pending_lookup: HashMap<PeerId, (Instant, LookupContext)>,
524    /// Currently active `FindNode` requests
525    pending_find_nodes: HashMap<PeerId, FindNodeRequest>,
526    /// Currently active ENR requests
527    pending_enr_requests: HashMap<PeerId, EnrRequestState>,
528    /// Copy of the sender half of the commands channel for [Discv4]
529    to_service: mpsc::UnboundedSender<Discv4Command>,
530    /// Receiver half of the commands channel for [Discv4]
531    commands_rx: mpsc::UnboundedReceiver<Discv4Command>,
532    /// All subscribers for table updates
533    update_listeners: Vec<mpsc::Sender<DiscoveryUpdate>>,
534    /// The interval when to trigger random lookups
535    lookup_interval: Interval,
536    /// Used to rotate targets to lookup
537    lookup_rotator: LookupTargetRotator,
538    /// Whether we still need to reset the lookup interval on the first bootnode pong.
539    pending_lookup_reset: bool,
540    /// Interval when to recheck active requests
541    evict_expired_requests_interval: Interval,
542    /// Interval when to resend pings.
543    ping_interval: Interval,
544    /// The interval at which to attempt resolving external IP again.
545    resolve_external_ip_interval: Option<ResolveNatInterval>,
546    /// How this services is configured
547    config: Discv4Config,
548    /// Buffered events populated during poll.
549    queued_events: VecDeque<Discv4Event>,
550    /// Keeps track of nodes from which we have received a `Pong` message.
551    received_pongs: PongTable,
552    /// Interval used to expire additionally tracked nodes
553    expire_interval: Interval,
554    /// Cached signed `FindNode` packet to avoid redundant ECDSA signing during lookups.
555    cached_find_node: Option<CachedFindNode>,
556}
557
558impl Discv4Service {
559    /// Create a new instance for a bound [`UdpSocket`].
560    ///
561    /// If `ingress_tx` is `Some`, the receive loop is spawned to read from the socket. If `None`,
562    /// the caller feeds packets into `ingress_rx` externally (shared socket mode).
563    pub(crate) fn new(
564        socket: Arc<UdpSocket>,
565        ingress_tx: Option<IngressSender>,
566        ingress_rx: IngressReceiver,
567        local_address: SocketAddr,
568        local_node_record: NodeRecord,
569        secret_key: SecretKey,
570        config: Discv4Config,
571    ) -> Self {
572        let (egress_tx, egress_rx) = mpsc::channel(config.udp_egress_message_buffer);
573        let mut tasks = JoinSet::<()>::new();
574
575        if let Some(ingress_tx) = ingress_tx {
576            let udp = Arc::clone(&socket);
577            tasks.spawn(receive_loop(udp, ingress_tx, local_node_record.id));
578        }
579
580        let udp = Arc::clone(&socket);
581        tasks.spawn(send_loop(udp, egress_rx));
582
583        let kbuckets = KBucketsTable::new(
584            NodeKey::from(&local_node_record).into(),
585            Duration::from_secs(60),
586            MAX_NODES_PER_BUCKET,
587            None,
588            None,
589        );
590
591        let self_lookup_interval = tokio::time::interval(config.lookup_interval);
592
593        // Wait `ping_interval` and then start pinging every `ping_interval` because we want to wait
594        // for
595        let ping_interval = tokio::time::interval_at(
596            tokio::time::Instant::now() + config.ping_interval,
597            config.ping_interval,
598        );
599
600        let evict_expired_requests_interval = tokio::time::interval_at(
601            tokio::time::Instant::now() + config.request_timeout,
602            config.request_timeout,
603        );
604
605        let lookup_rotator = if config.enable_dht_random_walk {
606            LookupTargetRotator::default()
607        } else {
608            LookupTargetRotator::local_only()
609        };
610
611        // for EIP-868 construct an ENR
612        let local_eip_868_enr = {
613            let mut builder = Enr::builder();
614            builder.ip(local_node_record.address);
615            if local_node_record.address.is_ipv4() {
616                builder.udp4(local_node_record.udp_port);
617                builder.tcp4(local_node_record.tcp_port);
618            } else {
619                builder.udp6(local_node_record.udp_port);
620                builder.tcp6(local_node_record.tcp_port);
621            }
622
623            for (key, val) in &config.additional_eip868_rlp_pairs {
624                builder.add_value_rlp(key, val.clone());
625            }
626            builder.build(&secret_key).expect("v4 is set")
627        };
628
629        let (to_service, commands_rx) = mpsc::unbounded_channel();
630
631        let shared_node_record = Arc::new(Mutex::new(local_node_record));
632
633        Self {
634            local_address,
635            local_eip_868_enr,
636            local_node_record,
637            shared_node_record,
638            _socket: socket,
639            kbuckets,
640            secret_key,
641            _tasks: tasks,
642            ingress: ingress_rx,
643            egress: egress_tx,
644            queued_pings: VecDeque::with_capacity(MAX_QUEUED_PINGS),
645            pending_pings: Default::default(),
646            pending_lookup: Default::default(),
647            pending_find_nodes: Default::default(),
648            pending_enr_requests: Default::default(),
649            commands_rx,
650            to_service,
651            update_listeners: Vec::with_capacity(1),
652            lookup_interval: self_lookup_interval,
653            ping_interval,
654            evict_expired_requests_interval,
655            lookup_rotator,
656            pending_lookup_reset: config.enable_lookup,
657            resolve_external_ip_interval: config.resolve_external_ip_interval(),
658            config,
659            queued_events: Default::default(),
660            received_pongs: Default::default(),
661            expire_interval: tokio::time::interval(EXPIRE_DURATION),
662            cached_find_node: None,
663        }
664    }
665
666    /// Returns the frontend handle that can communicate with the service via commands.
667    pub fn handle(&self) -> Discv4 {
668        Discv4 {
669            local_addr: self.local_address,
670            to_service: self.to_service.clone(),
671            node_record: self.shared_node_record.clone(),
672        }
673    }
674
675    /// Returns the current enr sequence of the local record.
676    fn enr_seq(&self) -> Option<u64> {
677        self.config.enable_eip868.then(|| self.local_eip_868_enr.seq())
678    }
679
680    /// Sets the [Interval] used for periodically looking up targets over the network
681    pub fn set_lookup_interval(&mut self, duration: Duration) {
682        self.lookup_interval = tokio::time::interval(duration);
683    }
684
685    /// Sets the external Ip to the configured external IP if [`NatResolver::ExternalIp`] or
686    /// [`NatResolver::ExternalAddr`]. In the case of [`NatResolver::ExternalAddr`], it will return
687    /// the first IP address found for the domain associated with the discv4 UDP port.
688    fn resolve_external_ip(&mut self) {
689        if let Some(r) = &self.resolve_external_ip_interval &&
690            let Some(external_ip) =
691                r.resolver().clone().as_external_ip(self.local_node_record.udp_port)
692        {
693            self.set_external_ip_addr(external_ip);
694        }
695    }
696
697    /// Sets the given ip address as the node's external IP in the node record announced in
698    /// discovery
699    pub fn set_external_ip_addr(&mut self, external_ip: IpAddr) {
700        if self.local_node_record.address != external_ip {
701            debug!(target: "discv4", ?external_ip, "Updating external ip");
702            self.local_node_record.address = external_ip;
703            let _ = self.local_eip_868_enr.set_ip(external_ip, &self.secret_key);
704            let mut lock = self.shared_node_record.lock();
705            *lock = self.local_node_record;
706            debug!(target: "discv4", enr=?self.local_eip_868_enr, "Updated local ENR");
707        }
708    }
709
710    /// Returns the [`PeerId`] that identifies this node
711    pub const fn local_peer_id(&self) -> &PeerId {
712        &self.local_node_record.id
713    }
714
715    /// Returns the address of the UDP socket
716    pub const fn local_addr(&self) -> SocketAddr {
717        self.local_address
718    }
719
720    /// Returns the ENR of this service.
721    ///
722    /// Note: this will include the external address if resolved.
723    pub const fn local_enr(&self) -> NodeRecord {
724        self.local_node_record
725    }
726
727    /// Returns mutable reference to ENR for testing.
728    #[cfg(test)]
729    pub const fn local_enr_mut(&mut self) -> &mut NodeRecord {
730        &mut self.local_node_record
731    }
732
733    /// Returns true if the given `PeerId` is currently in the bucket
734    pub fn contains_node(&self, id: PeerId) -> bool {
735        let key = kad_key(id);
736        self.kbuckets.get_index(&key).is_some()
737    }
738
739    /// Bootstraps the local node to join the DHT.
740    ///
741    /// Bootstrapping is a multi-step operation that starts with a lookup of the local node's
742    /// own ID in the DHT. This introduces the local node to the other nodes
743    /// in the DHT and populates its routing table with the closest proven neighbours.
744    ///
745    /// This inserts the configured bootnodes into the routing table and pings them. Once the
746    /// endpoint proof succeeds (pong received), a [`DiscoveryUpdate::Added`] event is emitted,
747    /// same as with [`Self::add_node`].
748    ///
749    /// **Note:** This is a noop if there are no bootnodes.
750    pub fn bootstrap(&mut self) {
751        for record in self.config.bootstrap_nodes.clone() {
752            debug!(target: "discv4", ?record, "pinging boot node");
753            let key = kad_key(record.id);
754            let entry = NodeEntry::new(record);
755
756            // insert the boot node in the table
757            match self.kbuckets.insert_or_update(
758                &key,
759                entry,
760                NodeStatus {
761                    state: ConnectionState::Disconnected,
762                    direction: ConnectionDirection::Outgoing,
763                },
764            ) {
765                InsertResult::Failed(_) => {}
766                _ => {
767                    self.try_ping(record, PingReason::InitialInsert);
768                }
769            }
770        }
771    }
772
773    /// Adds the node to the bootstrap set and to the routing table.
774    ///
775    /// Behaves like [`Self::add_node`] but also registers the node in the configured bootstrap
776    /// set so it is used for subsequent bootstrap attempts.
777    pub fn add_boot_node(&mut self, record: NodeRecord) -> bool {
778        self.config.bootstrap_nodes.insert(record);
779        self.add_node(record)
780    }
781
782    /// Spawns this services onto a new task
783    ///
784    /// Note: requires a running tokio runtime
785    pub fn spawn(mut self) -> JoinHandle<()> {
786        tokio::task::spawn(async move {
787            self.bootstrap();
788
789            while let Some(event) = self.next().await {
790                trace!(target: "discv4", ?event, "processed");
791            }
792            trace!(target: "discv4", "service terminated");
793        })
794    }
795
796    /// Creates a new bounded channel for [`DiscoveryUpdate`]s.
797    pub fn update_stream(&mut self) -> ReceiverStream<DiscoveryUpdate> {
798        let (tx, rx) = mpsc::channel(512);
799        self.update_listeners.push(tx);
800        ReceiverStream::new(rx)
801    }
802
803    /// Looks up the local node in the DHT.
804    pub fn lookup_self(&mut self) {
805        self.lookup(self.local_node_record.id)
806    }
807
808    /// Looks up the given node in the DHT
809    ///
810    /// A `FindNode` packet requests information about nodes close to target. The target is a
811    /// 64-byte secp256k1 public key. When `FindNode` is received, the recipient should reply
812    /// with Neighbors packets containing the closest 16 nodes to target found in its local
813    /// table.
814    //
815    // To guard against traffic amplification attacks, Neighbors replies should only be sent if the
816    // sender of FindNode has been verified by the endpoint proof procedure.
817    pub fn lookup(&mut self, target: PeerId) {
818        self.lookup_with(target, None)
819    }
820
821    /// Starts the recursive lookup process for the given target, <https://github.com/ethereum/devp2p/blob/master/discv4.md#recursive-lookup>.
822    ///
823    /// At first the `ALPHA` (==3, defined concurrency factor) nodes that are closest to the target
824    /// in the underlying DHT are selected to seed the lookup via `FindNode` requests. In the
825    /// recursive step, the initiator resends `FindNode` to nodes it has learned about from previous
826    /// queries.
827    ///
828    /// This takes an optional Sender through which all successfully discovered nodes are sent once
829    /// the request has finished.
830    fn lookup_with(&mut self, target: PeerId, tx: Option<NodeRecordSender>) {
831        trace!(target: "discv4", ?target, "Starting lookup");
832        let target_key = kad_key(target);
833
834        // Start a lookup context with the 16 (MAX_NODES_PER_BUCKET) closest nodes to which we have
835        // a valid endpoint proof
836        let ctx = LookupContext::new(
837            target_key.clone(),
838            self.kbuckets
839                .closest_values(&target_key)
840                .filter(|node| {
841                    node.value.has_endpoint_proof &&
842                        !self.pending_find_nodes.contains_key(&node.key.preimage().0)
843                })
844                .take(MAX_NODES_PER_BUCKET)
845                .map(|n| (target_key.distance(&n.key), n.value.record)),
846            tx,
847        );
848
849        // From those 16, pick the 3 closest to start the concurrent lookup.
850        let closest = ctx.closest(ALPHA);
851
852        if closest.is_empty() && self.pending_find_nodes.is_empty() {
853            // no closest nodes, and no lookup in progress: table is empty.
854            // This could happen if all records were deleted from the table due to missed pongs
855            // (e.g. connectivity problems over a long period of time, or issues during initial
856            // bootstrapping) so we attempt to bootstrap again
857            self.bootstrap();
858            return
859        }
860
861        trace!(target: "discv4", ?target, num = closest.len(), "Start lookup closest nodes");
862
863        for node in closest {
864            // here we still want to check against previous request failures and if necessary
865            // re-establish a new endpoint proof because it can be the case that the other node lost
866            // our entry and no longer has an endpoint proof on their end
867            self.find_node_checked(&node, ctx.clone());
868        }
869    }
870
871    /// Sends a new `FindNode` packet to the node with `target` as the lookup target.
872    ///
873    /// CAUTION: This expects there's a valid Endpoint proof to the given `node`.
874    fn find_node(&mut self, node: &NodeRecord, ctx: LookupContext) {
875        trace!(target: "discv4", ?node, lookup=?ctx.target(), "Sending FindNode");
876        ctx.mark_queried(node.id);
877        let (payload, hash) = self.find_node_packet(ctx.target());
878        let to = node.udp_addr();
879        trace!(target: "discv4", ?to, ?hash, "sending FindNode packet");
880        let _ = self.egress.try_send((payload, to)).map_err(|err| {
881            debug!(target: "discv4", %err, "dropped outgoing packet");
882        });
883        self.pending_find_nodes.insert(node.id, FindNodeRequest::new(ctx));
884    }
885
886    /// Sends a new `FindNode` packet to the node with `target` as the lookup target but checks
887    /// whether we should send a new ping first to renew the endpoint proof by checking the
888    /// previously failed findNode requests. It could be that the node is no longer reachable or
889    /// lost our entry.
890    fn find_node_checked(&mut self, node: &NodeRecord, ctx: LookupContext) {
891        let max_failures = self.config.max_find_node_failures;
892        let needs_ping = self
893            .on_entry(node.id, |entry| entry.exceeds_find_node_failures(max_failures))
894            .unwrap_or(true);
895        if needs_ping {
896            self.try_ping(*node, PingReason::Lookup(*node, ctx))
897        } else {
898            self.find_node(node, ctx)
899        }
900    }
901
902    /// Notifies all listeners.
903    ///
904    /// Removes all listeners that are closed.
905    fn notify(&mut self, update: DiscoveryUpdate) {
906        self.update_listeners.retain_mut(|listener| match listener.try_send(update.clone()) {
907            Ok(()) => true,
908            Err(err) => match err {
909                TrySendError::Full(_) => true,
910                TrySendError::Closed(_) => false,
911            },
912        });
913    }
914
915    /// Adds the ip to the ban list indefinitely
916    pub fn ban_ip(&mut self, ip: IpAddr) {
917        self.config.ban_list.ban_ip(ip);
918    }
919
920    /// Adds the peer to the ban list indefinitely.
921    pub fn ban_node(&mut self, node_id: PeerId) {
922        self.remove_node(node_id);
923        self.config.ban_list.ban_peer(node_id);
924    }
925
926    /// Adds the ip to the ban list until the given timestamp.
927    pub fn ban_ip_until(&mut self, ip: IpAddr, until: Instant) {
928        self.config.ban_list.ban_ip_until(ip, until);
929    }
930
931    /// Adds the peer to the ban list and bans it until the given timestamp
932    pub fn ban_node_until(&mut self, node_id: PeerId, until: Instant) {
933        self.remove_node(node_id);
934        self.config.ban_list.ban_peer_until(node_id, until);
935    }
936
937    /// Removes a `node_id` from the routing table.
938    ///
939    /// This allows applications, for whatever reason, to remove nodes from the local routing
940    /// table. Returns `true` if the node was in the table and `false` otherwise.
941    pub fn remove_node(&mut self, node_id: PeerId) -> bool {
942        let key = kad_key(node_id);
943        self.remove_key(node_id, key)
944    }
945
946    /// Removes a `node_id` from the routing table but only if there are enough other nodes in the
947    /// bucket (bucket must be at least half full)
948    ///
949    /// Returns `true` if the node was removed
950    pub fn soft_remove_node(&mut self, node_id: PeerId) -> bool {
951        let key = kad_key(node_id);
952        let Some(bucket) = self.kbuckets.get_bucket(&key) else { return false };
953        if bucket.num_entries() < MAX_NODES_PER_BUCKET / 2 {
954            // skip half empty bucket
955            return false
956        }
957        self.remove_key(node_id, key)
958    }
959
960    fn remove_key(&mut self, node_id: PeerId, key: discv5::Key<NodeKey>) -> bool {
961        let removed = self.kbuckets.remove(&key);
962        if removed {
963            trace!(target: "discv4", ?node_id, "removed node");
964            self.notify(DiscoveryUpdate::Removed(node_id));
965        }
966        removed
967    }
968
969    /// Gets the number of entries that are considered connected.
970    pub fn num_connected(&self) -> usize {
971        self.kbuckets.buckets_iter().fold(0, |count, bucket| count + bucket.num_connected())
972    }
973
974    /// Check if the peer has an active bond.
975    fn has_bond(&self, remote_id: PeerId, remote_ip: IpAddr) -> bool {
976        if let Some(timestamp) = self.received_pongs.last_pong(remote_id, remote_ip) &&
977            timestamp.elapsed() < self.config.bond_expiration
978        {
979            return true
980        }
981        false
982    }
983
984    /// Applies a closure on the pending or present [`NodeEntry`].
985    fn on_entry<F, R>(&mut self, peer_id: PeerId, f: F) -> Option<R>
986    where
987        F: FnOnce(&NodeEntry) -> R,
988    {
989        let key = kad_key(peer_id);
990        match self.kbuckets.entry(&key) {
991            BucketEntry::Present(entry, _) => Some(f(entry.value())),
992            BucketEntry::Pending(entry, _) => Some(f(entry.value())),
993            _ => None,
994        }
995    }
996
997    /// Update the entry on RE-ping.
998    ///
999    /// Invoked when we received the Pong to our [`PingReason::RePing`] ping.
1000    ///
1001    /// On re-ping we check for a changed `enr_seq` if eip868 is enabled and when it changed we sent
1002    /// a followup request to retrieve the updated ENR
1003    fn update_on_reping(&mut self, record: NodeRecord, mut last_enr_seq: Option<u64>) {
1004        if record.id == self.local_node_record.id {
1005            return
1006        }
1007
1008        // If EIP868 extension is disabled then we want to ignore this
1009        if !self.config.enable_eip868 {
1010            last_enr_seq = None;
1011        }
1012
1013        let key = kad_key(record.id);
1014        let old_enr = match self.kbuckets.entry(&key) {
1015            kbucket::Entry::Present(mut entry, _) => {
1016                entry.value_mut().update_with_enr(last_enr_seq)
1017            }
1018            kbucket::Entry::Pending(mut entry, _) => {
1019                entry.value_mut().update_with_enr(last_enr_seq)
1020            }
1021            _ => return,
1022        };
1023
1024        // Check if ENR was updated
1025        match (last_enr_seq, old_enr) {
1026            (Some(new), Some(old)) if new > old => {
1027                self.send_enr_request(record);
1028            }
1029            (Some(_), None) => {
1030                // got an ENR
1031                self.send_enr_request(record);
1032            }
1033            _ => {}
1034        };
1035    }
1036
1037    /// Callback invoked when we receive a pong from the peer.
1038    fn update_on_pong(&mut self, record: NodeRecord, mut last_enr_seq: Option<u64>) {
1039        if record.id == *self.local_peer_id() {
1040            return
1041        }
1042
1043        // If EIP868 extension is disabled then we want to ignore this
1044        if !self.config.enable_eip868 {
1045            last_enr_seq = None;
1046        }
1047
1048        // if the peer included a enr seq in the pong then we can try to request the ENR of that
1049        // node
1050        let has_enr_seq = last_enr_seq.is_some();
1051
1052        let key = kad_key(record.id);
1053        match self.kbuckets.entry(&key) {
1054            kbucket::Entry::Present(mut entry, old_status) => {
1055                // endpoint is now proven
1056                entry.value_mut().establish_proof();
1057                entry.value_mut().update_with_enr(last_enr_seq);
1058
1059                if !old_status.is_connected() {
1060                    let _ = entry.update(ConnectionState::Connected, Some(old_status.direction));
1061                    trace!(target: "discv4", ?record, "added after successful endpoint proof");
1062                    self.notify(DiscoveryUpdate::Added(record));
1063
1064                    if has_enr_seq {
1065                        // request the ENR of the node
1066                        self.send_enr_request(record);
1067                    }
1068                }
1069            }
1070            kbucket::Entry::Pending(mut entry, mut status) => {
1071                // endpoint is now proven
1072                entry.value_mut().establish_proof();
1073                entry.value_mut().update_with_enr(last_enr_seq);
1074
1075                if !status.is_connected() {
1076                    status.state = ConnectionState::Connected;
1077                    let _ = entry.update(status);
1078                    trace!(target: "discv4", ?record, "added after successful endpoint proof");
1079                    self.notify(DiscoveryUpdate::Added(record));
1080
1081                    if has_enr_seq {
1082                        // request the ENR of the node
1083                        self.send_enr_request(record);
1084                    }
1085                }
1086            }
1087            _ => {}
1088        };
1089    }
1090
1091    /// Adds all nodes
1092    ///
1093    /// See [`Self::add_node`]
1094    pub fn add_all_nodes(&mut self, records: impl IntoIterator<Item = NodeRecord>) {
1095        for record in records {
1096            self.add_node(record);
1097        }
1098    }
1099
1100    /// If the node's not in the table yet, this will add it to the table and start the endpoint
1101    /// proof by sending a ping to the node.
1102    ///
1103    /// Returns `true` if the record was added successfully, and `false` if the node is either
1104    /// already in the table or the record's bucket is full.
1105    pub fn add_node(&mut self, record: NodeRecord) -> bool {
1106        let key = kad_key(record.id);
1107        match self.kbuckets.entry(&key) {
1108            kbucket::Entry::Absent(entry) => {
1109                let node = NodeEntry::new(record);
1110                match entry.insert(
1111                    node,
1112                    NodeStatus {
1113                        direction: ConnectionDirection::Outgoing,
1114                        state: ConnectionState::Disconnected,
1115                    },
1116                ) {
1117                    BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1118                        trace!(target: "discv4", ?record, "inserted new record");
1119                    }
1120                    _ => return false,
1121                }
1122            }
1123            _ => return false,
1124        }
1125
1126        // send the initial ping to the _new_ node
1127        self.try_ping(record, PingReason::InitialInsert);
1128        true
1129    }
1130
1131    /// Encodes the packet, sends it and returns the hash.
1132    pub(crate) fn send_packet(&self, msg: Message, to: SocketAddr) -> B256 {
1133        let (payload, hash) = msg.encode(&self.secret_key);
1134        trace!(target: "discv4", r#type=?msg.msg_type(), ?to, ?hash, "sending packet");
1135        let _ = self.egress.try_send((payload, to)).map_err(|err| {
1136            debug!(
1137                target: "discv4",
1138                %err,
1139                "dropped outgoing packet",
1140            );
1141        });
1142        hash
1143    }
1144
1145    /// Returns a signed `FindNode` packet for `target`, reusing a cached payload when possible.
1146    fn find_node_packet(&mut self, target: PeerId) -> (Bytes, B256) {
1147        let expire = self.find_node_expiration();
1148        let cache_ttl = self.config.request_timeout / 4;
1149        CachedFindNode::get_or_sign(
1150            &mut self.cached_find_node,
1151            target,
1152            cache_ttl,
1153            &self.secret_key,
1154            expire,
1155        )
1156    }
1157
1158    /// Message handler for an incoming `Ping`
1159    fn on_ping(&mut self, ping: Ping, remote_addr: SocketAddr, remote_id: PeerId, hash: B256) {
1160        if self.is_expired(ping.expire) {
1161            // ping's expiration timestamp is in the past
1162            return
1163        }
1164
1165        // create the record
1166        let record = NodeRecord {
1167            address: remote_addr.ip(),
1168            udp_port: remote_addr.port(),
1169            tcp_port: ping.from.tcp_port,
1170            id: remote_id,
1171        }
1172        .into_ipv4_mapped();
1173
1174        let key = kad_key(record.id);
1175
1176        // See also <https://github.com/ethereum/devp2p/blob/master/discv4.md#ping-packet-0x01>:
1177        // > If no communication with the sender of this ping has occurred within the last 12h, a
1178        // > ping should be sent in addition to pong in order to receive an endpoint proof.
1179        //
1180        // Note: we only mark if the node is absent because the `last 12h` condition is handled by
1181        // the ping interval
1182        let mut is_new_insert = false;
1183        let mut needs_bond = false;
1184        let mut is_proven = false;
1185
1186        let old_enr = match self.kbuckets.entry(&key) {
1187            kbucket::Entry::Present(mut entry, _) => {
1188                if entry.value().is_expired() {
1189                    // If no communication with the sender has occurred within the last 12h, a ping
1190                    // should be sent in addition to pong in order to receive an endpoint proof.
1191                    needs_bond = true;
1192                } else {
1193                    is_proven = entry.value().has_endpoint_proof;
1194                }
1195                entry.value_mut().update_with_enr(ping.enr_sq)
1196            }
1197            kbucket::Entry::Pending(mut entry, _) => {
1198                if entry.value().is_expired() {
1199                    // If no communication with the sender has occurred within the last 12h, a ping
1200                    // should be sent in addition to pong in order to receive an endpoint proof.
1201                    needs_bond = true;
1202                } else {
1203                    is_proven = entry.value().has_endpoint_proof;
1204                }
1205                entry.value_mut().update_with_enr(ping.enr_sq)
1206            }
1207            kbucket::Entry::Absent(entry) => {
1208                let mut node = NodeEntry::new(record);
1209                node.last_enr_seq = ping.enr_sq;
1210
1211                match entry.insert(
1212                    node,
1213                    NodeStatus {
1214                        direction: ConnectionDirection::Incoming,
1215                        // mark as disconnected until endpoint proof established on pong
1216                        state: ConnectionState::Disconnected,
1217                    },
1218                ) {
1219                    BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1220                        // mark as new insert if insert was successful
1221                        is_new_insert = true;
1222                    }
1223                    BucketInsertResult::Full => {
1224                        // we received a ping but the corresponding bucket for the peer is already
1225                        // full, we can't add any additional peers to that bucket, but we still want
1226                        // to emit an event that we discovered the node
1227                        trace!(target: "discv4", ?record, "discovered new record but bucket is full");
1228                        self.notify(DiscoveryUpdate::DiscoveredAtCapacity(record));
1229                        needs_bond = true;
1230                    }
1231                    BucketInsertResult::TooManyIncoming | BucketInsertResult::NodeExists => {
1232                        needs_bond = true;
1233                        // insert unsuccessful but we still want to send the pong
1234                    }
1235                    BucketInsertResult::FailedFilter => return,
1236                }
1237
1238                None
1239            }
1240            kbucket::Entry::SelfEntry => return,
1241        };
1242
1243        // send the pong first, but the PONG and optionally PING don't need to be send in a
1244        // particular order
1245        let pong = Message::Pong(Pong {
1246            // we use the actual address of the peer
1247            to: record.into(),
1248            echo: hash,
1249            expire: ping.expire,
1250            enr_sq: self.enr_seq(),
1251        });
1252        self.send_packet(pong, remote_addr);
1253
1254        // if node was absent also send a ping to establish the endpoint proof from our end
1255        if is_new_insert {
1256            self.try_ping(record, PingReason::InitialInsert);
1257        } else if needs_bond {
1258            self.try_ping(record, PingReason::EstablishBond);
1259        } else {
1260            // if node has been proven, this means we've received a pong and verified its endpoint
1261            // proof. We've also sent a pong above to verify our endpoint proof, so we can now
1262            // send our find_nodes request if PingReason::Lookup
1263            if is_proven && let Some((_, ctx)) = self.pending_lookup.remove(&record.id) {
1264                if self.pending_find_nodes.contains_key(&record.id) {
1265                    // there's already another pending request, unmark it so the next round can
1266                    // try to send it
1267                    ctx.unmark_queried(record.id);
1268                } else {
1269                    // we just received a ping from that peer so we can send a find node request
1270                    // directly
1271                    self.find_node(&record, ctx);
1272                }
1273            }
1274
1275            // Request ENR if included in the ping
1276            match (ping.enr_sq, old_enr) {
1277                (Some(new), Some(old)) if new > old => {
1278                    self.send_enr_request(record);
1279                }
1280                (Some(_), None) => {
1281                    self.send_enr_request(record);
1282                }
1283                _ => {}
1284            };
1285        }
1286    }
1287
1288    // Guarding function for [`Self::send_ping`] that applies pre-checks
1289    fn try_ping(&mut self, node: NodeRecord, reason: PingReason) {
1290        if node.id == *self.local_peer_id() {
1291            // don't ping ourselves
1292            return
1293        }
1294
1295        if self.pending_pings.contains_key(&node.id) ||
1296            self.pending_find_nodes.contains_key(&node.id)
1297        {
1298            return
1299        }
1300
1301        if self.queued_pings.iter().any(|(n, _)| n.id == node.id) {
1302            return
1303        }
1304
1305        if self.pending_pings.len() < MAX_NODES_PING {
1306            self.send_ping(node, reason);
1307        } else if self.queued_pings.len() < MAX_QUEUED_PINGS {
1308            self.queued_pings.push_back((node, reason));
1309        }
1310    }
1311
1312    /// Sends a ping message to the node's UDP address.
1313    ///
1314    /// Returns the echo hash of the ping message.
1315    pub(crate) fn send_ping(&mut self, node: NodeRecord, reason: PingReason) -> B256 {
1316        let remote_addr = node.udp_addr();
1317        let id = node.id;
1318        let ping = Ping {
1319            from: self.local_node_record.into(),
1320            to: node.into(),
1321            expire: self.ping_expiration(),
1322            enr_sq: self.enr_seq(),
1323        };
1324        trace!(target: "discv4", ?ping, "sending ping");
1325        let echo_hash = self.send_packet(Message::Ping(ping), remote_addr);
1326
1327        self.pending_pings
1328            .insert(id, PingRequest { sent_at: Instant::now(), node, echo_hash, reason });
1329        echo_hash
1330    }
1331
1332    /// Sends an enr request message to the node's UDP address.
1333    ///
1334    /// Returns the echo hash of the ping message.
1335    pub(crate) fn send_enr_request(&mut self, node: NodeRecord) {
1336        if !self.config.enable_eip868 {
1337            return
1338        }
1339        let remote_addr = node.udp_addr();
1340        let enr_request = EnrRequest { expire: self.enr_request_expiration() };
1341
1342        trace!(target: "discv4", ?enr_request, "sending enr request");
1343        let echo_hash = self.send_packet(Message::EnrRequest(enr_request), remote_addr);
1344
1345        self.pending_enr_requests
1346            .insert(node.id, EnrRequestState { sent_at: Instant::now(), echo_hash });
1347    }
1348
1349    /// Message handler for an incoming `Pong`.
1350    fn on_pong(&mut self, pong: Pong, remote_addr: SocketAddr, remote_id: PeerId) {
1351        if self.is_expired(pong.expire) {
1352            return
1353        }
1354
1355        let PingRequest { node, reason, .. } = match self.pending_pings.entry(remote_id) {
1356            Entry::Occupied(entry) => {
1357                {
1358                    let request = entry.get();
1359                    if request.echo_hash != pong.echo {
1360                        trace!(target: "discv4", from=?remote_addr, expected=?request.echo_hash, echo_hash=?pong.echo,"Got unexpected Pong");
1361                        return
1362                    }
1363                }
1364                entry.remove()
1365            }
1366            Entry::Vacant(_) => return,
1367        };
1368
1369        // keep track of the pong
1370        self.received_pongs.on_pong(remote_id, remote_addr.ip());
1371
1372        match reason {
1373            PingReason::InitialInsert => {
1374                self.update_on_pong(node, pong.enr_sq);
1375                // Reset the lookup interval so the next poll_tick fires immediately,
1376                // rather than waiting the full ~20s for the first lookup.
1377                if self.pending_lookup_reset && self.config.bootstrap_nodes.contains(&node) {
1378                    self.pending_lookup_reset = false;
1379                    self.lookup_interval.reset();
1380                }
1381            }
1382            PingReason::EstablishBond => {
1383                // no initial lookup needed here since the node was already in the table.
1384                self.update_on_pong(node, pong.enr_sq);
1385            }
1386            PingReason::RePing => {
1387                self.update_on_reping(node, pong.enr_sq);
1388            }
1389            PingReason::Lookup(node, ctx) => {
1390                self.update_on_pong(node, pong.enr_sq);
1391                // insert node and assoc. lookup_context into the pending_lookup table to complete
1392                // our side of the endpoint proof verification.
1393                // Start the lookup timer here - and evict accordingly. Note that this is a separate
1394                // timer than the ping_request timer.
1395                self.pending_lookup.insert(node.id, (Instant::now(), ctx));
1396            }
1397        }
1398    }
1399
1400    /// Handler for an incoming `FindNode` message
1401    fn on_find_node(&mut self, msg: FindNode, remote_addr: SocketAddr, node_id: PeerId) {
1402        if self.is_expired(msg.expire) {
1403            // expiration timestamp is in the past
1404            return
1405        }
1406        if node_id == *self.local_peer_id() {
1407            // ignore find node requests to ourselves
1408            return
1409        }
1410
1411        if self.has_bond(node_id, remote_addr.ip()) {
1412            self.respond_closest(msg.id, remote_addr)
1413        }
1414    }
1415
1416    /// Handler for incoming `EnrResponse` message
1417    fn on_enr_response(&mut self, msg: EnrResponse, remote_addr: SocketAddr, id: PeerId) {
1418        trace!(target: "discv4", ?remote_addr, ?msg, "received ENR response");
1419        let Entry::Occupied(request) = self.pending_enr_requests.entry(id) else { return };
1420
1421        // A delayed response to an earlier request must not consume its replacement.
1422        // Also ensure the ENR's public key matches the expected node id.
1423        if request.get().echo_hash != msg.request_hash || id != pk2id(&msg.enr.public_key()) {
1424            return
1425        }
1426        request.remove();
1427
1428        let key = kad_key(id);
1429        let fork_id = msg.eth_fork_id();
1430        let (record, old_fork_id) = match self.kbuckets.entry(&key) {
1431            kbucket::Entry::Present(mut entry, _) => {
1432                let id = entry.value_mut().update_with_fork_id(fork_id);
1433                (entry.value().record, id)
1434            }
1435            kbucket::Entry::Pending(mut entry, _) => {
1436                let id = entry.value_mut().update_with_fork_id(fork_id);
1437                (entry.value().record, id)
1438            }
1439            _ => return,
1440        };
1441        match (fork_id, old_fork_id) {
1442            (Some(new), Some(old)) if new != old => {
1443                self.notify(DiscoveryUpdate::EnrForkId(record, new))
1444            }
1445            (Some(new), None) => self.notify(DiscoveryUpdate::EnrForkId(record, new)),
1446            _ => {}
1447        }
1448    }
1449
1450    /// Handler for incoming `EnrRequest` message
1451    fn on_enr_request(
1452        &self,
1453        msg: EnrRequest,
1454        remote_addr: SocketAddr,
1455        id: PeerId,
1456        request_hash: B256,
1457    ) {
1458        if !self.config.enable_eip868 || self.is_expired(msg.expire) {
1459            return
1460        }
1461
1462        if self.has_bond(id, remote_addr.ip()) {
1463            self.send_packet(
1464                Message::EnrResponse(EnrResponse {
1465                    request_hash,
1466                    enr: self.local_eip_868_enr.clone(),
1467                }),
1468                remote_addr,
1469            );
1470        }
1471    }
1472
1473    /// Handler for incoming `Neighbours` messages that are handled if they're responses to
1474    /// `FindNode` requests.
1475    fn on_neighbours(&mut self, msg: Neighbours, remote_addr: SocketAddr, node_id: PeerId) {
1476        if self.is_expired(msg.expire) {
1477            // response is expired
1478            return
1479        }
1480        // check if this request was expected
1481        let ctx = match self.pending_find_nodes.entry(node_id) {
1482            Entry::Occupied(mut entry) => {
1483                {
1484                    let request = entry.get_mut();
1485                    // Mark the request as answered
1486                    request.answered = true;
1487                    let total = request.response_count + msg.nodes.len();
1488
1489                    // Neighbours response is exactly 1 bucket (16 entries).
1490                    if total <= MAX_NODES_PER_BUCKET {
1491                        request.response_count = total;
1492                    } else {
1493                        trace!(target: "discv4", total, from=?remote_addr, "Received neighbors packet entries exceeds max nodes per bucket");
1494                        return
1495                    }
1496                };
1497
1498                if entry.get().response_count == MAX_NODES_PER_BUCKET {
1499                    // node responding with a full bucket of records
1500                    let ctx = entry.remove().lookup_context;
1501                    ctx.mark_responded(node_id);
1502                    ctx
1503                } else {
1504                    entry.get().lookup_context.clone()
1505                }
1506            }
1507            Entry::Vacant(_) => {
1508                // received neighbours response without requesting it
1509                trace!(target: "discv4", from=?remote_addr, "Received unsolicited Neighbours");
1510                return
1511            }
1512        };
1513
1514        // log the peers we discovered
1515        trace!(target: "discv4",
1516            target=format!("{:#?}", node_id),
1517            peers_count=msg.nodes.len(),
1518            peers=format!("[{:#}]", msg.nodes.iter()
1519                .map(|node_rec| node_rec.id
1520            ).format(", ")),
1521            "Received peers from Neighbours packet"
1522        );
1523
1524        // This is the recursive lookup step where we initiate new FindNode requests for new nodes
1525        // that were discovered.
1526        for node in msg.nodes.into_iter().map(NodeRecord::into_ipv4_mapped) {
1527            // prevent banned peers from being added to the context
1528            if self.config.ban_list.is_banned(&node.id, &node.address) {
1529                trace!(target: "discv4", peer_id=?node.id, ip=?node.address, "ignoring banned record");
1530                continue
1531            }
1532
1533            ctx.add_node(node);
1534        }
1535
1536        // get the next closest nodes, not yet queried nodes and start over.
1537        let closest =
1538            ctx.filter_closest(ALPHA, |node| !self.pending_find_nodes.contains_key(&node.id));
1539
1540        for closest in closest {
1541            let key = kad_key(closest.id);
1542            match self.kbuckets.entry(&key) {
1543                BucketEntry::Absent(entry) => {
1544                    // the node's endpoint is not proven yet, so we need to ping it first, on
1545                    // success, we will add the node to the pending_lookup table, and wait to send
1546                    // back a Pong before initiating a FindNode request.
1547                    // In order to prevent that this node is selected again on subsequent responses,
1548                    // while the ping is still active, we always mark it as queried.
1549                    ctx.mark_queried(closest.id);
1550                    let node = NodeEntry::new(closest);
1551                    match entry.insert(
1552                        node,
1553                        NodeStatus {
1554                            direction: ConnectionDirection::Outgoing,
1555                            state: ConnectionState::Disconnected,
1556                        },
1557                    ) {
1558                        BucketInsertResult::Inserted | BucketInsertResult::Pending { .. } => {
1559                            // only ping if the node was added to the table
1560                            self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1561                        }
1562                        BucketInsertResult::Full => {
1563                            // new node but the node's bucket is already full
1564                            self.notify(DiscoveryUpdate::DiscoveredAtCapacity(closest))
1565                        }
1566                        _ => {}
1567                    }
1568                }
1569                BucketEntry::SelfEntry => {
1570                    // we received our own node entry
1571                }
1572                BucketEntry::Present(entry, _) => {
1573                    if entry.value().has_endpoint_proof {
1574                        if entry
1575                            .value()
1576                            .exceeds_find_node_failures(self.config.max_find_node_failures)
1577                        {
1578                            self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1579                        } else {
1580                            self.find_node(&closest, ctx.clone());
1581                        }
1582                    }
1583                }
1584                BucketEntry::Pending(entry, _) => {
1585                    if entry.value().has_endpoint_proof {
1586                        if entry
1587                            .value()
1588                            .exceeds_find_node_failures(self.config.max_find_node_failures)
1589                        {
1590                            self.try_ping(closest, PingReason::Lookup(closest, ctx.clone()))
1591                        } else {
1592                            self.find_node(&closest, ctx.clone());
1593                        }
1594                    }
1595                }
1596            }
1597        }
1598    }
1599
1600    /// Sends a Neighbours packet for `target` to the given addr
1601    fn respond_closest(&mut self, target: PeerId, to: SocketAddr) {
1602        let key = kad_key(target);
1603        let expire = self.send_neighbours_expiration();
1604
1605        let enforce_eip868 = self.config.enable_eip868 && self.config.enforce_eip868_neighbours;
1606
1607        // get the MAX_NODES_PER_BUCKET closest nodes to the target, optionally filtering out
1608        // entries that have no EIP-868 fork ID
1609        let closest_nodes = self
1610            .kbuckets
1611            .closest_values(&key)
1612            .filter(|entry| !enforce_eip868 || entry.value.fork_id.is_some())
1613            .take(MAX_NODES_PER_BUCKET)
1614            .collect::<Vec<_>>();
1615
1616        if closest_nodes.is_empty() {
1617            // always respond so the requester does not treat this as a timeout
1618            let msg = Message::Neighbours(Neighbours { nodes: Vec::new(), expire });
1619            self.send_packet(msg, to);
1620            return;
1621        }
1622
1623        for nodes in closest_nodes.chunks(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS) {
1624            let nodes = nodes.iter().map(|node| node.value.record).collect::<Vec<NodeRecord>>();
1625            trace!(target: "discv4", len = nodes.len(), to=?to,"Sent neighbours packet");
1626            let msg = Message::Neighbours(Neighbours { nodes, expire });
1627            self.send_packet(msg, to);
1628        }
1629    }
1630
1631    fn evict_expired_requests(&mut self, now: Instant) {
1632        let mut failed_enr_requests = Vec::new();
1633        self.pending_enr_requests.retain(|node_id, enr_request| {
1634            if now.duration_since(enr_request.sent_at) < self.config.enr_expiration {
1635                return true
1636            }
1637            failed_enr_requests.push(*node_id);
1638            false
1639        });
1640
1641        // forget the announced enr seq, so the next ping or pong requests the ENR again.
1642        for node_id in failed_enr_requests {
1643            match self.kbuckets.entry(&kad_key(node_id)) {
1644                kbucket::Entry::Present(mut entry, _) => entry.value_mut().last_enr_seq = None,
1645                kbucket::Entry::Pending(mut entry, _) => entry.value_mut().last_enr_seq = None,
1646                _ => {}
1647            }
1648        }
1649
1650        let mut failed_pings = Vec::new();
1651        self.pending_pings.retain(|node_id, ping_request| {
1652            if now.duration_since(ping_request.sent_at) > self.config.ping_expiration {
1653                failed_pings.push(*node_id);
1654                return false
1655            }
1656            true
1657        });
1658
1659        if !failed_pings.is_empty() {
1660            // remove nodes that failed to pong
1661            trace!(target: "discv4", num=%failed_pings.len(), "evicting nodes due to failed pong");
1662            for node_id in failed_pings {
1663                self.remove_node(node_id);
1664            }
1665        }
1666
1667        let mut failed_lookups = Vec::new();
1668        self.pending_lookup.retain(|node_id, (lookup_sent_at, _)| {
1669            if now.duration_since(*lookup_sent_at) > self.config.request_timeout {
1670                failed_lookups.push(*node_id);
1671                return false
1672            }
1673            true
1674        });
1675
1676        if !failed_lookups.is_empty() {
1677            // remove nodes that failed the e2e lookup process, so we can restart it
1678            trace!(target: "discv4", num=%failed_lookups.len(), "evicting nodes due to failed lookup");
1679            for node_id in failed_lookups {
1680                self.remove_node(node_id);
1681            }
1682        }
1683
1684        self.evict_failed_find_nodes(now);
1685    }
1686
1687    /// Handles failed responses to `FindNode`
1688    fn evict_failed_find_nodes(&mut self, now: Instant) {
1689        let mut failed_find_nodes = Vec::new();
1690        self.pending_find_nodes.retain(|node_id, find_node_request| {
1691            if now.duration_since(find_node_request.sent_at) > self.config.neighbours_expiration {
1692                if !find_node_request.answered {
1693                    // node actually responded but with fewer entries than expected, but we don't
1694                    // treat this as an hard error since it responded.
1695                    failed_find_nodes.push(*node_id);
1696                }
1697                return false
1698            }
1699            true
1700        });
1701
1702        if failed_find_nodes.is_empty() {
1703            return
1704        }
1705
1706        trace!(target: "discv4", num=%failed_find_nodes.len(), "processing failed find nodes");
1707
1708        for node_id in failed_find_nodes {
1709            let key = kad_key(node_id);
1710            let failures = match self.kbuckets.entry(&key) {
1711                kbucket::Entry::Present(mut entry, _) => {
1712                    entry.value_mut().inc_failed_request();
1713                    entry.value().find_node_failures
1714                }
1715                kbucket::Entry::Pending(mut entry, _) => {
1716                    entry.value_mut().inc_failed_request();
1717                    entry.value().find_node_failures
1718                }
1719                _ => continue,
1720            };
1721
1722            // if the node failed to respond anything useful multiple times, remove the node from
1723            // the table, but only if there are enough other nodes in the bucket (bucket must be at
1724            // least half full)
1725            if failures > self.config.max_find_node_failures {
1726                self.soft_remove_node(node_id);
1727            }
1728        }
1729    }
1730
1731    /// Re-pings all nodes which endpoint proofs are considered expired: [`NodeEntry::is_expired`]
1732    ///
1733    /// This will send a `Ping` to the nodes, if a node fails to respond with a `Pong` to renew the
1734    /// endpoint proof it will be removed from the table.
1735    fn re_ping_oldest(&mut self) {
1736        let mut nodes = self
1737            .kbuckets
1738            .iter_ref()
1739            .filter(|entry| entry.node.value.is_expired())
1740            .map(|n| n.node.value)
1741            .collect::<Vec<_>>();
1742        nodes.sort_unstable_by_key(|a| a.last_seen);
1743        let to_ping = nodes.into_iter().map(|n| n.record).take(MAX_NODES_PING).collect::<Vec<_>>();
1744        for node in to_ping {
1745            self.try_ping(node, PingReason::RePing)
1746        }
1747    }
1748
1749    /// Returns true if the expiration timestamp is in the past.
1750    fn is_expired(&self, expiration: u64) -> bool {
1751        self.ensure_not_expired(expiration).is_err()
1752    }
1753
1754    /// Validate that given timestamp is not expired.
1755    ///
1756    /// Note: this accepts the timestamp as u64 because this is used by the wire protocol, but the
1757    /// UNIX timestamp (number of non-leap seconds since January 1, 1970 0:00:00 UTC) is supposed to
1758    /// be an i64.
1759    ///
1760    /// Returns an error if:
1761    ///  - invalid UNIX timestamp (larger than `i64::MAX`)
1762    ///  - timestamp is expired (lower than current local UNIX timestamp)
1763    fn ensure_not_expired(&self, timestamp: u64) -> Result<(), ()> {
1764        // ensure the timestamp is a valid UNIX timestamp
1765        let _ = i64::try_from(timestamp).map_err(drop)?;
1766
1767        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
1768        if self.config.enforce_expiration_timestamps && timestamp < now {
1769            trace!(target: "discv4", "Expired packet");
1770            return Err(())
1771        }
1772        Ok(())
1773    }
1774
1775    /// Pops buffered ping requests and sends them.
1776    fn ping_buffered(&mut self) {
1777        while self.pending_pings.len() < MAX_NODES_PING {
1778            match self.queued_pings.pop_front() {
1779                Some((next, reason)) => self.try_ping(next, reason),
1780                None => break,
1781            }
1782        }
1783    }
1784
1785    fn ping_expiration(&self) -> u64 {
1786        (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.ping_expiration)
1787            .as_secs()
1788    }
1789
1790    fn find_node_expiration(&self) -> u64 {
1791        (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.request_timeout)
1792            .as_secs()
1793    }
1794
1795    fn enr_request_expiration(&self) -> u64 {
1796        (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.enr_expiration)
1797            .as_secs()
1798    }
1799
1800    fn send_neighbours_expiration(&self) -> u64 {
1801        (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + self.config.neighbours_expiration)
1802            .as_secs()
1803    }
1804
1805    /// Polls the socket and advances the state.
1806    ///
1807    /// To prevent traffic amplification attacks, implementations must verify that the sender of a
1808    /// query participates in the discovery protocol. The sender of a packet is considered verified
1809    /// if it has sent a valid Pong response with matching ping hash within the last 12 hours.
1810    pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Discv4Event> {
1811        loop {
1812            // drain buffered events first
1813            if let Some(event) = self.queued_events.pop_front() {
1814                return Poll::Ready(event)
1815            }
1816
1817            // trigger self lookup
1818            if self.config.enable_lookup {
1819                while self.lookup_interval.poll_tick(cx).is_ready() {
1820                    let target = self.lookup_rotator.next(&self.local_node_record.id);
1821                    self.lookup_with(target, None);
1822                }
1823            }
1824
1825            // re-ping some peers
1826            while self.ping_interval.poll_tick(cx).is_ready() {
1827                self.re_ping_oldest();
1828            }
1829
1830            if let Some(Poll::Ready(Some(ip))) =
1831                self.resolve_external_ip_interval.as_mut().map(|r| r.poll_tick(cx))
1832            {
1833                self.set_external_ip_addr(ip);
1834            }
1835
1836            // drain all incoming `Discv4` commands, this channel can never close
1837            while let Poll::Ready(Some(cmd)) = self.commands_rx.poll_recv(cx) {
1838                match cmd {
1839                    Discv4Command::Add(enr) => {
1840                        self.add_node(enr);
1841                    }
1842                    Discv4Command::AddBootNode(record) => {
1843                        self.add_boot_node(record);
1844                    }
1845                    Discv4Command::Lookup { node_id, tx } => {
1846                        let node_id = node_id.unwrap_or(self.local_node_record.id);
1847                        self.lookup_with(node_id, tx);
1848                    }
1849                    Discv4Command::SetLookupInterval(duration) => {
1850                        self.set_lookup_interval(duration);
1851                    }
1852                    Discv4Command::Updates(tx) => {
1853                        let rx = self.update_stream();
1854                        let _ = tx.send(rx);
1855                    }
1856                    Discv4Command::BanPeer(node_id) => self.ban_node(node_id),
1857                    Discv4Command::Remove(node_id) => {
1858                        self.remove_node(node_id);
1859                    }
1860                    Discv4Command::Ban(node_id, ip) => {
1861                        self.ban_node(node_id);
1862                        self.ban_ip(ip);
1863                    }
1864                    Discv4Command::BanIp(ip) => {
1865                        self.ban_ip(ip);
1866                    }
1867                    Discv4Command::SetEIP868RLPPair { key, rlp } => {
1868                        debug!(target: "discv4", key=%String::from_utf8_lossy(&key), "Update EIP-868 extension pair");
1869
1870                        let _ = self.local_eip_868_enr.insert_raw_rlp(key, rlp, &self.secret_key);
1871                    }
1872                    Discv4Command::SetTcpPort(port) => {
1873                        debug!(target: "discv4", %port, "Update tcp port");
1874                        self.local_node_record.tcp_port = port;
1875                        if self.local_node_record.address.is_ipv4() {
1876                            let _ = self.local_eip_868_enr.set_tcp4(port, &self.secret_key);
1877                        } else {
1878                            let _ = self.local_eip_868_enr.set_tcp6(port, &self.secret_key);
1879                        }
1880                    }
1881
1882                    Discv4Command::Terminated => {
1883                        // terminate the service
1884                        self.queued_events.push_back(Discv4Event::Terminated);
1885                    }
1886                }
1887            }
1888
1889            // restricts how many messages we process in a single poll before yielding back control
1890            let mut udp_message_budget = UDP_MESSAGE_POLL_LOOP_BUDGET;
1891
1892            // process all incoming datagrams
1893            while let Poll::Ready(Some(event)) = self.ingress.poll_recv(cx) {
1894                match event {
1895                    IngressEvent::RecvError(err) => {
1896                        debug!(target: "discv4", %err, "failed to read datagram");
1897                    }
1898                    IngressEvent::BadPacket(from, err, data) => {
1899                        trace!(target: "discv4", ?from, %err, packet=?hex::encode(&data), "bad packet");
1900                    }
1901                    IngressEvent::Packet(remote_addr, Packet { msg, node_id, hash }) => {
1902                        trace!(target: "discv4", r#type=?msg.msg_type(), from=?remote_addr,"received packet");
1903                        let event = match msg {
1904                            Message::Ping(ping) => {
1905                                self.on_ping(ping, remote_addr, node_id, hash);
1906                                Discv4Event::Ping
1907                            }
1908                            Message::Pong(pong) => {
1909                                self.on_pong(pong, remote_addr, node_id);
1910                                Discv4Event::Pong
1911                            }
1912                            Message::FindNode(msg) => {
1913                                self.on_find_node(msg, remote_addr, node_id);
1914                                Discv4Event::FindNode
1915                            }
1916                            Message::Neighbours(msg) => {
1917                                self.on_neighbours(msg, remote_addr, node_id);
1918                                Discv4Event::Neighbours
1919                            }
1920                            Message::EnrRequest(msg) => {
1921                                self.on_enr_request(msg, remote_addr, node_id, hash);
1922                                Discv4Event::EnrRequest
1923                            }
1924                            Message::EnrResponse(msg) => {
1925                                self.on_enr_response(msg, remote_addr, node_id);
1926                                Discv4Event::EnrResponse
1927                            }
1928                        };
1929
1930                        self.queued_events.push_back(event);
1931                    }
1932                }
1933
1934                udp_message_budget -= 1;
1935                if udp_message_budget < 0 {
1936                    trace!(target: "discv4", budget=UDP_MESSAGE_POLL_LOOP_BUDGET, "exhausted message poll budget");
1937                    if self.queued_events.is_empty() {
1938                        // we've exceeded the message budget and have no events to process
1939                        // this will make sure we're woken up again
1940                        cx.waker().wake_by_ref();
1941                    }
1942                    break
1943                }
1944            }
1945
1946            // try resending buffered pings
1947            self.ping_buffered();
1948
1949            // evict expired requests
1950            while self.evict_expired_requests_interval.poll_tick(cx).is_ready() {
1951                self.evict_expired_requests(Instant::now());
1952            }
1953
1954            // evict expired nodes
1955            while self.expire_interval.poll_tick(cx).is_ready() {
1956                self.received_pongs.evict_expired(Instant::now(), EXPIRE_DURATION);
1957            }
1958
1959            if self.queued_events.is_empty() {
1960                return Poll::Pending
1961            }
1962        }
1963    }
1964}
1965
1966/// Endless future impl
1967impl Stream for Discv4Service {
1968    type Item = Discv4Event;
1969
1970    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1971        // Poll the internal poll method
1972        match ready!(self.get_mut().poll(cx)) {
1973            // if the service is terminated, return None to terminate the stream
1974            Discv4Event::Terminated => Poll::Ready(None),
1975            // For any other event, return Poll::Ready(Some(event))
1976            ev => Poll::Ready(Some(ev)),
1977        }
1978    }
1979}
1980
1981impl fmt::Debug for Discv4Service {
1982    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1983        f.debug_struct("Discv4Service")
1984            .field("local_address", &self.local_address)
1985            .field("local_peer_id", &self.local_peer_id())
1986            .field("local_node_record", &self.local_node_record)
1987            .field("queued_pings", &self.queued_pings)
1988            .field("pending_lookup", &self.pending_lookup)
1989            .field("pending_find_nodes", &self.pending_find_nodes)
1990            .field("lookup_interval", &self.lookup_interval)
1991            .finish_non_exhaustive()
1992    }
1993}
1994
1995/// The Event type the Service stream produces.
1996///
1997/// This is mainly used for testing purposes and represents messages the service processed
1998#[derive(Debug, Eq, PartialEq)]
1999pub enum Discv4Event {
2000    /// A `Ping` message was handled.
2001    Ping,
2002    /// A `Pong` message was handled.
2003    Pong,
2004    /// A `FindNode` message was handled.
2005    FindNode,
2006    /// A `Neighbours` message was handled.
2007    Neighbours,
2008    /// A `EnrRequest` message was handled.
2009    EnrRequest,
2010    /// A `EnrResponse` message was handled.
2011    EnrResponse,
2012    /// Service is being terminated
2013    Terminated,
2014}
2015
2016/// Continuously reads new messages from the channel and writes them to the socket
2017pub(crate) async fn send_loop(udp: Arc<UdpSocket>, rx: EgressReceiver) {
2018    let mut stream = ReceiverStream::new(rx);
2019    while let Some((payload, to)) = stream.next().await {
2020        match udp.send_to(&payload, to).await {
2021            Ok(size) => {
2022                trace!(target: "discv4", ?to, ?size,"sent payload");
2023            }
2024            Err(err) => {
2025                debug!(target: "discv4", ?to, %err,"Failed to send datagram.");
2026            }
2027        }
2028    }
2029}
2030
2031/// Rate limits the number of incoming packets from individual IPs to 1 packet/second
2032const MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP: usize = 60usize;
2033
2034/// Continuously awaits new incoming messages and sends them back through the channel.
2035///
2036/// The receive loop enforces primitive rate limiting for IPs to prevent message spams from
2037/// individual IPs.
2038pub(crate) async fn receive_loop(udp: Arc<UdpSocket>, tx: IngressSender, local_id: PeerId) {
2039    let mut handler = IngressHandler::new(tx, local_id);
2040    let mut buf = [0; MAX_PACKET_SIZE];
2041    loop {
2042        let res = udp.recv_from(&mut buf).await;
2043        match res {
2044            Err(err) => {
2045                debug!(target: "discv4", %err, "Failed to read datagram.");
2046                handler.send(IngressEvent::RecvError(err)).await;
2047            }
2048            Ok((read, remote_addr)) => {
2049                handler.handle_packet(&buf[..read], remote_addr).await;
2050            }
2051        }
2052    }
2053}
2054
2055/// Handles decoding, rate-limiting, and deduplication of incoming discv4 packets.
2056///
2057/// Used by both the standalone receive loop and the shared-port mode via
2058/// [`Discv4::bind_shared`].
2059#[derive(Debug)]
2060pub struct IngressHandler {
2061    tx: IngressSender,
2062    local_id: PeerId,
2063    tick: usize,
2064    tick_interval: Duration,
2065    cache: ReceiveCache,
2066    last_tick: Instant,
2067}
2068
2069impl IngressHandler {
2070    fn new(tx: IngressSender, local_id: PeerId) -> Self {
2071        let tick = MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP / 2;
2072        Self {
2073            tx,
2074            local_id,
2075            tick,
2076            tick_interval: Duration::from_secs(tick as u64),
2077            cache: ReceiveCache::default(),
2078            last_tick: Instant::now(),
2079        }
2080    }
2081
2082    async fn send(&self, event: IngressEvent) {
2083        let _ = self.tx.send(event).await.map_err(|err| {
2084            debug!(target: "discv4", %err, "failed send incoming packet");
2085        });
2086    }
2087
2088    /// Handles an incoming raw packet: decodes, rate-limits, deduplicates, and forwards to the
2089    /// discv4 service. Used in shared-port mode to process unrecognized frames from discv5.
2090    pub async fn handle_packet(&mut self, data: &[u8], src: SocketAddr) {
2091        if self.last_tick.elapsed() >= self.tick_interval {
2092            self.cache.tick_ips(self.tick);
2093            self.last_tick = Instant::now();
2094        }
2095
2096        // rate limit incoming packets by IP
2097        if self.cache.inc_ip(src.ip()) > MAX_INCOMING_PACKETS_PER_MINUTE_BY_IP {
2098            trace!(target: "discv4", ?src, "Too many incoming packets from IP.");
2099            return
2100        }
2101
2102        // A packet starts with the hash of everything that follows it, and `Message::decode`
2103        // rejects any packet whose contents do not hash to it. A repeat of a packet we already
2104        // accepted can therefore be recognised from those 32 bytes alone, without decoding.
2105        // Decoding runs an ECDSA recovery, so checking here keeps a replayed packet from costing
2106        // a signature verification.
2107        if data.len() >= MIN_PACKET_SIZE &&
2108            self.cache.contains_packet(B256::from_slice(&data[..32]))
2109        {
2110            trace!(target: "discv4", ?src, "Received duplicate packet.");
2111            return
2112        }
2113
2114        let event = match Message::decode(data) {
2115            Ok(packet) => {
2116                if packet.node_id == self.local_id {
2117                    debug!(target: "discv4", ?src, "Received own packet.");
2118                    return
2119                }
2120
2121                // Only packets that decoded are remembered, so a peer cannot suppress a packet we
2122                // have not seen yet by guessing its hash.
2123                self.cache.insert_packet(packet.hash);
2124
2125                IngressEvent::Packet(src, packet)
2126            }
2127            Err(err) => {
2128                trace!(target: "discv4", %err, "Failed to decode packet");
2129                IngressEvent::BadPacket(src, err, data.to_vec())
2130            }
2131        };
2132
2133        self.send(event).await;
2134    }
2135}
2136
2137/// A cache for received packets and their source address.
2138///
2139/// This is used to discard duplicated packets and rate limit messages from the same source.
2140#[derive(Debug)]
2141struct ReceiveCache {
2142    /// keeps track of how many messages we've received from a given IP address since the last
2143    /// tick.
2144    ///
2145    /// This is used to count the number of messages received from a given IP address within an
2146    /// interval.
2147    ip_messages: HashMap<IpAddr, usize>,
2148    // keeps track of unique packet hashes
2149    unique_packets: schnellru::LruMap<B256, ()>,
2150}
2151
2152impl ReceiveCache {
2153    /// Updates the counter for each IP address and removes IPs that have exceeded the limit.
2154    ///
2155    /// This will decrement the counter for each IP address and remove IPs that have reached 0.
2156    fn tick_ips(&mut self, tick: usize) {
2157        self.ip_messages.retain(|_, count| {
2158            if let Some(reset) = count.checked_sub(tick) {
2159                *count = reset;
2160                true
2161            } else {
2162                false
2163            }
2164        });
2165    }
2166
2167    /// Increases the counter for the given IP address and returns the new count.
2168    fn inc_ip(&mut self, ip: IpAddr) -> usize {
2169        let ctn = self.ip_messages.entry(ip).or_default();
2170        *ctn = ctn.saturating_add(1);
2171        *ctn
2172    }
2173
2174    /// Returns true if we previously received the packet.
2175    ///
2176    /// A hit refreshes the entry so that a packet being replayed repeatedly stays cached.
2177    fn contains_packet(&mut self, hash: B256) -> bool {
2178        self.unique_packets.get(&hash).is_some()
2179    }
2180
2181    /// Remembers a packet we accepted.
2182    fn insert_packet(&mut self, hash: B256) {
2183        self.unique_packets.insert(hash, ());
2184    }
2185}
2186
2187impl Default for ReceiveCache {
2188    fn default() -> Self {
2189        Self {
2190            ip_messages: Default::default(),
2191            unique_packets: schnellru::LruMap::new(schnellru::ByLength::new(32)),
2192        }
2193    }
2194}
2195
2196/// The commands sent from the frontend [Discv4] to the service [`Discv4Service`].
2197enum Discv4Command {
2198    Add(NodeRecord),
2199    AddBootNode(NodeRecord),
2200    SetTcpPort(u16),
2201    SetEIP868RLPPair { key: Vec<u8>, rlp: Bytes },
2202    Ban(PeerId, IpAddr),
2203    BanPeer(PeerId),
2204    BanIp(IpAddr),
2205    Remove(PeerId),
2206    Lookup { node_id: Option<PeerId>, tx: Option<NodeRecordSender> },
2207    SetLookupInterval(Duration),
2208    Updates(OneshotSender<ReceiverStream<DiscoveryUpdate>>),
2209    Terminated,
2210}
2211
2212/// Event type receiver produces
2213#[derive(Debug)]
2214pub(crate) enum IngressEvent {
2215    /// Encountered an error when reading a datagram message.
2216    RecvError(io::Error),
2217    /// Received a bad message
2218    BadPacket(SocketAddr, DecodePacketError, Vec<u8>),
2219    /// Received a datagram from an address.
2220    Packet(SocketAddr, Packet),
2221}
2222
2223/// Tracks a sent ping
2224#[derive(Debug)]
2225struct PingRequest {
2226    // Timestamp when the request was sent.
2227    sent_at: Instant,
2228    // Node to which the request was sent.
2229    node: NodeRecord,
2230    // Hash sent in the Ping request
2231    echo_hash: B256,
2232    /// Why this ping was sent.
2233    reason: PingReason,
2234}
2235
2236/// Rotates the `PeerId` that is periodically looked up.
2237///
2238/// By selecting different targets, the lookups will be seeded with different ALPHA seed nodes.
2239#[derive(Debug)]
2240struct LookupTargetRotator {
2241    interval: usize,
2242    counter: usize,
2243}
2244
2245// === impl LookupTargetRotator ===
2246
2247impl LookupTargetRotator {
2248    /// Returns a rotator that always returns the local target.
2249    const fn local_only() -> Self {
2250        Self { interval: 1, counter: 0 }
2251    }
2252}
2253
2254impl Default for LookupTargetRotator {
2255    fn default() -> Self {
2256        Self {
2257            // every 4th lookup is our own node
2258            interval: 4,
2259            counter: 3,
2260        }
2261    }
2262}
2263
2264impl LookupTargetRotator {
2265    /// This will return the next node id to lookup
2266    fn next(&mut self, local: &PeerId) -> PeerId {
2267        self.counter += 1;
2268        self.counter %= self.interval;
2269        if self.counter == 0 {
2270            return *local
2271        }
2272        PeerId::random()
2273    }
2274}
2275
2276/// Tracks lookups across multiple `FindNode` requests.
2277///
2278/// If this type is dropped by all Clones, it will send all the discovered nodes to the listener, if
2279/// one is present.
2280#[derive(Clone, Debug)]
2281struct LookupContext {
2282    inner: Rc<LookupContextInner>,
2283}
2284
2285impl LookupContext {
2286    /// Create new context for a recursive lookup
2287    fn new(
2288        target: discv5::Key<NodeKey>,
2289        nearest_nodes: impl IntoIterator<Item = (Distance, NodeRecord)>,
2290        listener: Option<NodeRecordSender>,
2291    ) -> Self {
2292        let closest_nodes = nearest_nodes
2293            .into_iter()
2294            .map(|(distance, record)| {
2295                (distance, QueryNode { record, queried: false, responded: false })
2296            })
2297            .collect();
2298
2299        let inner = Rc::new(LookupContextInner {
2300            target,
2301            closest_nodes: RefCell::new(closest_nodes),
2302            listener,
2303        });
2304        Self { inner }
2305    }
2306
2307    /// Returns the target of this lookup
2308    fn target(&self) -> PeerId {
2309        self.inner.target.preimage().0
2310    }
2311
2312    fn closest(&self, num: usize) -> Vec<NodeRecord> {
2313        self.inner
2314            .closest_nodes
2315            .borrow()
2316            .iter()
2317            .filter(|(_, node)| !node.queried)
2318            .map(|(_, n)| n.record)
2319            .take(num)
2320            .collect()
2321    }
2322
2323    /// Returns the closest nodes that have not been queried yet.
2324    fn filter_closest<P>(&self, num: usize, filter: P) -> Vec<NodeRecord>
2325    where
2326        P: FnMut(&NodeRecord) -> bool,
2327    {
2328        self.inner
2329            .closest_nodes
2330            .borrow()
2331            .iter()
2332            .filter(|(_, node)| !node.queried)
2333            .map(|(_, n)| n.record)
2334            .filter(filter)
2335            .take(num)
2336            .collect()
2337    }
2338
2339    /// Inserts the node if it's missing
2340    fn add_node(&self, record: NodeRecord) {
2341        let distance = self.inner.target.distance(&kad_key(record.id));
2342        let mut closest = self.inner.closest_nodes.borrow_mut();
2343        if let btree_map::Entry::Vacant(entry) = closest.entry(distance) {
2344            entry.insert(QueryNode { record, queried: false, responded: false });
2345        }
2346    }
2347
2348    fn set_queried(&self, id: PeerId, val: bool) {
2349        if let Some((_, node)) =
2350            self.inner.closest_nodes.borrow_mut().iter_mut().find(|(_, node)| node.record.id == id)
2351        {
2352            node.queried = val;
2353        }
2354    }
2355
2356    /// Marks the node as queried
2357    fn mark_queried(&self, id: PeerId) {
2358        self.set_queried(id, true)
2359    }
2360
2361    /// Marks the node as not queried
2362    fn unmark_queried(&self, id: PeerId) {
2363        self.set_queried(id, false)
2364    }
2365
2366    /// Marks the node as responded
2367    fn mark_responded(&self, id: PeerId) {
2368        if let Some((_, node)) =
2369            self.inner.closest_nodes.borrow_mut().iter_mut().find(|(_, node)| node.record.id == id)
2370        {
2371            node.responded = true;
2372        }
2373    }
2374}
2375
2376// SAFETY: The [`Discv4Service`] is intended to be spawned as task which requires `Send`.
2377// The `LookupContext` is shared by all active `FindNode` requests that are part of the lookup step.
2378// Which can modify the context. The shared context is only ever accessed mutably when a `Neighbour`
2379// response is processed and all Clones are stored inside [`Discv4Service`], in other words it is
2380// guaranteed that there's only 1 owner ([`Discv4Service`]) of all possible [`Rc`] clones of
2381// [`LookupContext`].
2382unsafe impl Send for LookupContext {}
2383#[derive(Debug)]
2384struct LookupContextInner {
2385    /// The target to lookup.
2386    target: discv5::Key<NodeKey>,
2387    /// The closest nodes
2388    closest_nodes: RefCell<BTreeMap<Distance, QueryNode>>,
2389    /// A listener for all the nodes retrieved in this lookup
2390    ///
2391    /// This is present if the lookup was triggered manually via [Discv4] and we want to return all
2392    /// the nodes once the lookup finishes.
2393    listener: Option<NodeRecordSender>,
2394}
2395
2396impl Drop for LookupContextInner {
2397    fn drop(&mut self) {
2398        if let Some(tx) = self.listener.take() {
2399            // there's only 1 instance shared across `FindNode` requests, if this is dropped then
2400            // all requests finished, and we can send all results back
2401            let nodes = self
2402                .closest_nodes
2403                .take()
2404                .into_values()
2405                .filter(|node| node.responded)
2406                .map(|node| node.record)
2407                .collect();
2408            let _ = tx.send(nodes);
2409        }
2410    }
2411}
2412
2413/// Tracks the state of a recursive lookup step
2414#[derive(Debug, Clone, Copy)]
2415struct QueryNode {
2416    record: NodeRecord,
2417    queried: bool,
2418    responded: bool,
2419}
2420
2421#[derive(Debug)]
2422struct FindNodeRequest {
2423    // Timestamp when the request was sent.
2424    sent_at: Instant,
2425    // Number of items sent by the node
2426    response_count: usize,
2427    // Whether the request has been answered yet.
2428    answered: bool,
2429    /// Response buffer
2430    lookup_context: LookupContext,
2431}
2432
2433// === impl FindNodeRequest ===
2434
2435impl FindNodeRequest {
2436    fn new(resp: LookupContext) -> Self {
2437        Self { sent_at: Instant::now(), response_count: 0, answered: false, lookup_context: resp }
2438    }
2439}
2440
2441/// Cached signed `FindNode` packet to avoid redundant ECDSA signing during Kademlia lookups.
2442#[derive(Debug)]
2443struct CachedFindNode {
2444    target: PeerId,
2445    payload: Bytes,
2446    hash: B256,
2447    cached_at: Instant,
2448}
2449
2450impl CachedFindNode {
2451    /// Returns the cached `(payload, hash)` if the target matches and the cache is still fresh,
2452    /// or signs a new packet, updates the cache, and returns it.
2453    fn get_or_sign(
2454        cache: &mut Option<Self>,
2455        target: PeerId,
2456        ttl: Duration,
2457        secret_key: &secp256k1::SecretKey,
2458        expire: u64,
2459    ) -> (Bytes, B256) {
2460        if let Some(c) = cache.as_ref() &&
2461            c.target == target &&
2462            c.cached_at.elapsed() < ttl
2463        {
2464            return (c.payload.clone(), c.hash);
2465        }
2466
2467        let msg = Message::FindNode(FindNode { id: target, expire });
2468        let (payload, hash) = msg.encode(secret_key);
2469
2470        *cache = Some(Self { target, payload: payload.clone(), hash, cached_at: Instant::now() });
2471
2472        (payload, hash)
2473    }
2474}
2475
2476#[derive(Debug)]
2477struct EnrRequestState {
2478    // Timestamp when the request was sent.
2479    sent_at: Instant,
2480    // Hash sent in the Ping request
2481    echo_hash: B256,
2482}
2483
2484/// Stored node info.
2485#[derive(Debug, Clone, Eq, PartialEq)]
2486struct NodeEntry {
2487    /// Node record info.
2488    record: NodeRecord,
2489    /// Timestamp of last pong.
2490    last_seen: Instant,
2491    /// Last enr seq we retrieved via a ENR request.
2492    last_enr_seq: Option<u64>,
2493    /// `ForkId` if retrieved via ENR requests.
2494    fork_id: Option<ForkId>,
2495    /// Counter for failed _consecutive_ findNode requests.
2496    find_node_failures: u8,
2497    /// Whether the endpoint of the peer is proven.
2498    has_endpoint_proof: bool,
2499}
2500
2501// === impl NodeEntry ===
2502
2503impl NodeEntry {
2504    /// Creates a new, unpopulated entry
2505    fn new(record: NodeRecord) -> Self {
2506        Self {
2507            record,
2508            last_seen: Instant::now(),
2509            last_enr_seq: None,
2510            fork_id: None,
2511            find_node_failures: 0,
2512            has_endpoint_proof: false,
2513        }
2514    }
2515
2516    #[cfg(test)]
2517    fn new_proven(record: NodeRecord) -> Self {
2518        let mut node = Self::new(record);
2519        node.has_endpoint_proof = true;
2520        node
2521    }
2522
2523    /// Marks the entry with an established proof and resets the consecutive failure counter.
2524    const fn establish_proof(&mut self) {
2525        self.has_endpoint_proof = true;
2526        self.find_node_failures = 0;
2527    }
2528
2529    /// Returns true if the tracked find node failures exceed the max amount
2530    const fn exceeds_find_node_failures(&self, max_failures: u8) -> bool {
2531        self.find_node_failures >= max_failures
2532    }
2533
2534    /// Updates the last timestamp and sets the enr seq
2535    fn update_with_enr(&mut self, last_enr_seq: Option<u64>) -> Option<u64> {
2536        self.update_now(|s| std::mem::replace(&mut s.last_enr_seq, last_enr_seq))
2537    }
2538
2539    /// Increases the failed request counter
2540    const fn inc_failed_request(&mut self) {
2541        self.find_node_failures += 1;
2542    }
2543
2544    /// Updates the last timestamp and sets the enr seq
2545    fn update_with_fork_id(&mut self, fork_id: Option<ForkId>) -> Option<ForkId> {
2546        self.update_now(|s| std::mem::replace(&mut s.fork_id, fork_id))
2547    }
2548
2549    /// Updates the `last_seen` timestamp and calls the closure
2550    fn update_now<F, R>(&mut self, f: F) -> R
2551    where
2552        F: FnOnce(&mut Self) -> R,
2553    {
2554        self.last_seen = Instant::now();
2555        f(self)
2556    }
2557}
2558
2559// === impl NodeEntry ===
2560
2561impl NodeEntry {
2562    /// Returns true if the node should be re-pinged.
2563    fn is_expired(&self) -> bool {
2564        self.last_seen.elapsed() > (ENDPOINT_PROOF_EXPIRATION / 2)
2565    }
2566}
2567
2568/// Represents why a ping is issued
2569#[derive(Debug)]
2570enum PingReason {
2571    /// Initial ping to a previously unknown peer that was inserted into the table.
2572    InitialInsert,
2573    /// A ping to a peer to establish a bond (endpoint proof).
2574    EstablishBond,
2575    /// Re-ping a peer.
2576    RePing,
2577    /// Part of a lookup to ensure endpoint is proven before we can send a `FindNode` request.
2578    Lookup(NodeRecord, LookupContext),
2579}
2580
2581/// Represents node related updates state changes in the underlying node table
2582#[derive(Debug, Clone)]
2583pub enum DiscoveryUpdate {
2584    /// A new node was discovered _and_ added to the table.
2585    Added(NodeRecord),
2586    /// A new node was discovered but _not_ added to the table because it is currently full.
2587    DiscoveredAtCapacity(NodeRecord),
2588    /// Received a [`ForkId`] via EIP-868 for the given [`NodeRecord`].
2589    EnrForkId(NodeRecord, ForkId),
2590    /// Node that was removed from the table
2591    Removed(PeerId),
2592    /// A series of updates
2593    Batch(Vec<Self>),
2594}
2595
2596#[cfg(test)]
2597mod tests {
2598    use super::*;
2599    use crate::test_utils::{create_discv4, create_discv4_with_config, rng_endpoint, rng_record};
2600    use alloy_primitives::hex;
2601    use alloy_rlp::{Decodable, Encodable};
2602    use rand_08::Rng;
2603    use reth_ethereum_forks::{EnrForkIdEntry, ForkHash};
2604    use reth_network_peers::mainnet_nodes;
2605    use secp256k1::SECP256K1;
2606    use std::future::poll_fn;
2607
2608    #[tokio::test]
2609    async fn test_duplicate_packet_rejected_without_decoding() {
2610        let secret_key = SecretKey::new(&mut rand_08::thread_rng());
2611        let local_id = pk2id(&secret_key.public_key(SECP256K1));
2612        let (tx, mut rx) = mpsc::channel(16);
2613        let mut handler = IngressHandler::new(tx, local_id);
2614
2615        let remote_key = SecretKey::new(&mut rand_08::thread_rng());
2616        let msg = Message::Ping(Ping {
2617            from: rng_endpoint(&mut rand_08::thread_rng()),
2618            to: rng_endpoint(&mut rand_08::thread_rng()),
2619            expire: u64::MAX,
2620            enr_sq: None,
2621        });
2622        let (packet, hash) = msg.encode(&remote_key);
2623        let src = "10.0.0.1:30303".parse().unwrap();
2624
2625        handler.handle_packet(&packet, src).await;
2626        assert!(matches!(rx.try_recv(), Ok(IngressEvent::Packet(_, _))));
2627
2628        // the replay is dropped on the hash prefix alone
2629        handler.handle_packet(&packet, src).await;
2630        assert!(rx.try_recv().is_err());
2631        assert!(handler.cache.contains_packet(hash));
2632    }
2633
2634    #[tokio::test]
2635    async fn test_undecodable_packet_does_not_poison_cache() {
2636        let secret_key = SecretKey::new(&mut rand_08::thread_rng());
2637        let local_id = pk2id(&secret_key.public_key(SECP256K1));
2638        let (tx, mut rx) = mpsc::channel(16);
2639        let mut handler = IngressHandler::new(tx, local_id);
2640
2641        let remote_key = SecretKey::new(&mut rand_08::thread_rng());
2642        let msg = Message::Ping(Ping {
2643            from: rng_endpoint(&mut rand_08::thread_rng()),
2644            to: rng_endpoint(&mut rand_08::thread_rng()),
2645            expire: u64::MAX,
2646            enr_sq: None,
2647        });
2648        let (packet, hash) = msg.encode(&remote_key);
2649        let src = "10.0.0.1:30303".parse().unwrap();
2650
2651        // a packet carrying the right hash prefix but a corrupt body must not be remembered,
2652        // otherwise it could be used to suppress the real packet
2653        let mut forged = packet.to_vec();
2654        let last = forged.len() - 1;
2655        forged[last] ^= 0xff;
2656        handler.handle_packet(&forged, src).await;
2657        assert!(matches!(rx.try_recv(), Ok(IngressEvent::BadPacket(..))));
2658        assert!(!handler.cache.contains_packet(hash));
2659
2660        // so the genuine packet still gets through
2661        handler.handle_packet(&packet, src).await;
2662        assert!(matches!(rx.try_recv(), Ok(IngressEvent::Packet(_, _))));
2663    }
2664
2665    #[tokio::test]
2666    async fn test_configured_enr_forkid_entry() {
2667        let fork: ForkId = ForkId { hash: ForkHash([220, 233, 108, 45]), next: 0u64 };
2668        let mut disc_conf = Discv4Config::default();
2669        disc_conf.add_eip868_pair("eth", EnrForkIdEntry::from(fork));
2670        let (_discv4, service) = create_discv4_with_config(disc_conf).await;
2671        let eth = service.local_eip_868_enr.get_raw_rlp(b"eth").unwrap();
2672        let fork_entry_id = EnrForkIdEntry::decode(&mut &eth[..]).unwrap();
2673
2674        let raw: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2675        let decoded = EnrForkIdEntry::decode(&mut &raw[..]).unwrap();
2676        let expected = EnrForkIdEntry {
2677            fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2678        };
2679        assert_eq!(expected, fork_entry_id);
2680        assert_eq!(expected, decoded);
2681    }
2682
2683    #[test]
2684    fn test_enr_forkid_entry_decode() {
2685        let raw: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2686        let decoded = EnrForkIdEntry::decode(&mut &raw[..]).unwrap();
2687        let expected = EnrForkIdEntry {
2688            fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2689        };
2690        assert_eq!(expected, decoded);
2691    }
2692
2693    #[test]
2694    fn test_enr_forkid_entry_encode() {
2695        let original = EnrForkIdEntry {
2696            fork_id: ForkId { hash: ForkHash([0xdc, 0xe9, 0x6c, 0x2d]), next: 0 },
2697        };
2698        let expected: [u8; 8] = [0xc7, 0xc6, 0x84, 0xdc, 0xe9, 0x6c, 0x2d, 0x80];
2699        let mut encoded = Vec::with_capacity(expected.len());
2700        original.encode(&mut encoded);
2701        assert_eq!(&expected[..], encoded.as_slice());
2702    }
2703
2704    #[tokio::test]
2705    async fn test_request_enr_of_proven_node() {
2706        reth_tracing::init_test_tracing();
2707        let mut rng = rand_08::thread_rng();
2708        let (_discv4, mut service) = create_discv4().await;
2709
2710        let id = PeerId::random();
2711        let addr: SocketAddr = (Ipv4Addr::new(10, 0, 0, 1), DEFAULT_DISCOVERY_PORT).into();
2712        let mut entry = NodeEntry::new_proven(NodeRecord::new(addr, id));
2713        entry.last_enr_seq = Some(1);
2714        let _ = service.kbuckets.insert_or_update(
2715            &kad_key(id),
2716            entry,
2717            NodeStatus {
2718                direction: ConnectionDirection::Outgoing,
2719                state: ConnectionState::Connected,
2720            },
2721        );
2722        let ping = |service: &Discv4Service, rng: &mut rand_08::rngs::ThreadRng| Ping {
2723            from: rng_endpoint(rng),
2724            to: rng_endpoint(rng),
2725            expire: service.ping_expiration(),
2726            enr_sq: Some(2),
2727        };
2728
2729        // the node announces a newer record.
2730        service.on_ping(ping(&service, &mut rng), addr, id, B256::random());
2731        assert!(service.pending_enr_requests.contains_key(&id));
2732
2733        // the request goes unanswered, so the next ping has to ask again.
2734        service.evict_expired_requests(Instant::now() + service.config.enr_expiration * 2);
2735        assert!(!service.pending_enr_requests.contains_key(&id));
2736        service.on_ping(ping(&service, &mut rng), addr, id, B256::random());
2737        assert!(service.pending_enr_requests.contains_key(&id));
2738    }
2739
2740    #[test]
2741    fn test_local_rotator() {
2742        let id = PeerId::random();
2743        let mut rotator = LookupTargetRotator::local_only();
2744        assert_eq!(rotator.next(&id), id);
2745        assert_eq!(rotator.next(&id), id);
2746    }
2747
2748    #[test]
2749    fn test_rotator() {
2750        let id = PeerId::random();
2751        let mut rotator = LookupTargetRotator::default();
2752        assert_eq!(rotator.next(&id), id);
2753        assert_ne!(rotator.next(&id), id);
2754        assert_ne!(rotator.next(&id), id);
2755        assert_ne!(rotator.next(&id), id);
2756        assert_eq!(rotator.next(&id), id);
2757    }
2758
2759    #[tokio::test]
2760    async fn test_pending_ping() {
2761        let (_, mut service) = create_discv4().await;
2762
2763        let local_addr = service.local_addr();
2764
2765        let mut num_inserted = 0;
2766        loop {
2767            let node = NodeRecord::new(local_addr, PeerId::random());
2768            if service.add_node(node) {
2769                num_inserted += 1;
2770                assert!(service.pending_pings.contains_key(&node.id));
2771                assert_eq!(service.pending_pings.len(), num_inserted);
2772                if num_inserted == MAX_NODES_PING {
2773                    break
2774                }
2775            }
2776        }
2777
2778        // `pending_pings` is full, insert into `queued_pings`.
2779        num_inserted = 0;
2780        for _ in 0..MAX_NODES_PING {
2781            let node = NodeRecord::new(local_addr, PeerId::random());
2782            if service.add_node(node) {
2783                num_inserted += 1;
2784                assert!(!service.pending_pings.contains_key(&node.id));
2785                assert_eq!(service.pending_pings.len(), MAX_NODES_PING);
2786                assert_eq!(service.queued_pings.len(), num_inserted);
2787            }
2788        }
2789    }
2790
2791    // Bootstraps with mainnet boot nodes
2792    #[tokio::test(flavor = "multi_thread")]
2793    #[ignore]
2794    async fn test_mainnet_lookup() {
2795        reth_tracing::init_test_tracing();
2796        let fork_id = ForkId { hash: ForkHash(hex!("743f3d89")), next: 16191202 };
2797
2798        let all_nodes = mainnet_nodes();
2799        let config = Discv4Config::builder()
2800            .add_boot_nodes(all_nodes)
2801            .lookup_interval(Duration::from_secs(1))
2802            .add_eip868_pair("eth", fork_id)
2803            .build();
2804        let (_discv4, mut service) = create_discv4_with_config(config).await;
2805
2806        let mut updates = service.update_stream();
2807
2808        let _handle = service.spawn();
2809
2810        let mut table = HashMap::new();
2811        while let Some(update) = updates.next().await {
2812            match update {
2813                DiscoveryUpdate::EnrForkId(record, fork_id) => {
2814                    println!("{record:?}, {fork_id:?}");
2815                }
2816                DiscoveryUpdate::Added(record) => {
2817                    table.insert(record.id, record);
2818                }
2819                DiscoveryUpdate::Removed(id) => {
2820                    table.remove(&id);
2821                }
2822                _ => {}
2823            }
2824            println!("total peers {}", table.len());
2825        }
2826    }
2827
2828    #[tokio::test]
2829    async fn test_mapped_ipv4() {
2830        reth_tracing::init_test_tracing();
2831        let mut rng = rand_08::thread_rng();
2832        let config = Discv4Config::builder().build();
2833        let (_discv4, mut service) = create_discv4_with_config(config).await;
2834
2835        let v4: Ipv4Addr = "0.0.0.0".parse().unwrap();
2836        let v6 = v4.to_ipv6_mapped();
2837        let addr: SocketAddr = (v6, DEFAULT_DISCOVERY_PORT).into();
2838
2839        let ping = Ping {
2840            from: rng_endpoint(&mut rng),
2841            to: rng_endpoint(&mut rng),
2842            expire: service.ping_expiration(),
2843            enr_sq: Some(rng.r#gen()),
2844        };
2845
2846        let id = PeerId::random();
2847        service.on_ping(ping, addr, id, B256::random());
2848
2849        let key = kad_key(id);
2850        match service.kbuckets.entry(&key) {
2851            kbucket::Entry::Present(entry, _) => {
2852                let node_addr = entry.value().record.address;
2853                assert!(node_addr.is_ipv4());
2854                assert_eq!(node_addr, IpAddr::from(v4));
2855            }
2856            _ => unreachable!(),
2857        };
2858    }
2859
2860    #[tokio::test]
2861    async fn test_respect_ping_expiration() {
2862        reth_tracing::init_test_tracing();
2863        let mut rng = rand_08::thread_rng();
2864        let config = Discv4Config::builder().build();
2865        let (_discv4, mut service) = create_discv4_with_config(config).await;
2866
2867        let v4: Ipv4Addr = "0.0.0.0".parse().unwrap();
2868        let v6 = v4.to_ipv6_mapped();
2869        let addr: SocketAddr = (v6, DEFAULT_DISCOVERY_PORT).into();
2870
2871        let ping = Ping {
2872            from: rng_endpoint(&mut rng),
2873            to: rng_endpoint(&mut rng),
2874            expire: SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() - 1,
2875            enr_sq: Some(rng.r#gen()),
2876        };
2877
2878        let id = PeerId::random();
2879        service.on_ping(ping, addr, id, B256::random());
2880
2881        let key = kad_key(id);
2882        match service.kbuckets.entry(&key) {
2883            kbucket::Entry::Absent(_) => {}
2884            _ => unreachable!(),
2885        };
2886    }
2887
2888    #[tokio::test]
2889    async fn test_single_lookups() {
2890        reth_tracing::init_test_tracing();
2891
2892        let config = Discv4Config::builder().build();
2893        let (_discv4, mut service) = create_discv4_with_config(config.clone()).await;
2894
2895        let id = PeerId::random();
2896        let key = kad_key(id);
2897        let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
2898
2899        let _ = service.kbuckets.insert_or_update(
2900            &key,
2901            NodeEntry::new_proven(record),
2902            NodeStatus {
2903                direction: ConnectionDirection::Incoming,
2904                state: ConnectionState::Connected,
2905            },
2906        );
2907
2908        service.lookup_self();
2909        assert_eq!(service.pending_find_nodes.len(), 1);
2910
2911        poll_fn(|cx| {
2912            let _ = service.poll(cx);
2913            assert_eq!(service.pending_find_nodes.len(), 1);
2914
2915            Poll::Ready(())
2916        })
2917        .await;
2918    }
2919
2920    #[tokio::test]
2921    async fn test_on_neighbours_recursive_lookup() {
2922        reth_tracing::init_test_tracing();
2923
2924        let config = Discv4Config::builder().build();
2925        let (_discv4, mut service) = create_discv4_with_config(config.clone()).await;
2926        let (_discv4, mut service2) = create_discv4_with_config(config).await;
2927
2928        let id = PeerId::random();
2929        let key = kad_key(id);
2930        let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
2931
2932        let _ = service.kbuckets.insert_or_update(
2933            &key,
2934            NodeEntry::new_proven(record),
2935            NodeStatus {
2936                direction: ConnectionDirection::Incoming,
2937                state: ConnectionState::Connected,
2938            },
2939        );
2940        // Needed in this test to populate self.pending_find_nodes for as a prereq to a valid
2941        // on_neighbours request
2942        service.lookup_self();
2943        assert_eq!(service.pending_find_nodes.len(), 1);
2944
2945        poll_fn(|cx| {
2946            let _ = service.poll(cx);
2947            assert_eq!(service.pending_find_nodes.len(), 1);
2948
2949            Poll::Ready(())
2950        })
2951        .await;
2952
2953        let expiry = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs() +
2954            10000000000000;
2955        let msg = Neighbours { nodes: vec![service2.local_node_record], expire: expiry };
2956        service.on_neighbours(msg, record.tcp_addr(), id);
2957        // wait for the processed ping
2958        let event = poll_fn(|cx| service2.poll(cx)).await;
2959        assert_eq!(event, Discv4Event::Ping);
2960        // assert that no find_node req has been added here on top of the initial one, since both
2961        // sides of the endpoint proof is not completed here
2962        assert_eq!(service.pending_find_nodes.len(), 1);
2963        // we now wait for PONG
2964        let event = poll_fn(|cx| service.poll(cx)).await;
2965        assert_eq!(event, Discv4Event::Pong);
2966        // Ideally we want to assert against service.pending_lookup.len() here - but because the
2967        // service2 sends Pong and Ping consecutivley on_ping(), the pending_lookup table gets
2968        // drained almost immediately - and no way to grab the handle to its intermediary state here
2969        // :(
2970        let event = poll_fn(|cx| service.poll(cx)).await;
2971        assert_eq!(event, Discv4Event::Ping);
2972        // assert that we've added the find_node req here after both sides of the endpoint proof is
2973        // done
2974        assert_eq!(service.pending_find_nodes.len(), 2);
2975    }
2976
2977    #[tokio::test]
2978    async fn test_no_local_in_closest() {
2979        reth_tracing::init_test_tracing();
2980
2981        let config = Discv4Config::builder().build();
2982        let (_discv4, mut service) = create_discv4_with_config(config).await;
2983
2984        let target_key = kad_key(PeerId::random());
2985
2986        let id = PeerId::random();
2987        let key = kad_key(id);
2988        let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
2989
2990        let _ = service.kbuckets.insert_or_update(
2991            &key,
2992            NodeEntry::new(record),
2993            NodeStatus {
2994                direction: ConnectionDirection::Incoming,
2995                state: ConnectionState::Connected,
2996            },
2997        );
2998
2999        let closest = service
3000            .kbuckets
3001            .closest_values(&target_key)
3002            .map(|n| n.value.record)
3003            .take(MAX_NODES_PER_BUCKET)
3004            .collect::<Vec<_>>();
3005
3006        assert_eq!(closest.len(), 1);
3007        assert!(!closest.iter().any(|r| r.id == *service.local_peer_id()));
3008    }
3009
3010    #[tokio::test]
3011    async fn test_random_lookup() {
3012        reth_tracing::init_test_tracing();
3013
3014        let config = Discv4Config::builder().build();
3015        let (_discv4, mut service) = create_discv4_with_config(config).await;
3016
3017        let target = PeerId::random();
3018
3019        let id = PeerId::random();
3020        let key = kad_key(id);
3021        let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3022
3023        let _ = service.kbuckets.insert_or_update(
3024            &key,
3025            NodeEntry::new_proven(record),
3026            NodeStatus {
3027                direction: ConnectionDirection::Incoming,
3028                state: ConnectionState::Connected,
3029            },
3030        );
3031
3032        service.lookup(target);
3033        assert_eq!(service.pending_find_nodes.len(), 1);
3034
3035        let ctx = service.pending_find_nodes.values().next().unwrap().lookup_context.clone();
3036
3037        assert_eq!(ctx.target(), target);
3038        assert_eq!(ctx.inner.closest_nodes.borrow().len(), 1);
3039
3040        ctx.add_node(record);
3041        assert_eq!(ctx.inner.closest_nodes.borrow().len(), 1);
3042    }
3043
3044    #[tokio::test]
3045    async fn test_reping_on_find_node_failures() {
3046        reth_tracing::init_test_tracing();
3047
3048        let config = Discv4Config::builder().build();
3049        let (_discv4, mut service) = create_discv4_with_config(config).await;
3050
3051        let target = PeerId::random();
3052
3053        let id = PeerId::random();
3054        let key = kad_key(id);
3055        let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3056
3057        let mut entry = NodeEntry::new_proven(record);
3058        entry.find_node_failures = u8::MAX;
3059        let _ = service.kbuckets.insert_or_update(
3060            &key,
3061            entry,
3062            NodeStatus {
3063                direction: ConnectionDirection::Incoming,
3064                state: ConnectionState::Connected,
3065            },
3066        );
3067
3068        service.lookup(target);
3069        assert_eq!(service.pending_find_nodes.len(), 0);
3070        assert_eq!(service.pending_pings.len(), 1);
3071
3072        service.update_on_pong(record, None);
3073
3074        service
3075            .on_entry(record.id, |entry| {
3076                // reset on pong
3077                assert_eq!(entry.find_node_failures, 0);
3078                assert!(entry.has_endpoint_proof);
3079            })
3080            .unwrap();
3081    }
3082
3083    #[tokio::test]
3084    async fn test_service_commands() {
3085        reth_tracing::init_test_tracing();
3086
3087        let config = Discv4Config::builder().build();
3088        let (discv4, mut service) = create_discv4_with_config(config).await;
3089
3090        service.lookup_self();
3091
3092        let _handle = service.spawn();
3093        discv4.send_lookup_self();
3094        let _ = discv4.lookup_self().await;
3095    }
3096
3097    #[tokio::test]
3098    async fn test_requests_timeout() {
3099        reth_tracing::init_test_tracing();
3100        let fork_id = ForkId { hash: ForkHash(hex!("743f3d89")), next: 16191202 };
3101
3102        let config = Discv4Config::builder()
3103            .request_timeout(Duration::from_millis(200))
3104            .ping_expiration(Duration::from_millis(200))
3105            .lookup_neighbours_expiration(Duration::from_millis(200))
3106            .add_eip868_pair("eth", fork_id)
3107            .build();
3108        let (_disv4, mut service) = create_discv4_with_config(config).await;
3109
3110        let id = PeerId::random();
3111        let key = kad_key(id);
3112        let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), id);
3113
3114        let _ = service.kbuckets.insert_or_update(
3115            &key,
3116            NodeEntry::new_proven(record),
3117            NodeStatus {
3118                direction: ConnectionDirection::Incoming,
3119                state: ConnectionState::Connected,
3120            },
3121        );
3122
3123        service.lookup_self();
3124        assert_eq!(service.pending_find_nodes.len(), 1);
3125
3126        let ctx = service.pending_find_nodes.values().next().unwrap().lookup_context.clone();
3127
3128        service.pending_lookup.insert(record.id, (Instant::now(), ctx));
3129
3130        assert_eq!(service.pending_lookup.len(), 1);
3131
3132        let ping = Ping {
3133            from: service.local_node_record.into(),
3134            to: record.into(),
3135            expire: service.ping_expiration(),
3136            enr_sq: service.enr_seq(),
3137        };
3138        let echo_hash = service.send_packet(Message::Ping(ping), record.udp_addr());
3139        let ping_request = PingRequest {
3140            sent_at: Instant::now(),
3141            node: record,
3142            echo_hash,
3143            reason: PingReason::InitialInsert,
3144        };
3145        service.pending_pings.insert(record.id, ping_request);
3146
3147        assert_eq!(service.pending_pings.len(), 1);
3148
3149        tokio::time::sleep(Duration::from_secs(1)).await;
3150
3151        poll_fn(|cx| {
3152            let _ = service.poll(cx);
3153
3154            assert_eq!(service.pending_find_nodes.len(), 0);
3155            assert_eq!(service.pending_lookup.len(), 0);
3156            assert_eq!(service.pending_pings.len(), 0);
3157
3158            Poll::Ready(())
3159        })
3160        .await;
3161    }
3162
3163    // sends a PING packet with wrong 'to' field and expects a PONG response.
3164    #[tokio::test(flavor = "multi_thread")]
3165    async fn test_check_wrong_to() {
3166        reth_tracing::init_test_tracing();
3167
3168        let config = Discv4Config::builder().external_ip_resolver(None).build();
3169        let (_discv4, mut service_1) = create_discv4_with_config(config.clone()).await;
3170        let (_discv4, mut service_2) = create_discv4_with_config(config).await;
3171
3172        // ping node 2 with wrong to field
3173        let mut ping = Ping {
3174            from: service_1.local_node_record.into(),
3175            to: service_2.local_node_record.into(),
3176            expire: service_1.ping_expiration(),
3177            enr_sq: service_1.enr_seq(),
3178        };
3179        ping.to.address = "192.0.2.0".parse().unwrap();
3180
3181        let echo_hash = service_1.send_packet(Message::Ping(ping), service_2.local_addr());
3182        let ping_request = PingRequest {
3183            sent_at: Instant::now(),
3184            node: service_2.local_node_record,
3185            echo_hash,
3186            reason: PingReason::InitialInsert,
3187        };
3188        service_1.pending_pings.insert(*service_2.local_peer_id(), ping_request);
3189
3190        // wait for the processed ping
3191        let event = poll_fn(|cx| service_2.poll(cx)).await;
3192        assert_eq!(event, Discv4Event::Ping);
3193
3194        // we now wait for PONG
3195        let event = poll_fn(|cx| service_1.poll(cx)).await;
3196        assert_eq!(event, Discv4Event::Pong);
3197        // followed by a ping
3198        let event = poll_fn(|cx| service_1.poll(cx)).await;
3199        assert_eq!(event, Discv4Event::Ping);
3200    }
3201
3202    #[tokio::test(flavor = "multi_thread")]
3203    async fn test_check_ping_pong() {
3204        reth_tracing::init_test_tracing();
3205
3206        let config = Discv4Config::builder().external_ip_resolver(None).build();
3207        let (_discv4, mut service_1) = create_discv4_with_config(config.clone()).await;
3208        let (_discv4, mut service_2) = create_discv4_with_config(config).await;
3209
3210        // send ping from 1 -> 2
3211        service_1.add_node(service_2.local_node_record);
3212
3213        // wait for the processed ping
3214        let event = poll_fn(|cx| service_2.poll(cx)).await;
3215        assert_eq!(event, Discv4Event::Ping);
3216
3217        // node is now in the table but not connected yet
3218        let key1 = kad_key(*service_1.local_peer_id());
3219        match service_2.kbuckets.entry(&key1) {
3220            kbucket::Entry::Present(_entry, status) => {
3221                assert!(!status.is_connected());
3222            }
3223            _ => unreachable!(),
3224        }
3225
3226        // we now wait for PONG
3227        let event = poll_fn(|cx| service_1.poll(cx)).await;
3228        assert_eq!(event, Discv4Event::Pong);
3229
3230        // endpoint is proven
3231        let key2 = kad_key(*service_2.local_peer_id());
3232        match service_1.kbuckets.entry(&key2) {
3233            kbucket::Entry::Present(_entry, status) => {
3234                assert!(status.is_connected());
3235            }
3236            _ => unreachable!(),
3237        }
3238
3239        // we now wait for the PING initiated by 2
3240        let event = poll_fn(|cx| service_1.poll(cx)).await;
3241        assert_eq!(event, Discv4Event::Ping);
3242
3243        // Drain events from service_2 until we see the Pong. Intermediate EnrRequest and
3244        // FindNode events are expected: ENR requests come from the ping handshake, and FindNode
3245        // arrives because service_1 resets its lookup interval on the first bootnode pong.
3246        tokio::time::timeout(Duration::from_secs(5), async {
3247            loop {
3248                let event = poll_fn(|cx| service_2.poll(cx)).await;
3249                match event {
3250                    Discv4Event::Pong => break,
3251                    Discv4Event::EnrRequest | Discv4Event::FindNode => {}
3252                    ev => unreachable!("{ev:?}"),
3253                }
3254            }
3255        })
3256        .await
3257        .expect("timed out waiting for Pong from service_2");
3258
3259        // endpoint is proven
3260        match service_2.kbuckets.entry(&key1) {
3261            kbucket::Entry::Present(_entry, status) => {
3262                assert!(status.is_connected());
3263            }
3264            ev => unreachable!("{ev:?}"),
3265        }
3266    }
3267
3268    #[test]
3269    fn test_insert() {
3270        let local_node_record = rng_record(&mut rand_08::thread_rng());
3271        let mut kbuckets: KBucketsTable<NodeKey, NodeEntry> = KBucketsTable::new(
3272            NodeKey::from(&local_node_record).into(),
3273            Duration::from_secs(60),
3274            MAX_NODES_PER_BUCKET,
3275            None,
3276            None,
3277        );
3278
3279        let new_record = rng_record(&mut rand_08::thread_rng());
3280        let key = kad_key(new_record.id);
3281        match kbuckets.entry(&key) {
3282            kbucket::Entry::Absent(entry) => {
3283                let node = NodeEntry::new(new_record);
3284                let _ = entry.insert(
3285                    node,
3286                    NodeStatus {
3287                        direction: ConnectionDirection::Outgoing,
3288                        state: ConnectionState::Disconnected,
3289                    },
3290                );
3291            }
3292            _ => {
3293                unreachable!()
3294            }
3295        };
3296        match kbuckets.entry(&key) {
3297            kbucket::Entry::Present(_, _) => {}
3298            _ => {
3299                unreachable!()
3300            }
3301        }
3302    }
3303
3304    #[tokio::test]
3305    async fn test_bootnode_not_in_update_stream() {
3306        reth_tracing::init_test_tracing();
3307        let (_, service_1) = create_discv4().await;
3308        let peerid_1 = *service_1.local_peer_id();
3309
3310        let config = Discv4Config::builder().add_boot_node(service_1.local_node_record).build();
3311        service_1.spawn();
3312
3313        let (_, mut service_2) = create_discv4_with_config(config).await;
3314
3315        let mut updates = service_2.update_stream();
3316
3317        service_2.spawn();
3318
3319        // Poll for events for a reasonable time
3320        let mut bootnode_appeared = false;
3321        let timeout = tokio::time::sleep(Duration::from_secs(1));
3322        tokio::pin!(timeout);
3323
3324        loop {
3325            tokio::select! {
3326                Some(update) = updates.next() => {
3327                    if let DiscoveryUpdate::Added(record) = update
3328                        && record.id == peerid_1 {
3329                            bootnode_appeared = true;
3330                            break;
3331                        }
3332                }
3333                _ = &mut timeout => break,
3334            }
3335        }
3336
3337        // Assert bootnode did not appear in update stream
3338        assert!(bootnode_appeared, "Bootnode should appear in update stream");
3339    }
3340
3341    fn insert_proven_node(service: &mut Discv4Service, record: NodeRecord) {
3342        let key = kad_key(record.id);
3343        let _ = service.kbuckets.insert_or_update(
3344            &key,
3345            NodeEntry::new_proven(record),
3346            NodeStatus {
3347                direction: ConnectionDirection::Incoming,
3348                state: ConnectionState::Connected,
3349            },
3350        );
3351    }
3352
3353    fn insert_initial_ping(service: &mut Discv4Service, record: NodeRecord) -> B256 {
3354        let echo_hash = B256::random();
3355        service.pending_pings.insert(
3356            record.id,
3357            PingRequest {
3358                sent_at: Instant::now(),
3359                node: record,
3360                echo_hash,
3361                reason: PingReason::InitialInsert,
3362            },
3363        );
3364        echo_hash
3365    }
3366
3367    fn make_pong(service: &Discv4Service, echo_hash: B256) -> Pong {
3368        Pong {
3369            to: rng_endpoint(&mut rand_08::thread_rng()),
3370            echo: echo_hash,
3371            expire: service.ping_expiration(),
3372            enr_sq: None,
3373        }
3374    }
3375
3376    #[tokio::test]
3377    async fn test_lookup_reset_on_first_bootnode_pong() {
3378        let record = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3379        let config = Discv4Config::builder().add_boot_node(record).build();
3380        let (_discv4, mut service) = create_discv4_with_config(config).await;
3381
3382        // 1. initial state
3383        assert!(service.pending_lookup_reset);
3384
3385        // 2. setup: proven bootnode + pending InitialInsert ping
3386        insert_proven_node(&mut service, record);
3387        let echo_hash = insert_initial_ping(&mut service, record);
3388
3389        // 3. input: pong arrives
3390        service.on_pong(make_pong(&service, echo_hash), record.udp_addr(), record.id);
3391
3392        // 4. flag should be consumed — interval was reset
3393        assert!(!service.pending_lookup_reset, "flag should be consumed");
3394    }
3395
3396    #[tokio::test]
3397    async fn test_lookup_reset_fires_only_once() {
3398        let records: Vec<_> = (0..2)
3399            .map(|_| NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random()))
3400            .collect();
3401        let config = Discv4Config::builder().add_boot_nodes(records.clone()).build();
3402        let (_discv4, mut service) = create_discv4_with_config(config).await;
3403
3404        // 1. setup: two proven bootnodes with pending InitialInsert pings
3405        for &r in &records {
3406            insert_proven_node(&mut service, r);
3407        }
3408        let hashes: Vec<_> =
3409            records.iter().map(|r| insert_initial_ping(&mut service, *r)).collect();
3410
3411        // 2. first pong -> consumes the flag (resets the interval)
3412        service.on_pong(make_pong(&service, hashes[0]), records[0].udp_addr(), records[0].id);
3413        assert!(!service.pending_lookup_reset);
3414
3415        // 3. second pong -> flag already consumed, no second reset
3416        service.on_pong(make_pong(&service, hashes[1]), records[1].udp_addr(), records[1].id);
3417        assert!(!service.pending_lookup_reset);
3418    }
3419
3420    #[tokio::test]
3421    async fn test_lookup_reset_not_triggered_by_non_bootnode() {
3422        let bootnode = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3423        let config = Discv4Config::builder().add_boot_node(bootnode).build();
3424        let (_discv4, mut service) = create_discv4_with_config(config).await;
3425
3426        assert!(service.pending_lookup_reset);
3427
3428        // a non-bootnode pong should not consume the flag
3429        let stranger = NodeRecord::new("0.0.0.0:0".parse().unwrap(), PeerId::random());
3430        insert_proven_node(&mut service, stranger);
3431        let echo_hash = insert_initial_ping(&mut service, stranger);
3432        service.on_pong(make_pong(&service, echo_hash), stranger.udp_addr(), stranger.id);
3433
3434        assert!(service.pending_lookup_reset, "flag should not be consumed by non-bootnode");
3435    }
3436
3437    #[tokio::test]
3438    async fn test_lookup_reset_disabled_when_lookup_disabled() {
3439        let config = Discv4Config::builder().enable_lookup(false).build();
3440        let (_discv4, service) = create_discv4_with_config(config).await;
3441
3442        // flag should be false when lookups are disabled
3443        assert!(!service.pending_lookup_reset);
3444    }
3445
3446    #[tokio::test]
3447    async fn test_enr_response_preserves_pending_request() {
3448        let fork_id = ForkId { hash: ForkHash([1, 2, 3, 4]), next: 0 };
3449        let mut config = Discv4Config::default();
3450        config.add_eip868_pair("eth", EnrForkIdEntry::from(fork_id));
3451        let (_remote, mut remote) = create_discv4_with_config(config).await;
3452        let (_discv4, mut service) = create_discv4().await;
3453        let record = remote.local_node_record;
3454        let id = record.id;
3455        let addr = record.udp_addr();
3456        insert_proven_node(&mut service, record);
3457        service.update_on_reping(record, remote.enr_seq());
3458        let old_enr = remote.local_eip_868_enr.clone();
3459
3460        // Model a request sent in the previous second without waiting for the wall clock.
3461        let (_, old_hash) =
3462            Message::EnrRequest(EnrRequest { expire: service.enr_request_expiration() - 1 })
3463                .encode(&service.secret_key);
3464        service.pending_enr_requests.get_mut(&id).unwrap().echo_hash = old_hash;
3465
3466        // A newer sequence replaces the pending request before its response arrives.
3467        remote.local_eip_868_enr.set_tcp4(30304, &remote.secret_key).unwrap();
3468        service.on_ping(
3469            Ping {
3470                from: record.into(),
3471                to: service.local_node_record.into(),
3472                expire: service.ping_expiration(),
3473                enr_sq: remote.enr_seq(),
3474            },
3475            addr,
3476            id,
3477            B256::random(),
3478        );
3479        let request_hash = service.pending_enr_requests[&id].echo_hash;
3480        assert_ne!(request_hash, old_hash);
3481
3482        service.on_enr_response(EnrResponse { request_hash: old_hash, enr: old_enr }, addr, id);
3483        assert_eq!(service.pending_enr_requests[&id].echo_hash, request_hash);
3484        assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(None));
3485
3486        // A response with the right hash but a different node identity is also ignored.
3487        service.on_enr_response(
3488            EnrResponse { request_hash, enr: service.local_eip_868_enr.clone() },
3489            addr,
3490            id,
3491        );
3492        assert_eq!(service.pending_enr_requests[&id].echo_hash, request_hash);
3493        assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(None));
3494
3495        let mut updates = service.update_stream().into_inner();
3496        service.on_enr_response(
3497            EnrResponse { request_hash, enr: remote.local_eip_868_enr.clone() },
3498            addr,
3499            id,
3500        );
3501        assert!(!service.pending_enr_requests.contains_key(&id));
3502        assert_eq!(service.on_entry(id, |entry| entry.fork_id), Some(Some(fork_id)));
3503        assert_matches::assert_matches!(
3504            updates.try_recv().unwrap(),
3505            DiscoveryUpdate::EnrForkId(found_record, found_fork_id) => {
3506                assert_eq!(found_record, record);
3507                assert_eq!(found_fork_id, fork_id);
3508            }
3509        );
3510    }
3511}