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_network_peers::NodeRecord;
27use schnellru::{ByLength, LruMap};
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}
120
121// === impl DnsDiscoveryService ===
122
123impl<R: Resolver> DnsDiscoveryService<R> {
124    /// Creates a new instance of the [`DnsDiscoveryService`] using the given settings.
125    ///
126    /// ```
127    /// use reth_dns_discovery::{DnsDiscoveryService, DnsResolver};
128    /// use std::sync::Arc;
129    /// # fn t() {
130    /// let service = DnsDiscoveryService::new(
131    ///     Arc::new(DnsResolver::from_system_conf().unwrap()),
132    ///     Default::default(),
133    /// );
134    /// # }
135    /// ```
136    pub fn new(resolver: Arc<R>, config: DnsDiscoveryConfig) -> Self {
137        let DnsDiscoveryConfig {
138            lookup_timeout,
139            max_requests_per_sec,
140            recheck_interval,
141            dns_record_cache_limit,
142            bootstrap_dns_networks,
143        } = config;
144        let queries = QueryPool::new(resolver, max_requests_per_sec, lookup_timeout);
145        let (command_tx, command_rx) = mpsc::unbounded_channel();
146        Self {
147            command_tx,
148            command_rx: UnboundedReceiverStream::new(command_rx),
149            node_record_listeners: Default::default(),
150            trees: Default::default(),
151            queries,
152            dns_record_cache: LruMap::new(ByLength::new(dns_record_cache_limit.get())),
153            queued_events: Default::default(),
154            recheck_interval,
155            recheck_tick: recheck_tick(recheck_interval),
156            bootstrap_dns_networks: bootstrap_dns_networks.unwrap_or_default(),
157        }
158    }
159
160    /// Spawns this services onto a new task
161    ///
162    /// Note: requires a running runtime
163    pub fn spawn(mut self) -> JoinHandle<()> {
164        tokio::task::spawn(async move {
165            self.bootstrap();
166
167            while let Some(event) = self.next().await {
168                trace!(target: "disc::dns", ?event, "processed");
169            }
170        })
171    }
172
173    /// Starts discovery with all configured bootstrap links
174    pub fn bootstrap(&mut self) {
175        for link in self.bootstrap_dns_networks.clone() {
176            self.sync_tree_with_link(link);
177        }
178    }
179
180    /// Same as [`DnsDiscoveryService::new`] but also returns a new handle that's connected to the
181    /// service
182    pub fn new_pair(resolver: Arc<R>, config: DnsDiscoveryConfig) -> (Self, DnsDiscoveryHandle) {
183        let service = Self::new(resolver, config);
184        let handle = service.handle();
185        (service, handle)
186    }
187
188    /// Returns a new [`DnsDiscoveryHandle`] that can send commands to this type.
189    pub fn handle(&self) -> DnsDiscoveryHandle {
190        DnsDiscoveryHandle { to_service: self.command_tx.clone() }
191    }
192
193    /// Creates a new channel for [`NodeRecord`]s.
194    pub fn node_record_stream(&mut self) -> ReceiverStream<DnsNodeRecordUpdate> {
195        let (tx, rx) = mpsc::channel(256);
196        self.node_record_listeners.push(tx);
197        ReceiverStream::new(rx)
198    }
199
200    /// Sends  the event to all listeners.
201    ///
202    /// Remove channels that got closed.
203    fn notify(&mut self, record: DnsNodeRecordUpdate) {
204        self.node_record_listeners.retain_mut(|listener| match listener.try_send(record.clone()) {
205            Ok(()) => true,
206            Err(err) => match err {
207                TrySendError::Full(_) => true,
208                TrySendError::Closed(_) => false,
209            },
210        });
211    }
212
213    /// Starts syncing the given link to a tree.
214    pub fn sync_tree(&mut self, link: &str) -> Result<(), ParseDnsEntryError> {
215        self.sync_tree_with_link(link.parse()?);
216        Ok(())
217    }
218
219    /// Starts syncing the given link to a tree.
220    pub fn sync_tree_with_link(&mut self, link: LinkEntry) {
221        self.queries.resolve_root(link);
222    }
223
224    /// Resolves an entry
225    fn resolve_entry(&mut self, link: LinkEntry<SecretKey>, hash: String, kind: ResolveKind) {
226        if let Some(entry) = self.dns_record_cache.get(&hash).cloned() {
227            // already resolved
228            let cached = ResolveEntryResult { entry: Some(Ok(entry)), link, hash, kind };
229            self.on_resolved_entry(cached);
230            return
231        }
232        self.queries.resolve_entry(link, hash, kind)
233    }
234
235    fn on_resolved_root(&mut self, resp: ResolveRootResult<SecretKey>) {
236        match resp {
237            Ok((root, link)) => match self.trees.entry(link.clone()) {
238                Entry::Occupied(mut entry) => {
239                    entry.get_mut().update_root(root);
240                }
241                Entry::Vacant(entry) => {
242                    entry.insert(SyncTree::new(root, link));
243                }
244            },
245            Err((err, link)) => {
246                debug!(target: "disc::dns",%err, ?link, "Failed to lookup root");
247                // an already synced tree asked for this lookup and is waiting on it, so it has to
248                // be released even though the lookup failed
249                if let Some(tree) = self.trees.get_mut(&link) {
250                    tree.root_update_failed();
251                }
252            }
253        }
254    }
255
256    fn on_resolved_enr(&mut self, enr: Enr<SecretKey>) {
257        if let Some(record) = convert_enr_node_record(&enr) {
258            self.notify(record);
259        }
260        self.queued_events.push_back(DnsDiscoveryEvent::Enr(enr))
261    }
262
263    fn on_resolved_entry(&mut self, resp: ResolveEntryResult<SecretKey>) {
264        let ResolveEntryResult { entry, link, hash, kind } = resp;
265
266        match entry {
267            Some(Err(err)) => {
268                debug!(target: "disc::dns",%err, domain=%link.domain, ?hash, "Failed to lookup entry")
269            }
270            None => {
271                trace!(target: "disc::dns",domain=%link.domain, ?hash, "No dns entry")
272            }
273            Some(Ok(entry)) => {
274                // cache entry
275                self.dns_record_cache.insert(hash.clone(), entry.clone());
276
277                match entry {
278                    DnsEntry::Root(root) => {
279                        debug!(target: "disc::dns",%root, domain=%link.domain, ?hash, "resolved unexpected root entry");
280                    }
281                    DnsEntry::Link(link_entry) => {
282                        if kind.is_link() {
283                            self.sync_tree_with_link(link_entry)
284                        } else {
285                            debug!(target: "disc::dns",%link_entry, domain=%link.domain, ?hash, "resolved unexpected Link entry");
286                        }
287                    }
288                    DnsEntry::Branch(branch_entry) => {
289                        if let Some(tree) = self.trees.get_mut(&link) {
290                            tree.extend_children(kind, branch_entry.children)
291                        }
292                    }
293                    DnsEntry::Node(entry) => {
294                        if kind.is_link() {
295                            debug!(target: "disc::dns",domain=%link.domain, ?hash, "resolved unexpected enr entry");
296                        } else {
297                            self.on_resolved_enr(entry.enr)
298                        }
299                    }
300                }
301            }
302        }
303    }
304
305    /// Advances the state of the DNS discovery service by polling,triggering lookups
306    pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<DnsDiscoveryEvent> {
307        loop {
308            // drain buffered events first
309            if let Some(event) = self.queued_events.pop_front() {
310                return Poll::Ready(event)
311            }
312
313            // process all incoming commands
314            while let Poll::Ready(Some(cmd)) = Pin::new(&mut self.command_rx).poll_next(cx) {
315                match cmd {
316                    DnsDiscoveryCommand::SyncTree(link) => {
317                        self.sync_tree_with_link(link);
318                    }
319                    DnsDiscoveryCommand::NodeRecordUpdates(tx) => {
320                        let _ = tx.send(self.node_record_stream());
321                    }
322                }
323            }
324
325            while let Poll::Ready(outcome) = self.queries.poll(cx) {
326                // handle query outcome
327                match outcome {
328                    QueryOutcome::Root(resp) => self.on_resolved_root(resp),
329                    QueryOutcome::Entry(resp) => self.on_resolved_entry(resp),
330                }
331            }
332
333            let mut progress = false;
334            let now = Instant::now();
335            let mut pending_resolves = Vec::new();
336            let mut pending_updates = Vec::new();
337            for tree in self.trees.values_mut() {
338                while let Some(action) = tree.poll(now, self.recheck_interval) {
339                    progress = true;
340                    match action {
341                        SyncAction::UpdateRoot => {
342                            pending_updates.push(tree.link().clone());
343                        }
344                        SyncAction::Enr(hash) => {
345                            pending_resolves.push((tree.link().clone(), hash, ResolveKind::Enr));
346                        }
347                        SyncAction::Link(hash) => {
348                            pending_resolves.push((tree.link().clone(), hash, ResolveKind::Link));
349                        }
350                    }
351                }
352            }
353
354            for (domain, hash, kind) in pending_resolves {
355                self.resolve_entry(domain, hash, kind)
356            }
357
358            for link in pending_updates {
359                self.sync_tree_with_link(link)
360            }
361
362            // Register the wake-up for the next recheck. Polled to `Pending` so the waker is
363            // always armed; a ready tick needs no extra work because the trees were already polled
364            // with a current timestamp above.
365            while self.recheck_tick.poll_tick(cx).is_ready() {}
366
367            if !progress && self.queued_events.is_empty() {
368                return Poll::Pending
369            }
370        }
371    }
372}
373
374/// Builds the interval that wakes the service to look for trees due a recheck.
375fn recheck_tick(recheck_interval: Duration) -> Interval {
376    let mut tick = tokio::time::interval(recheck_interval);
377    // a service that was starved must not then fire a burst of catch-up ticks
378    tick.set_missed_tick_behavior(MissedTickBehavior::Delay);
379    tick
380}
381
382/// A Stream events, mainly used for debugging
383impl<R: Resolver> Stream for DnsDiscoveryService<R> {
384    type Item = DnsDiscoveryEvent;
385
386    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
387        Poll::Ready(Some(ready!(self.get_mut().poll(cx))))
388    }
389}
390
391/// The converted discovered [Enr] object
392#[derive(Debug, Clone, Eq, PartialEq)]
393pub struct DnsNodeRecordUpdate {
394    /// Discovered node and it's addresses
395    pub node_record: NodeRecord,
396    /// The forkid of the node, if present in the ENR
397    pub fork_id: Option<ForkId>,
398    /// Original [`Enr`].
399    pub enr: Enr<SecretKey>,
400}
401
402/// Commands sent from [`DnsDiscoveryHandle`] to [`DnsDiscoveryService`]
403enum DnsDiscoveryCommand {
404    /// Sync a tree
405    SyncTree(LinkEntry),
406    NodeRecordUpdates(oneshot::Sender<ReceiverStream<DnsNodeRecordUpdate>>),
407}
408
409/// Represents dns discovery related update events.
410#[derive(Debug, Clone)]
411pub enum DnsDiscoveryEvent {
412    /// Resolved an Enr entry via DNS.
413    Enr(Enr<SecretKey>),
414}
415
416/// Converts an [Enr] into a [`NodeRecord`]
417fn convert_enr_node_record(enr: &Enr<SecretKey>) -> Option<DnsNodeRecordUpdate> {
418    // DNS discovery yields RLPx dial targets, so records without a tcp endpoint are skipped.
419    let node_record = NodeRecord::try_from(enr).ok().filter(NodeRecord::has_rlpx_endpoint)?;
420
421    let fork_id =
422        enr.get_decodable::<EnrForkIdEntry>(b"eth").transpose().ok().flatten().map(Into::into);
423
424    Some(DnsNodeRecordUpdate { node_record, fork_id, enr: enr.clone() })
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use crate::tree::TreeRootEntry;
431    use alloy_chains::Chain;
432    use alloy_primitives::keccak256;
433    use alloy_rlp::{Decodable, Encodable};
434    use data_encoding::BASE32_NOPAD;
435    use enr::EnrKey;
436    use reth_chainspec::MAINNET;
437    use reth_ethereum_forks::{EthereumHardfork, ForkHash};
438    use secp256k1::rand::thread_rng;
439    use std::{
440        future::poll_fn,
441        net::{IpAddr, Ipv4Addr},
442        num::NonZeroUsize,
443    };
444
445    fn entry_hash(entry_txt: &str) -> String {
446        BASE32_NOPAD.encode(&keccak256(entry_txt.as_bytes()).as_slice()[..16])
447    }
448
449    #[test]
450    fn test_convert_enr_node_record() {
451        // rig
452        let secret_key = SecretKey::new(&mut secp256k1::rand::thread_rng());
453        let enr = Enr::builder()
454            .ip("127.0.0.1".parse().unwrap())
455            .udp4(9000)
456            .tcp4(30303)
457            .add_value(b"eth", &EnrForkIdEntry::from(MAINNET.latest_fork_id()))
458            .build(&secret_key)
459            .unwrap();
460
461        // test
462        let node_record_update = convert_enr_node_record(&enr).unwrap();
463
464        assert_eq!(node_record_update.node_record.address, "127.0.0.1".parse::<IpAddr>().unwrap());
465        assert_eq!(node_record_update.node_record.tcp_port, 30303);
466        assert_eq!(node_record_update.node_record.udp_port, 9000);
467        assert_eq!(node_record_update.fork_id, Some(MAINNET.latest_fork_id()));
468        assert_eq!(node_record_update.enr, enr);
469    }
470
471    #[test]
472    fn test_decode_and_convert_enr_node_record() {
473        // rig
474
475        let secret_key = SecretKey::new(&mut secp256k1::rand::thread_rng());
476        let enr = Enr::builder()
477            .ip("127.0.0.1".parse().unwrap())
478            .udp4(9000)
479            .tcp4(30303)
480            .add_value(b"eth", &EnrForkIdEntry::from(MAINNET.latest_fork_id()))
481            .add_value(b"opstack", &ForkId { hash: ForkHash(rand::random()), next: rand::random() })
482            .build(&secret_key)
483            .unwrap();
484
485        let mut encoded_enr = vec![];
486        enr.encode(&mut encoded_enr);
487
488        // test
489        let decoded_enr = Enr::decode(&mut &encoded_enr[..]).unwrap();
490
491        let node_record_update = convert_enr_node_record(&decoded_enr).unwrap();
492
493        assert_eq!(node_record_update.node_record.address, "127.0.0.1".parse::<IpAddr>().unwrap());
494        assert_eq!(node_record_update.node_record.tcp_port, 30303);
495        assert_eq!(node_record_update.node_record.udp_port, 9000);
496        assert_eq!(node_record_update.fork_id, Some(MAINNET.latest_fork_id()));
497        assert_eq!(node_record_update.enr, enr);
498    }
499
500    #[tokio::test]
501    async fn test_start_root_sync() {
502        reth_tracing::init_test_tracing();
503
504        let secret_key = SecretKey::new(&mut thread_rng());
505        let resolver = MapResolver::default();
506        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
507        let mut root: TreeRootEntry = s.parse().unwrap();
508        root.sign(&secret_key).unwrap();
509
510        let link =
511            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
512        resolver.insert(link.domain.clone(), root.to_string());
513
514        let mut service = DnsDiscoveryService::new(Arc::new(resolver), Default::default());
515
516        service.sync_tree_with_link(link.clone());
517
518        poll_fn(|cx| {
519            let _ = service.poll(cx);
520            Poll::Ready(())
521        })
522        .await;
523
524        let tree = service.trees.get(&link).unwrap();
525        assert_eq!(tree.root().clone(), root);
526    }
527
528    #[tokio::test(flavor = "multi_thread")]
529    async fn test_get_node() {
530        reth_tracing::init_test_tracing();
531
532        let secret_key = SecretKey::new(&mut thread_rng());
533        let resolver = MapResolver::default();
534        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
535        let mut root: TreeRootEntry = s.parse().unwrap();
536
537        let link =
538            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
539
540        let mut builder = Enr::builder();
541        let fork_id = MAINNET.hardfork_fork_id(EthereumHardfork::Frontier).unwrap();
542        builder
543            .ip4(Ipv4Addr::LOCALHOST)
544            .udp4(30303)
545            .tcp4(30303)
546            .add_value(b"eth", &EnrForkIdEntry::from(fork_id));
547        let enr = builder.build(&secret_key).unwrap();
548        let enr_txt = enr.to_base64();
549
550        root.enr_root = entry_hash(&enr_txt);
551        root.sign(&secret_key).unwrap();
552
553        resolver.insert(link.domain.clone(), root.to_string());
554        resolver.insert(format!("{}.{}", root.enr_root.clone(), link.domain), enr_txt);
555
556        let mut service = DnsDiscoveryService::new(Arc::new(resolver), Default::default());
557
558        let mut node_records = service.node_record_stream();
559
560        let task = tokio::task::spawn(async move {
561            let record = node_records.next().await.unwrap();
562            assert_eq!(record.fork_id, Some(fork_id));
563        });
564
565        service.sync_tree_with_link(link.clone());
566
567        let event = poll_fn(|cx| service.poll(cx)).await;
568
569        match event {
570            DnsDiscoveryEvent::Enr(discovered) => {
571                assert_eq!(discovered, enr);
572            }
573        }
574
575        poll_fn(|cx| {
576            assert!(service.poll(cx).is_pending());
577            Poll::Ready(())
578        })
579        .await;
580
581        task.await.unwrap();
582    }
583
584    #[tokio::test]
585    async fn test_recheck_tree() {
586        reth_tracing::init_test_tracing();
587
588        let config = DnsDiscoveryConfig {
589            recheck_interval: Duration::from_millis(750),
590            ..Default::default()
591        };
592
593        let secret_key = SecretKey::new(&mut thread_rng());
594        let resolver = Arc::new(MapResolver::default());
595        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
596        let mut root: TreeRootEntry = s.parse().unwrap();
597        root.sign(&secret_key).unwrap();
598
599        let link =
600            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
601        resolver.insert(link.domain.clone(), root.to_string());
602
603        let mut service = DnsDiscoveryService::new(Arc::clone(&resolver), config.clone());
604
605        service.sync_tree_with_link(link.clone());
606
607        poll_fn(|cx| {
608            assert!(service.poll(cx).is_pending());
609            Poll::Ready(())
610        })
611        .await;
612
613        // await recheck timeout
614        tokio::time::sleep(config.recheck_interval).await;
615
616        let mut new_root = root.clone();
617        new_root.sequence_number = new_root.sequence_number.saturating_add(1);
618
619        let enr = Enr::empty(&secret_key).unwrap();
620        let enr_txt = enr.to_base64();
621        new_root.enr_root = entry_hash(&enr_txt);
622        new_root.sign(&secret_key).unwrap();
623        resolver.insert(link.domain.clone(), new_root.to_string());
624        resolver.insert(format!("{}.{}", new_root.enr_root.clone(), link.domain), enr_txt);
625
626        let event = poll_fn(|cx| service.poll(cx)).await;
627
628        match event {
629            DnsDiscoveryEvent::Enr(discovered) => {
630                assert_eq!(discovered, enr);
631            }
632        }
633
634        poll_fn(|cx| {
635            assert!(service.poll(cx).is_pending());
636            Poll::Ready(())
637        })
638        .await;
639    }
640
641    #[tokio::test]
642    async fn spawned_service_rechecks_without_outside_wakeups() {
643        reth_tracing::init_test_tracing();
644
645        let config = DnsDiscoveryConfig {
646            recheck_interval: Duration::from_millis(100),
647            max_requests_per_sec: NonZeroUsize::new(50).unwrap(),
648            ..Default::default()
649        };
650
651        let secret_key = SecretKey::new(&mut thread_rng());
652        let resolver = Arc::new(MapResolver::default());
653        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
654        let mut root: TreeRootEntry = s.parse().unwrap();
655        root.sign(&secret_key).unwrap();
656
657        let link =
658            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
659        resolver.insert(link.domain.clone(), root.to_string());
660
661        let mut service = DnsDiscoveryService::new(Arc::clone(&resolver), config.clone());
662        let mut node_records = service.node_record_stream();
663        service.sync_tree_with_link(link.clone());
664
665        // drive it the way the production path does: nothing polls the service except its own
666        // event stream, so only a wake-up the service arms itself can start a recheck
667        let handle = service.spawn();
668
669        // let the initial walk finish and the service go idle
670        tokio::time::sleep(config.recheck_interval * 3).await;
671
672        // publish a change that only a self-scheduled recheck can pick up
673        let mut new_root = root.clone();
674        new_root.sequence_number = new_root.sequence_number.saturating_add(1);
675        // needs an address, an ENR without one yields no node record
676        let enr = Enr::builder()
677            .ip4(Ipv4Addr::LOCALHOST)
678            .udp4(30303)
679            .tcp4(30303)
680            .build(&secret_key)
681            .unwrap();
682        let enr_txt = enr.to_base64();
683        new_root.enr_root = entry_hash(&enr_txt);
684        new_root.sign(&secret_key).unwrap();
685        resolver.insert(link.domain.clone(), new_root.to_string());
686        resolver.insert(format!("{}.{}", new_root.enr_root.clone(), link.domain), enr_txt);
687
688        let update = tokio::time::timeout(Duration::from_secs(10), node_records.next())
689            .await
690            .expect("service never woke up to recheck the tree")
691            .expect("record stream closed");
692        assert_eq!(update.enr, enr);
693
694        handle.abort();
695    }
696
697    #[tokio::test]
698    async fn test_hash_mismatch_is_not_cached_and_does_not_poison_same_hash() {
699        let secret_key = SecretKey::new(&mut thread_rng());
700        let resolver = MapResolver::default();
701
702        let invalid_entry = "enrtree-branch:AAAAAAAAAAAAAAAAAAAA".to_string();
703        let valid_entry = "enrtree-branch:YNEGZIWHOM7TOOSUATAPTM".to_string();
704
705        let hash = entry_hash(&valid_entry);
706
707        let bad_link =
708            LinkEntry { domain: "bad.example.org".to_string(), pubkey: secret_key.public() };
709        let good_link =
710            LinkEntry { domain: "good.example.org".to_string(), pubkey: secret_key.public() };
711
712        resolver.insert(format!("{}.{}", hash, bad_link.domain), invalid_entry);
713        resolver.insert(format!("{}.{}", hash, good_link.domain), valid_entry.clone());
714
715        let mut service = DnsDiscoveryService::new(Arc::new(resolver), Default::default());
716
717        service.resolve_entry(bad_link, hash.clone(), ResolveKind::Enr);
718        poll_fn(|cx| {
719            let _ = service.poll(cx);
720            Poll::Ready(())
721        })
722        .await;
723
724        assert!(service.dns_record_cache.get(&hash).is_none());
725
726        service.resolve_entry(good_link, hash.clone(), ResolveKind::Enr);
727        poll_fn(|cx| {
728            let _ = service.poll(cx);
729            Poll::Ready(())
730        })
731        .await;
732
733        let cached = service.dns_record_cache.get(&hash).cloned();
734        assert_eq!(cached.map(|entry| entry.to_string()), Some(valid_entry));
735    }
736
737    #[tokio::test]
738    #[ignore]
739    async fn test_dns_resolver() {
740        reth_tracing::init_test_tracing();
741
742        let mut service = DnsDiscoveryService::new(
743            Arc::new(DnsResolver::from_system_conf().unwrap()),
744            Default::default(),
745        );
746
747        service.sync_tree(&Chain::mainnet().public_dns_network_protocol().unwrap()).unwrap();
748
749        while let Some(event) = service.next().await {
750            match event {
751                DnsDiscoveryEvent::Enr(enr) => {
752                    println!("discovered enr {}", enr.to_base64());
753                }
754            }
755        }
756    }
757}
reth_dns_discovery/lib.rs

reth_dns_discovery/
lib.rs
