Skip to main content

reth_dns_discovery/
lib.rs

1//! Implementation of [EIP-1459](https://eips.ethereum.org/EIPS/eip-1459) Node Discovery via DNS.
2//!
3//! ## Feature Flags
4//!
5//! - `serde` (default): Enable serde support
6//! - `test-utils`: Export utilities for testing
7
8#![doc(
9    html_logo_url = "https://raw.githubusercontent.com/paradigmxyz/reth/main/assets/reth-docs.png",
10    html_favicon_url = "https://avatars0.githubusercontent.com/u/97369466?s=256",
11    issue_tracker_base_url = "https://github.com/paradigmxyz/reth/issues/"
12)]
13#![cfg_attr(not(test), warn(unused_crate_dependencies))]
14#![cfg_attr(docsrs, feature(doc_cfg))]
15
16pub use crate::resolver::{DnsResolver, MapResolver, Resolver};
17use crate::{
18    query::{QueryOutcome, QueryPool, ResolveEntryResult, ResolveRootResult},
19    sync::{ResolveKind, SyncAction},
20    tree::{DnsEntry, LinkEntry},
21};
22pub use config::DnsDiscoveryConfig;
23use enr::Enr;
24pub use error::ParseDnsEntryError;
25use reth_ethereum_forks::{EnrForkIdEntry, ForkId};
26use reth_lru::LruMap;
27use reth_network_peers::NodeRecord;
28use secp256k1::SecretKey;
29use std::{
30    collections::{hash_map::Entry, HashMap, HashSet, VecDeque},
31    pin::Pin,
32    sync::Arc,
33    task::{ready, Context, Poll},
34    time::{Duration, Instant},
35};
36use sync::SyncTree;
37use tokio::{
38    sync::{
39        mpsc,
40        mpsc::{error::TrySendError, UnboundedSender},
41        oneshot,
42    },
43    task::JoinHandle,
44    time::{Interval, MissedTickBehavior},
45};
46use tokio_stream::{
47    wrappers::{ReceiverStream, UnboundedReceiverStream},
48    Stream, StreamExt,
49};
50use tracing::{debug, trace};
51
52mod config;
53mod error;
54mod query;
55pub mod resolver;
56mod sync;
57pub mod tree;
58
59/// [`DnsDiscoveryService`] front-end.
60#[derive(Clone, Debug)]
61pub struct DnsDiscoveryHandle {
62    /// Channel for sending commands to the service.
63    to_service: UnboundedSender<DnsDiscoveryCommand>,
64}
65
66// === impl DnsDiscovery ===
67
68impl DnsDiscoveryHandle {
69    /// Starts syncing the given link to a tree.
70    pub fn sync_tree(&self, link: &str) -> Result<(), ParseDnsEntryError> {
71        self.sync_tree_with_link(link.parse()?);
72        Ok(())
73    }
74
75    /// Starts syncing the given link to a tree.
76    pub fn sync_tree_with_link(&self, link: LinkEntry) {
77        let _ = self.to_service.send(DnsDiscoveryCommand::SyncTree(link));
78    }
79
80    /// Returns the receiver half of new listener channel that streams discovered [`NodeRecord`]s.
81    pub async fn node_record_stream(
82        &self,
83    ) -> Result<ReceiverStream<DnsNodeRecordUpdate>, oneshot::error::RecvError> {
84        let (tx, rx) = oneshot::channel();
85        let cmd = DnsDiscoveryCommand::NodeRecordUpdates(tx);
86        let _ = self.to_service.send(cmd);
87        rx.await
88    }
89}
90
91/// A client that discovers nodes via DNS.
92#[must_use = "Service does nothing unless polled"]
93#[expect(missing_debug_implementations)]
94pub struct DnsDiscoveryService<R: Resolver = DnsResolver> {
95    /// Copy of the sender half, so new [`DnsDiscoveryHandle`] can be created on demand.
96    command_tx: UnboundedSender<DnsDiscoveryCommand>,
97    /// Receiver half of the command channel.
98    command_rx: UnboundedReceiverStream<DnsDiscoveryCommand>,
99    /// All subscribers for resolved [`NodeRecord`]s.
100    node_record_listeners: Vec<mpsc::Sender<DnsNodeRecordUpdate>>,
101    /// All the trees that can be synced.
102    trees: HashMap<LinkEntry, SyncTree>,
103    /// All queries currently in progress
104    queries: QueryPool<R, SecretKey>,
105    /// Cached dns records
106    dns_record_cache: LruMap<String, DnsEntry<SecretKey>>,
107    /// all buffered events
108    queued_events: VecDeque<DnsDiscoveryEvent>,
109    /// The rate at which trees should be updated.
110    recheck_interval: Duration,
111    /// Wakes the service up so due trees are rechecked.
112    ///
113    /// [`SyncTree::poll`] decides that a recheck is due by comparing timestamps, which registers
114    /// no waker, and once the initial walk is done there is nothing else to poll. Without this the
115    /// spawned service would sleep until an unrelated command or query happened to wake it.
116    recheck_tick: Interval,
117    /// Links to the DNS networks to bootstrap.
118    bootstrap_dns_networks: HashSet<LinkEntry>,
119    /// Links whose initial root lookup failed, retried on every recheck tick until a tree exists
120    /// for them.
121    failed_root_lookups: HashSet<LinkEntry>,
122}
123
124// === impl DnsDiscoveryService ===
125
126impl<R: Resolver> DnsDiscoveryService<R> {
127    /// Creates a new instance of the [`DnsDiscoveryService`] using the given settings.
128    ///
129    /// ```
130    /// use reth_dns_discovery::{DnsDiscoveryService, DnsResolver};
131    /// use std::sync::Arc;
132    /// # fn t() {
133    /// let service = DnsDiscoveryService::new(
134    ///     Arc::new(DnsResolver::from_system_conf().unwrap()),
135    ///     Default::default(),
136    /// );
137    /// # }
138    /// ```
139    pub fn new(resolver: Arc<R>, config: DnsDiscoveryConfig) -> Self {
140        let DnsDiscoveryConfig {
141            lookup_timeout,
142            max_requests_per_sec,
143            recheck_interval,
144            dns_record_cache_limit,
145            bootstrap_dns_networks,
146        } = config;
147        let queries = QueryPool::new(resolver, max_requests_per_sec, lookup_timeout);
148        let (command_tx, command_rx) = mpsc::unbounded_channel();
149        Self {
150            command_tx,
151            command_rx: UnboundedReceiverStream::new(command_rx),
152            node_record_listeners: Default::default(),
153            trees: Default::default(),
154            queries,
155            dns_record_cache: LruMap::new(dns_record_cache_limit.get()),
156            queued_events: Default::default(),
157            recheck_interval,
158            recheck_tick: recheck_tick(recheck_interval),
159            bootstrap_dns_networks: bootstrap_dns_networks.unwrap_or_default(),
160            failed_root_lookups: Default::default(),
161        }
162    }
163
164    /// Spawns this services onto a new task
165    ///
166    /// Note: requires a running runtime
167    pub fn spawn(mut self) -> JoinHandle<()> {
168        let task = Box::pin(async move {
169            self.bootstrap();
170
171            while let Some(event) = self.next().await {
172                trace!(target: "disc::dns", ?event, "processed");
173            }
174        });
175        tokio::task::spawn(task as Pin<Box<dyn Future<Output = ()> + Send>>)
176    }
177
178    /// Starts discovery with all configured bootstrap links
179    pub fn bootstrap(&mut self) {
180        for link in self.bootstrap_dns_networks.clone() {
181            self.sync_tree_with_link(link);
182        }
183    }
184
185    /// Same as [`DnsDiscoveryService::new`] but also returns a new handle that's connected to the
186    /// service
187    pub fn new_pair(resolver: Arc<R>, config: DnsDiscoveryConfig) -> (Self, DnsDiscoveryHandle) {
188        let service = Self::new(resolver, config);
189        let handle = service.handle();
190        (service, handle)
191    }
192
193    /// Returns a new [`DnsDiscoveryHandle`] that can send commands to this type.
194    pub fn handle(&self) -> DnsDiscoveryHandle {
195        DnsDiscoveryHandle { to_service: self.command_tx.clone() }
196    }
197
198    /// Creates a new channel for [`NodeRecord`]s.
199    pub fn node_record_stream(&mut self) -> ReceiverStream<DnsNodeRecordUpdate> {
200        let (tx, rx) = mpsc::channel(256);
201        self.node_record_listeners.push(tx);
202        ReceiverStream::new(rx)
203    }
204
205    /// Sends  the event to all listeners.
206    ///
207    /// Remove channels that got closed.
208    fn notify(&mut self, record: DnsNodeRecordUpdate) {
209        self.node_record_listeners.retain_mut(|listener| match listener.try_send(record.clone()) {
210            Ok(()) => true,
211            Err(err) => match err {
212                TrySendError::Full(_) => true,
213                TrySendError::Closed(_) => false,
214            },
215        });
216    }
217
218    /// Starts syncing the given link to a tree.
219    pub fn sync_tree(&mut self, link: &str) -> Result<(), ParseDnsEntryError> {
220        self.sync_tree_with_link(link.parse()?);
221        Ok(())
222    }
223
224    /// Starts syncing the given link to a tree.
225    pub fn sync_tree_with_link(&mut self, link: LinkEntry) {
226        self.queries.resolve_root(link);
227    }
228
229    /// Resolves an entry
230    fn resolve_entry(&mut self, link: LinkEntry<SecretKey>, hash: String, kind: ResolveKind) {
231        if let Some(entry) = self.dns_record_cache.get(&hash).cloned() {
232            // already resolved
233            let cached = ResolveEntryResult { entry: Some(Ok(entry)), link, hash, kind };
234            self.on_resolved_entry(cached);
235            return
236        }
237        self.queries.resolve_entry(link, hash, kind)
238    }
239
240    fn on_resolved_root(&mut self, resp: ResolveRootResult<SecretKey>) {
241        match resp {
242            Ok((root, link)) => match self.trees.entry(link.clone()) {
243                Entry::Occupied(mut entry) => {
244                    entry.get_mut().update_root(root);
245                }
246                Entry::Vacant(entry) => {
247                    self.failed_root_lookups.remove(&link);
248                    entry.insert(SyncTree::new(root, link));
249                }
250            },
251            Err((err, link)) => {
252                debug!(target: "disc::dns",%err, ?link, "Failed to lookup root");
253                // an already synced tree asked for this lookup and is waiting on it, so it has to
254                // be released even though the lookup failed
255                if let Some(tree) = self.trees.get_mut(&link) {
256                    tree.root_update_failed();
257                } else {
258                    self.failed_root_lookups.insert(link);
259                }
260            }
261        }
262    }
263
264    fn on_resolved_enr(&mut self, enr: Enr<SecretKey>) {
265        if let Some(record) = convert_enr_node_record(&enr) {
266            self.notify(record);
267        }
268        self.queued_events.push_back(DnsDiscoveryEvent::Enr(enr))
269    }
270
271    fn on_resolved_entry(&mut self, resp: ResolveEntryResult<SecretKey>) {
272        let ResolveEntryResult { entry, link, hash, kind } = resp;
273
274        match entry {
275            Some(Err(err)) => {
276                debug!(target: "disc::dns",%err, domain=%link.domain, ?hash, "Failed to lookup entry")
277            }
278            None => {
279                trace!(target: "disc::dns",domain=%link.domain, ?hash, "No dns entry")
280            }
281            Some(Ok(entry)) => {
282                // cache entry
283                self.dns_record_cache.insert(hash.clone(), entry.clone());
284
285                match entry {
286                    DnsEntry::Root(root) => {
287                        debug!(target: "disc::dns",%root, domain=%link.domain, ?hash, "resolved unexpected root entry");
288                    }
289                    DnsEntry::Link(link_entry) => {
290                        if kind.is_link() {
291                            self.sync_tree_with_link(link_entry)
292                        } else {
293                            debug!(target: "disc::dns",%link_entry, domain=%link.domain, ?hash, "resolved unexpected Link entry");
294                        }
295                    }
296                    DnsEntry::Branch(branch_entry) => {
297                        if let Some(tree) = self.trees.get_mut(&link) {
298                            tree.extend_children(kind, branch_entry.children)
299                        }
300                    }
301                    DnsEntry::Node(entry) => {
302                        if kind.is_link() {
303                            debug!(target: "disc::dns",domain=%link.domain, ?hash, "resolved unexpected enr entry");
304                        } else {
305                            self.on_resolved_enr(entry.enr)
306                        }
307                    }
308                }
309            }
310        }
311    }
312
313    /// Advances the state of the DNS discovery service by polling,triggering lookups
314    pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<DnsDiscoveryEvent> {
315        loop {
316            // drain buffered events first
317            if let Some(event) = self.queued_events.pop_front() {
318                return Poll::Ready(event)
319            }
320
321            // process all incoming commands
322            while let Poll::Ready(Some(cmd)) = Pin::new(&mut self.command_rx).poll_next(cx) {
323                match cmd {
324                    DnsDiscoveryCommand::SyncTree(link) => {
325                        self.sync_tree_with_link(link);
326                    }
327                    DnsDiscoveryCommand::NodeRecordUpdates(tx) => {
328                        let _ = tx.send(self.node_record_stream());
329                    }
330                }
331            }
332
333            while let Poll::Ready(outcome) = self.queries.poll(cx) {
334                // handle query outcome
335                match outcome {
336                    QueryOutcome::Root(resp) => self.on_resolved_root(resp),
337                    QueryOutcome::Entry(resp) => self.on_resolved_entry(resp),
338                }
339            }
340
341            let mut progress = false;
342            let now = Instant::now();
343            let mut pending_resolves = Vec::new();
344            let mut pending_updates = Vec::new();
345            for tree in self.trees.values_mut() {
346                while let Some(action) = tree.poll(now, self.recheck_interval) {
347                    progress = true;
348                    match action {
349                        SyncAction::UpdateRoot => {
350                            pending_updates.push(tree.link().clone());
351                        }
352                        SyncAction::Enr(hash) => {
353                            pending_resolves.push((tree.link().clone(), hash, ResolveKind::Enr));
354                        }
355                        SyncAction::Link(hash) => {
356                            pending_resolves.push((tree.link().clone(), hash, ResolveKind::Link));
357                        }
358                    }
359                }
360            }
361
362            for (domain, hash, kind) in pending_resolves {
363                self.resolve_entry(domain, hash, kind)
364            }
365
366            for link in pending_updates {
367                self.sync_tree_with_link(link)
368            }
369
370            // Register the wake-up for the next recheck. Polled to `Pending` so the waker is
371            // always armed; the trees were already polled with a current timestamp above, so a
372            // ready tick only retries the links whose initial root lookup failed.
373            while self.recheck_tick.poll_tick(cx).is_ready() {
374                for link in self.failed_root_lookups.drain() {
375                    progress = true;
376                    self.queries.resolve_root(link);
377                }
378            }
379
380            if !progress && self.queued_events.is_empty() {
381                return Poll::Pending
382            }
383        }
384    }
385}
386
387/// Builds the interval that wakes the service to look for trees due a recheck.
388fn recheck_tick(recheck_interval: Duration) -> Interval {
389    let mut tick = tokio::time::interval(recheck_interval);
390    // a service that was starved must not then fire a burst of catch-up ticks
391    tick.set_missed_tick_behavior(MissedTickBehavior::Delay);
392    // the first tick would otherwise complete immediately and retry failed root lookups right away.
393    tick.reset();
394    tick
395}
396
397/// A Stream events, mainly used for debugging
398impl<R: Resolver> Stream for DnsDiscoveryService<R> {
399    type Item = DnsDiscoveryEvent;
400
401    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
402        Poll::Ready(Some(ready!(self.get_mut().poll(cx))))
403    }
404}
405
406/// The converted discovered [Enr] object
407#[derive(Debug, Clone, Eq, PartialEq)]
408pub struct DnsNodeRecordUpdate {
409    /// Discovered node and it's addresses
410    pub node_record: NodeRecord,
411    /// The forkid of the node, if present in the ENR
412    pub fork_id: Option<ForkId>,
413    /// Original [`Enr`].
414    pub enr: Enr<SecretKey>,
415}
416
417/// Commands sent from [`DnsDiscoveryHandle`] to [`DnsDiscoveryService`]
418enum DnsDiscoveryCommand {
419    /// Sync a tree
420    SyncTree(LinkEntry),
421    NodeRecordUpdates(oneshot::Sender<ReceiverStream<DnsNodeRecordUpdate>>),
422}
423
424/// Represents dns discovery related update events.
425#[derive(Debug, Clone)]
426pub enum DnsDiscoveryEvent {
427    /// Resolved an Enr entry via DNS.
428    Enr(Enr<SecretKey>),
429}
430
431/// Converts an [Enr] into a [`NodeRecord`]
432fn convert_enr_node_record(enr: &Enr<SecretKey>) -> Option<DnsNodeRecordUpdate> {
433    // DNS discovery yields RLPx dial targets, so records without a tcp endpoint are skipped.
434    let node_record = NodeRecord::try_from(enr).ok().filter(NodeRecord::has_rlpx_endpoint)?;
435
436    let fork_id =
437        enr.get_decodable::<EnrForkIdEntry>(b"eth").transpose().ok().flatten().map(Into::into);
438
439    Some(DnsNodeRecordUpdate { node_record, fork_id, enr: enr.clone() })
440}
441
442#[cfg(test)]
443mod tests {
444    use super::*;
445    use crate::tree::TreeRootEntry;
446    use alloy_chains::Chain;
447    use alloy_primitives::keccak256;
448    use alloy_rlp::{Decodable, Encodable};
449    use data_encoding::BASE32_NOPAD;
450    use enr::EnrKey;
451    use reth_chainspec::MAINNET;
452    use reth_ethereum_forks::{EthereumHardfork, ForkHash};
453    use secp256k1::rand::rng;
454    use std::{
455        future::poll_fn,
456        net::{IpAddr, Ipv4Addr},
457        num::NonZeroUsize,
458    };
459
460    fn entry_hash(entry_txt: &str) -> String {
461        BASE32_NOPAD.encode(&keccak256(entry_txt.as_bytes()).as_slice()[..16])
462    }
463
464    #[test]
465    fn test_convert_enr_node_record() {
466        // rig
467        let secret_key = SecretKey::new(&mut secp256k1::rand::rng());
468        let enr = Enr::builder()
469            .ip("127.0.0.1".parse().unwrap())
470            .udp4(9000)
471            .tcp4(30303)
472            .add_value(b"eth", &EnrForkIdEntry::from(MAINNET.latest_fork_id()))
473            .build(&secret_key)
474            .unwrap();
475
476        // test
477        let node_record_update = convert_enr_node_record(&enr).unwrap();
478
479        assert_eq!(node_record_update.node_record.address, "127.0.0.1".parse::<IpAddr>().unwrap());
480        assert_eq!(node_record_update.node_record.tcp_port, 30303);
481        assert_eq!(node_record_update.node_record.udp_port, 9000);
482        assert_eq!(node_record_update.fork_id, Some(MAINNET.latest_fork_id()));
483        assert_eq!(node_record_update.enr, enr);
484    }
485
486    #[test]
487    fn test_decode_and_convert_enr_node_record() {
488        // rig
489
490        let secret_key = SecretKey::new(&mut secp256k1::rand::rng());
491        let enr = Enr::builder()
492            .ip("127.0.0.1".parse().unwrap())
493            .udp4(9000)
494            .tcp4(30303)
495            .add_value(b"eth", &EnrForkIdEntry::from(MAINNET.latest_fork_id()))
496            .add_value(b"opstack", &ForkId { hash: ForkHash(rand::random()), next: rand::random() })
497            .build(&secret_key)
498            .unwrap();
499
500        let mut encoded_enr = vec![];
501        enr.encode(&mut encoded_enr);
502
503        // test
504        let decoded_enr = Enr::decode(&mut &encoded_enr[..]).unwrap();
505
506        let node_record_update = convert_enr_node_record(&decoded_enr).unwrap();
507
508        assert_eq!(node_record_update.node_record.address, "127.0.0.1".parse::<IpAddr>().unwrap());
509        assert_eq!(node_record_update.node_record.tcp_port, 30303);
510        assert_eq!(node_record_update.node_record.udp_port, 9000);
511        assert_eq!(node_record_update.fork_id, Some(MAINNET.latest_fork_id()));
512        assert_eq!(node_record_update.enr, enr);
513    }
514
515    #[tokio::test]
516    async fn test_start_root_sync() {
517        reth_tracing::init_test_tracing();
518
519        let secret_key = SecretKey::new(&mut rng());
520        let resolver = MapResolver::default();
521        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
522        let mut root: TreeRootEntry = s.parse().unwrap();
523        root.sign(&secret_key).unwrap();
524
525        let link =
526            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
527        resolver.insert(link.domain.clone(), root.to_string());
528
529        let mut service = DnsDiscoveryService::new(Arc::new(resolver), Default::default());
530
531        service.sync_tree_with_link(link.clone());
532
533        poll_fn(|cx| {
534            let _ = service.poll(cx);
535            Poll::Ready(())
536        })
537        .await;
538
539        let tree = service.trees.get(&link).unwrap();
540        assert_eq!(tree.root().clone(), root);
541    }
542
543    #[tokio::test(flavor = "multi_thread")]
544    async fn test_get_node() {
545        reth_tracing::init_test_tracing();
546
547        let secret_key = SecretKey::new(&mut rng());
548        let resolver = MapResolver::default();
549        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
550        let mut root: TreeRootEntry = s.parse().unwrap();
551
552        let link =
553            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
554
555        let mut builder = Enr::builder();
556        let fork_id = MAINNET.hardfork_fork_id(EthereumHardfork::Frontier).unwrap();
557        builder
558            .ip4(Ipv4Addr::LOCALHOST)
559            .udp4(30303)
560            .tcp4(30303)
561            .add_value(b"eth", &EnrForkIdEntry::from(fork_id));
562        let enr = builder.build(&secret_key).unwrap();
563        let enr_txt = enr.to_base64();
564
565        root.enr_root = entry_hash(&enr_txt);
566        root.sign(&secret_key).unwrap();
567
568        resolver.insert(link.domain.clone(), root.to_string());
569        resolver.insert(format!("{}.{}", root.enr_root.clone(), link.domain), enr_txt);
570
571        let mut service = DnsDiscoveryService::new(Arc::new(resolver), Default::default());
572
573        let mut node_records = service.node_record_stream();
574
575        let task = tokio::task::spawn(async move {
576            let record = node_records.next().await.unwrap();
577            assert_eq!(record.fork_id, Some(fork_id));
578        });
579
580        service.sync_tree_with_link(link.clone());
581
582        let event = poll_fn(|cx| service.poll(cx)).await;
583
584        match event {
585            DnsDiscoveryEvent::Enr(discovered) => {
586                assert_eq!(discovered, enr);
587            }
588        }
589
590        poll_fn(|cx| {
591            assert!(service.poll(cx).is_pending());
592            Poll::Ready(())
593        })
594        .await;
595
596        task.await.unwrap();
597    }
598
599    #[tokio::test]
600    async fn test_recheck_tree() {
601        reth_tracing::init_test_tracing();
602
603        let config = DnsDiscoveryConfig {
604            recheck_interval: Duration::from_millis(750),
605            ..Default::default()
606        };
607
608        let secret_key = SecretKey::new(&mut rng());
609        let resolver = Arc::new(MapResolver::default());
610        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
611        let mut root: TreeRootEntry = s.parse().unwrap();
612        root.sign(&secret_key).unwrap();
613
614        let link =
615            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
616        resolver.insert(link.domain.clone(), root.to_string());
617
618        let mut service = DnsDiscoveryService::new(Arc::clone(&resolver), config.clone());
619
620        service.sync_tree_with_link(link.clone());
621
622        poll_fn(|cx| {
623            assert!(service.poll(cx).is_pending());
624            Poll::Ready(())
625        })
626        .await;
627
628        // await recheck timeout
629        tokio::time::sleep(config.recheck_interval).await;
630
631        let mut new_root = root.clone();
632        new_root.sequence_number = new_root.sequence_number.saturating_add(1);
633
634        let enr = Enr::empty(&secret_key).unwrap();
635        let enr_txt = enr.to_base64();
636        new_root.enr_root = entry_hash(&enr_txt);
637        new_root.sign(&secret_key).unwrap();
638        resolver.insert(link.domain.clone(), new_root.to_string());
639        resolver.insert(format!("{}.{}", new_root.enr_root.clone(), link.domain), enr_txt);
640
641        let event = poll_fn(|cx| service.poll(cx)).await;
642
643        match event {
644            DnsDiscoveryEvent::Enr(discovered) => {
645                assert_eq!(discovered, enr);
646            }
647        }
648
649        poll_fn(|cx| {
650            assert!(service.poll(cx).is_pending());
651            Poll::Ready(())
652        })
653        .await;
654    }
655
656    #[tokio::test]
657    async fn spawned_service_rechecks_without_outside_wakeups() {
658        reth_tracing::init_test_tracing();
659
660        let config = DnsDiscoveryConfig {
661            recheck_interval: Duration::from_millis(100),
662            max_requests_per_sec: NonZeroUsize::new(50).unwrap(),
663            ..Default::default()
664        };
665
666        let secret_key = SecretKey::new(&mut rng());
667        let resolver = Arc::new(MapResolver::default());
668        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
669        let mut root: TreeRootEntry = s.parse().unwrap();
670        root.sign(&secret_key).unwrap();
671
672        let link =
673            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
674        resolver.insert(link.domain.clone(), root.to_string());
675
676        let mut service = DnsDiscoveryService::new(Arc::clone(&resolver), config.clone());
677        let mut node_records = service.node_record_stream();
678        service.sync_tree_with_link(link.clone());
679
680        // drive it the way the production path does: nothing polls the service except its own
681        // event stream, so only a wake-up the service arms itself can start a recheck
682        let handle = service.spawn();
683
684        // let the initial walk finish and the service go idle
685        tokio::time::sleep(config.recheck_interval * 3).await;
686
687        // publish a change that only a self-scheduled recheck can pick up
688        let mut new_root = root.clone();
689        new_root.sequence_number = new_root.sequence_number.saturating_add(1);
690        // needs an address, an ENR without one yields no node record
691        let enr = Enr::builder()
692            .ip4(Ipv4Addr::LOCALHOST)
693            .udp4(30303)
694            .tcp4(30303)
695            .build(&secret_key)
696            .unwrap();
697        let enr_txt = enr.to_base64();
698        new_root.enr_root = entry_hash(&enr_txt);
699        new_root.sign(&secret_key).unwrap();
700        resolver.insert(link.domain.clone(), new_root.to_string());
701        resolver.insert(format!("{}.{}", new_root.enr_root.clone(), link.domain), enr_txt);
702
703        let update = tokio::time::timeout(Duration::from_secs(10), node_records.next())
704            .await
705            .expect("service never woke up to recheck the tree")
706            .expect("record stream closed");
707        assert_eq!(update.enr, enr);
708
709        handle.abort();
710    }
711
712    #[tokio::test]
713    async fn test_retry_failed_initial_root_lookup() {
714        reth_tracing::init_test_tracing();
715
716        let config = DnsDiscoveryConfig {
717            recheck_interval: Duration::from_millis(100),
718            max_requests_per_sec: NonZeroUsize::new(50).unwrap(),
719            ..Default::default()
720        };
721
722        let secret_key = SecretKey::new(&mut rng());
723        let resolver = Arc::new(MapResolver::default());
724        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
725        let mut root: TreeRootEntry = s.parse().unwrap();
726        root.sign(&secret_key).unwrap();
727
728        let link =
729            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
730
731        let mut service = DnsDiscoveryService::new(Arc::clone(&resolver), config.clone());
732
733        // the root is not published yet, so the initial lookup fails.
734        service.sync_tree_with_link(link.clone());
735        poll_fn(|cx| {
736            assert!(service.poll(cx).is_pending());
737            Poll::Ready(())
738        })
739        .await;
740        assert!(service.trees.is_empty());
741
742        resolver.insert(link.domain.clone(), root.to_string());
743        tokio::time::sleep(config.recheck_interval).await;
744
745        poll_fn(|cx| {
746            let _ = service.poll(cx);
747            Poll::Ready(())
748        })
749        .await;
750
751        let tree = service.trees.get(&link).unwrap();
752        assert_eq!(tree.root().clone(), root);
753    }
754
755    #[tokio::test]
756    async fn test_hash_mismatch_is_not_cached_and_does_not_poison_same_hash() {
757        let secret_key = SecretKey::new(&mut rng());
758        let resolver = MapResolver::default();
759
760        let invalid_entry = "enrtree-branch:AAAAAAAAAAAAAAAAAAAA".to_string();
761        let valid_entry = "enrtree-branch:YNEGZIWHOM7TOOSUATAPTM".to_string();
762
763        let hash = entry_hash(&valid_entry);
764
765        let bad_link =
766            LinkEntry { domain: "bad.example.org".to_string(), pubkey: secret_key.public() };
767        let good_link =
768            LinkEntry { domain: "good.example.org".to_string(), pubkey: secret_key.public() };
769
770        resolver.insert(format!("{}.{}", hash, bad_link.domain), invalid_entry);
771        resolver.insert(format!("{}.{}", hash, good_link.domain), valid_entry.clone());
772
773        let mut service = DnsDiscoveryService::new(Arc::new(resolver), Default::default());
774
775        service.resolve_entry(bad_link, hash.clone(), ResolveKind::Enr);
776        poll_fn(|cx| {
777            let _ = service.poll(cx);
778            Poll::Ready(())
779        })
780        .await;
781
782        assert!(service.dns_record_cache.get(&hash).is_none());
783
784        service.resolve_entry(good_link, hash.clone(), ResolveKind::Enr);
785        poll_fn(|cx| {
786            let _ = service.poll(cx);
787            Poll::Ready(())
788        })
789        .await;
790
791        let cached = service.dns_record_cache.get(&hash).cloned();
792        assert_eq!(cached.map(|entry| entry.to_string()), Some(valid_entry));
793    }
794
795    #[tokio::test]
796    #[ignore]
797    async fn test_dns_resolver() {
798        reth_tracing::init_test_tracing();
799
800        let mut service = DnsDiscoveryService::new(
801            Arc::new(DnsResolver::from_system_conf().unwrap()),
802            Default::default(),
803        );
804
805        service.sync_tree(&Chain::mainnet().public_dns_network_protocol().unwrap()).unwrap();
806
807        while let Some(event) = service.next().await {
808            match event {
809                DnsDiscoveryEvent::Enr(enr) => {
810                    println!("discovered enr {}", enr.to_base64());
811                }
812            }
813        }
814    }
815}