reth_dns_discovery/
resolver.rs

1//! Perform DNS lookups
2
3use dashmap::DashMap;
4pub use hickory_resolver::{net::NetError, TokioResolver};
5use hickory_resolver::{proto::rr::RData, ConnectionProvider};
6use std::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.answers().iter().find_map(|r| match &r.data {
27                    RData::TXT(txt) => Some(txt),
28                    _ => None,
29                })?;
30                let entry = txt.txt_data.first()?;
31                String::from_utf8(entry.to_vec()).ok()
32            }
33        }
34    }
35}
36
37/// An asynchronous DNS resolver
38///
39/// See also [`TokioResolver`]
40///
41/// ```
42/// # fn t() {
43/// use reth_dns_discovery::resolver::DnsResolver;
44/// let resolver = DnsResolver::from_system_conf().unwrap();
45/// # }
46/// ```
47///
48/// Note: This [Resolver] can send multiple lookup attempts, See also
49/// [`ResolverOpts`](hickory_resolver::config::ResolverOpts) which configures 2 attempts (1 retry)
50/// by default.
51#[derive(Clone, Debug)]
52pub struct DnsResolver(TokioResolver);
53
54// === impl DnsResolver ===
55
56impl DnsResolver {
57    /// Create a new resolver by wrapping the given [`TokioResolver`].
58    pub const fn new(resolver: TokioResolver) -> Self {
59        Self(resolver)
60    }
61
62    /// Constructs a new Tokio based Resolver with the system configuration.
63    ///
64    /// This will use `/etc/resolv.conf` on Unix OSes and the registry on Windows.
65    pub fn from_system_conf() -> Result<Self, NetError> {
66        TokioResolver::builder_tokio()?.build().map(Self::new)
67    }
68}
69
70impl Resolver for DnsResolver {
71    async fn lookup_txt(&self, query: &str) -> Option<String> {
72        Resolver::lookup_txt(&self.0, query).await
73    }
74}
75
76/// A [Resolver] that uses an in memory map to lookup entries
77#[derive(Debug, Default)]
78pub struct MapResolver(DashMap<String, String>);
79
80// === impl MapResolver ===
81
82impl MapResolver {
83    /// Inserts a key-value pair into the map.
84    pub fn insert(&self, k: String, v: String) -> Option<String> {
85        self.0.insert(k, v)
86    }
87
88    /// Returns the value corresponding to the key
89    pub fn get(&self, k: &str) -> Option<String> {
90        self.0.get(k).map(|entry| entry.value().clone())
91    }
92
93    /// Removes a key from the map, returning the value at the key if the key was previously in the
94    /// map.
95    pub fn remove(&self, k: &str) -> Option<String> {
96        self.0.remove(k).map(|(_, v)| v)
97    }
98}
99
100impl Resolver for MapResolver {
101    async fn lookup_txt(&self, query: &str) -> Option<String> {
102        self.get(query)
103    }
104}
105
106/// A Resolver that always times out.
107#[cfg(test)]
108pub(crate) struct TimeoutResolver(pub(crate) std::time::Duration);
109
110#[cfg(test)]
111impl Resolver for TimeoutResolver {
112    async fn lookup_txt(&self, _query: &str) -> Option<String> {
113        tokio::time::sleep(self.0).await;
114        None
115    }
116}
reth_dns_discovery/resolver.rs

reth_dns_discovery/
resolver.rs
