reth_dns_discovery/
resolver.rs

1//! Perform DNS lookups
2
3use crate::tree::has_entry_prefix;
4use dashmap::DashMap;
5pub use hickory_resolver::{net::NetError, TokioResolver};
6use hickory_resolver::{
7    proto::rr::{rdata::TXT, RData, Record},
8    ConnectionProvider,
9};
10use std::future::Future;
11use tracing::trace;
12
13/// A type that can lookup DNS entries
14pub trait Resolver: Send + Sync + Unpin + 'static {
15    /// Performs a textual lookup and returns the first text
16    fn lookup_txt(&self, query: &str) -> impl Future<Output = Option<String>> + Send;
17}
18
19impl<P: ConnectionProvider> Resolver for hickory_resolver::Resolver<P> {
20    async fn lookup_txt(&self, query: &str) -> Option<String> {
21        // See: [AsyncResolver::txt_lookup]
22        // > *hint* queries that end with a '.' are fully qualified names and are cheaper lookups
23        let fqn = if query.ends_with('.') { query.to_string() } else { format!("{query}.") };
24        match self.txt_lookup(fqn).await {
25            Err(err) => {
26                trace!(target: "disc::dns", %err, ?query, "dns lookup failed");
27                None
28            }
29            Ok(lookup) => find_txt_entry(lookup.answers()),
30        }
31    }
32}
33
34/// Returns the first TXT record with a recognized EIP-1459 entry prefix.
35fn find_txt_entry(records: &[Record]) -> Option<String> {
36    records
37        .iter()
38        .filter_map(|record| {
39            let RData::TXT(txt) = &record.data else { return None };
40            txt_entry(txt)
41        })
42        .find(|entry| has_entry_prefix(entry))
43}
44
45/// Joins all `<character-string>`s of a TXT record into a single entry.
46///
47/// [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035#section-3.3) limits a single
48/// `<character-string>` to 255 bytes, while an
49/// [EIP-1459](https://eips.ethereum.org/EIPS/eip-1459) entry is only bounded by the 512 byte DNS
50/// UDP limit. Entries above 255 bytes are therefore published as several `<character-string>`s
51/// which have to be rejoined without a separator to recover the entry.
52fn txt_entry(txt: &TXT) -> Option<String> {
53    String::from_utf8(txt.txt_data.concat()).ok()
54}
55
56/// An asynchronous DNS resolver
57///
58/// See also [`TokioResolver`]
59///
60/// ```
61/// # fn t() {
62/// use reth_dns_discovery::resolver::DnsResolver;
63/// let resolver = DnsResolver::from_system_conf().unwrap();
64/// # }
65/// ```
66///
67/// Note: This [Resolver] can send multiple lookup attempts, See also
68/// [`ResolverOpts`](hickory_resolver::config::ResolverOpts) which configures 2 attempts (1 retry)
69/// by default.
70#[derive(Clone, Debug)]
71pub struct DnsResolver(TokioResolver);
72
73// === impl DnsResolver ===
74
75impl DnsResolver {
76    /// Create a new resolver by wrapping the given [`TokioResolver`].
77    pub const fn new(resolver: TokioResolver) -> Self {
78        Self(resolver)
79    }
80
81    /// Constructs a new Tokio based Resolver with the system configuration.
82    ///
83    /// This will use `/etc/resolv.conf` on Unix OSes and the registry on Windows.
84    pub fn from_system_conf() -> Result<Self, NetError> {
85        TokioResolver::builder_tokio()?.build().map(Self::new)
86    }
87}
88
89impl Resolver for DnsResolver {
90    async fn lookup_txt(&self, query: &str) -> Option<String> {
91        Resolver::lookup_txt(&self.0, query).await
92    }
93}
94
95/// A [Resolver] that uses an in memory map to lookup entries
96#[derive(Debug, Default)]
97pub struct MapResolver(DashMap<String, String>);
98
99// === impl MapResolver ===
100
101impl MapResolver {
102    /// Inserts a key-value pair into the map.
103    pub fn insert(&self, k: String, v: String) -> Option<String> {
104        self.0.insert(k, v)
105    }
106
107    /// Returns the value corresponding to the key
108    pub fn get(&self, k: &str) -> Option<String> {
109        self.0.get(k).map(|entry| entry.value().clone())
110    }
111
112    /// Removes a key from the map, returning the value at the key if the key was previously in the
113    /// map.
114    pub fn remove(&self, k: &str) -> Option<String> {
115        self.0.remove(k).map(|(_, v)| v)
116    }
117}
118
119impl Resolver for MapResolver {
120    async fn lookup_txt(&self, query: &str) -> Option<String> {
121        self.get(query)
122    }
123}
124
125/// A Resolver that always times out.
126#[cfg(test)]
127pub(crate) struct TimeoutResolver(pub(crate) std::time::Duration);
128
129#[cfg(test)]
130impl Resolver for TimeoutResolver {
131    async fn lookup_txt(&self, _query: &str) -> Option<String> {
132        tokio::time::sleep(self.0).await;
133        None
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use alloy_primitives::keccak256;
141    use data_encoding::BASE32_NOPAD;
142    use hickory_resolver::{
143        config::{ConnectionConfig, NameServerConfig, ResolverConfig},
144        net::runtime::TokioRuntimeProvider,
145        proto::{
146            op::{Message, OpCode},
147            rr::{Name, Record},
148            serialize::binary::{BinDecodable, BinEncodable},
149        },
150    };
151    use std::net::{Ipv4Addr, SocketAddr};
152    use tokio::net::UdpSocket;
153
154    /// Maximum size of a single RFC 1035 `<character-string>`.
155    const MAX_CHARACTER_STRING: usize = 255;
156
157    /// A branch entry of the size go-ethereum's writer emits, which exceeds what a single
158    /// `<character-string>` can hold.
159    fn long_branch_entry() -> String {
160        let children = (0u8..13)
161            .map(|i| BASE32_NOPAD.encode(&keccak256([i]).as_slice()[..16]))
162            .collect::<Vec<_>>()
163            .join(",");
164        format!("enrtree-branch:{children}")
165    }
166
167    fn character_strings(entry: &str) -> Vec<String> {
168        entry
169            .as_bytes()
170            .chunks(MAX_CHARACTER_STRING)
171            .map(|chunk| String::from_utf8(chunk.to_vec()).unwrap())
172            .collect()
173    }
174
175    /// Answers every query with a single TXT record made up of `strings`.
176    async fn spawn_txt_server(strings: Vec<String>) -> SocketAddr {
177        let socket = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
178        let addr = socket.local_addr().unwrap();
179        tokio::spawn(async move {
180            let mut buf = [0u8; 512];
181            while let Ok((len, from)) = socket.recv_from(&mut buf).await {
182                let request = Message::from_bytes(&buf[..len]).unwrap();
183                let query = request.queries.first().unwrap().clone();
184
185                let mut response = Message::response(request.metadata.id, OpCode::Query);
186                response.metadata.authoritative = true;
187                response.metadata.recursion_desired = request.metadata.recursion_desired;
188                response.metadata.recursion_available = true;
189                response.answers.push(Record::from_rdata(
190                    query.name().clone(),
191                    60,
192                    RData::TXT(TXT::new(strings.clone())),
193                ));
194                response.queries.push(query);
195
196                socket.send_to(&response.to_bytes().unwrap(), from).await.unwrap();
197            }
198        });
199        addr
200    }
201
202    fn resolver_for(addr: SocketAddr) -> DnsResolver {
203        let mut connection = ConnectionConfig::udp();
204        connection.port = addr.port();
205        let config = ResolverConfig::from_parts(
206            None,
207            vec![],
208            vec![NameServerConfig::new(addr.ip(), true, vec![connection])],
209        );
210        DnsResolver::new(
211            TokioResolver::builder_with_config(config, TokioRuntimeProvider::default())
212                .build()
213                .unwrap(),
214        )
215    }
216
217    #[test]
218    fn txt_entry_joins_character_strings() {
219        let entry = long_branch_entry();
220        let strings = character_strings(&entry);
221        assert!(entry.len() > MAX_CHARACTER_STRING);
222        assert!(strings.len() > 1);
223
224        assert_eq!(txt_entry(&TXT::new(strings)), Some(entry));
225    }
226
227    #[tokio::test]
228    async fn lookup_txt_reads_record_split_over_character_strings() {
229        let entry = long_branch_entry();
230        let addr = spawn_txt_server(character_strings(&entry)).await;
231
232        let resolved = resolver_for(addr).lookup_txt("YNEGZIWHOM7TOOSUATAPTM.example.org").await;
233
234        assert_eq!(resolved, Some(entry));
235    }
236
237    fn txt_record(entry: &str) -> Record {
238        Record::from_rdata(Name::root(), 60, RData::TXT(TXT::new(vec![entry.to_string()])))
239    }
240
241    #[test]
242    fn find_txt_entry_skips_unrelated_records() {
243        let entry = "enrtree-root:v1 e=enr-root l=link-root seq=1 sig=signature";
244        let records = [
245            txt_record("v=spf1 -all"),
246            txt_record(entry),
247            txt_record("enrtree-branch:YNEGZIWHOM7TOOSUATAPTM"),
248        ];
249
250        assert_eq!(find_txt_entry(&records).as_deref(), Some(entry));
251    }
252}
reth_dns_discovery/resolver.rs

reth_dns_discovery/
resolver.rs
