reth_dns_discovery/
resolver.rs

Help
1//! Perform DNS lookups
2
3use hickory_resolver::name_server::ConnectionProvider;
4pub use hickory_resolver::{ResolveError, TokioResolver};
5use parking_lot::RwLock;
6use std::{collections::HashMap, future::Future};
7use tracing::trace;
8
9/// A type that can lookup DNS entries
10pub trait Resolver: Send + Sync + Unpin + 'static {
11    /// Performs a textual lookup and returns the first text
12    fn lookup_txt(&self, query: &str) -> impl Future<Output = Option<String>> + Send;
13}
14
15impl<P: ConnectionProvider> Resolver for hickory_resolver::Resolver<P> {
16    async fn lookup_txt(&self, query: &str) -> Option<String> {
17        // See: [AsyncResolver::txt_lookup]
18        // > *hint* queries that end with a '.' are fully qualified names and are cheaper lookups
19        let fqn = if query.ends_with('.') { query.to_string() } else { format!("{query}.") };
20        match self.txt_lookup(fqn).await {
21            Err(err) => {
22                trace!(target: "disc::dns", %err, ?query, "dns lookup failed");
23                None
24            }
25            Ok(lookup) => {
26                let txt = lookup.into_iter().next()?;
27                let entry = txt.iter().next()?;
28                String::from_utf8(entry.to_vec()).ok()
29            }
30        }
31    }
32}
33
34/// An asynchronous DNS resolver
35///
36/// See also [`TokioResolver`]
37///
38/// ```
39/// # fn t() {
40/// use reth_dns_discovery::resolver::DnsResolver;
41/// let resolver = DnsResolver::from_system_conf().unwrap();
42/// # }
43/// ```
44///
45/// Note: This [Resolver] can send multiple lookup attempts, See also
46/// [`ResolverOpts`](hickory_resolver::config::ResolverOpts) which configures 2 attempts (1 retry)
47/// by default.
48#[derive(Clone, Debug)]
49pub struct DnsResolver(TokioResolver);
50
51// === impl DnsResolver ===
52
53impl DnsResolver {
54    /// Create a new resolver by wrapping the given [`TokioResolver`].
55    pub const fn new(resolver: TokioResolver) -> Self {
56        Self(resolver)
57    }
58
59    /// Constructs a new Tokio based Resolver with the system configuration.
60    ///
61    /// This will use `/etc/resolv.conf` on Unix OSes and the registry on Windows.
62    pub fn from_system_conf() -> Result<Self, ResolveError> {
63        TokioResolver::builder_tokio().map(|builder| Self::new(builder.build()))
64    }
65}
66
67impl Resolver for DnsResolver {
68    async fn lookup_txt(&self, query: &str) -> Option<String> {
69        Resolver::lookup_txt(&self.0, query).await
70    }
71}
72
73/// A [Resolver] that uses an in memory map to lookup entries
74#[derive(Debug, Default)]
75pub struct MapResolver(RwLock<HashMap<String, String>>);
76
77// === impl MapResolver ===
78
79impl MapResolver {
80    /// Inserts a key-value pair into the map.
81    pub fn insert(&self, k: String, v: String) -> Option<String> {
82        self.0.write().insert(k, v)
83    }
84
85    /// Returns the value corresponding to the key
86    pub fn get(&self, k: &str) -> Option<String> {
87        self.0.read().get(k).cloned()
88    }
89
90    /// Removes a key from the map, returning the value at the key if the key was previously in the
91    /// map.
92    pub fn remove(&self, k: &str) -> Option<String> {
93        self.0.write().remove(k)
94    }
95}
96
97impl Resolver for MapResolver {
98    async fn lookup_txt(&self, query: &str) -> Option<String> {
99        self.get(query)
100    }
101}
102
103/// A Resolver that always times out.
104#[cfg(test)]
105pub(crate) struct TimeoutResolver(pub(crate) std::time::Duration);
106
107#[cfg(test)]
108impl Resolver for TimeoutResolver {
109    async fn lookup_txt(&self, _query: &str) -> Option<String> {
110        tokio::time::sleep(self.0).await;
111        None
112    }
113}
reth_dns_discovery/resolver.rs

reth_dns_discovery/
resolver.rs
