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