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