This repository has no description
2.3 kB
86 lines
1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use tokio::sync::OnceCell;
6
7use crate::http::NetworkError;
8
9pub type DnsFuture = Pin<Box<dyn Future<Output = Result<Vec<String>, NetworkError>> + Send>>;
10
11pub trait DnsTxtResolver: Send + Sync + 'static {
12 fn lookup_txt(&self, name: String) -> DnsFuture;
13}
14
15pub struct SystemDns {
16 resolver: Arc<OnceCell<hickory_resolver::TokioAsyncResolver>>,
17}
18
19impl SystemDns {
20 pub fn new() -> Self {
21 Self {
22 resolver: Arc::new(OnceCell::new()),
23 }
24 }
25}
26
27impl Default for SystemDns {
28 fn default() -> Self {
29 Self::new()
30 }
31}
32
33impl DnsTxtResolver for SystemDns {
34 fn lookup_txt(&self, name: String) -> DnsFuture {
35 let cell = self.resolver.clone();
36 Box::pin(async move {
37 let resolver = cell
38 .get_or_try_init(|| async {
39 hickory_resolver::TokioAsyncResolver::tokio_from_system_conf()
40 .map_err(|error| NetworkError::Build(error.to_string()))
41 })
42 .await?;
43 match resolver.txt_lookup(name).await {
44 Ok(lookup) => Ok(lookup.iter().map(render_txt).collect()),
45 Err(error) => match error.kind() {
46 hickory_resolver::error::ResolveErrorKind::NoRecordsFound { .. } => {
47 Ok(Vec::new())
48 }
49 _ => Err(NetworkError::Request(error.to_string())),
50 },
51 }
52 })
53 }
54}
55
56fn render_txt(record: &hickory_resolver::proto::rr::rdata::TXT) -> String {
57 let bytes: Vec<u8> = record
58 .txt_data()
59 .iter()
60 .flat_map(|chunk| chunk.iter().copied())
61 .collect();
62 String::from_utf8_lossy(&bytes).into_owned()
63}
64
65pub struct FakeDns<F> {
66 responder: F,
67}
68
69impl<F> FakeDns<F>
70where
71 F: Fn(&str) -> Result<Vec<String>, NetworkError> + Send + Sync + 'static,
72{
73 pub fn new(responder: F) -> Self {
74 Self { responder }
75 }
76}
77
78impl<F> DnsTxtResolver for FakeDns<F>
79where
80 F: Fn(&str) -> Result<Vec<String>, NetworkError> + Send + Sync + 'static,
81{
82 fn lookup_txt(&self, name: String) -> DnsFuture {
83 let result = (self.responder)(&name);
84 Box::pin(async move { result })
85 }
86}