ate_comms/dns/
client.rs

1use std::net::SocketAddr;
2use std::net::ToSocketAddrs;
3use derivative::*;
4use tokio::net::TcpStream as TokioTcpStream;
5#[allow(unused_imports, dead_code)]
6use tracing::{debug, error, info, instrument, span, trace, warn, Level};
7use tokio::sync::Mutex;
8use {
9    trust_dns_client::client::*, trust_dns_client::op::DnsResponse, trust_dns_client::tcp::*,
10    trust_dns_proto::iocompat::AsyncIoTokioAsStd, trust_dns_proto::DnssecDnsHandle,
11};
12
13pub use {trust_dns_client::error::ClientError, trust_dns_client::rr::*};
14
15#[derive(Derivative)]
16#[derivative(Debug)]
17pub enum DnsClient {
18    Dns {
19        dns_server: String,
20        #[derivative(Debug = "ignore")]
21        #[cfg(feature = "dns")]
22        client: Mutex<MemoizeClientHandle<AsyncClient>>,
23    },
24    DnsSec {
25        dns_server: String,
26        #[derivative(Debug = "ignore")]
27        #[cfg(feature = "dns")]
28        client: Mutex<DnssecDnsHandle<MemoizeClientHandle<AsyncClient>>>,
29    },
30}
31
32impl DnsClient {
33    pub async fn connect(dns_server: &str, dns_sec: bool) -> Self {
34        debug!("using DNS server: {}", dns_server);
35        let addr: SocketAddr = (dns_server.to_string(), 53)
36            .to_socket_addrs()
37            .unwrap()
38            .next()
39            .unwrap();
40
41        let (stream, sender) = TcpClientStream::<AsyncIoTokioAsStd<TokioTcpStream>>::new(addr);
42        let client = AsyncClient::new(stream, sender, None);
43        let (client, bg) = client.await.expect("client failed to connect");
44        wasmer_bus::task::spawn(bg);
45
46        let client = MemoizeClientHandle::new(client);
47
48        match dns_sec {
49            false => {
50                debug!("configured for DNSSec");
51                Self::Dns {
52                    dns_server: dns_server.to_string(),
53                    client: Mutex::new(client),
54                }
55            }
56            true => {
57                debug!("configured for plain DNS");
58                Self::DnsSec {
59                    dns_server: dns_server.to_string(),
60                    client: Mutex::new(DnssecDnsHandle::new(client.clone())),
61                }
62            }
63        }
64    }
65
66    pub async fn reconnect(&mut self) {
67        let (dns_server, dns_sec) = match self {
68            Self::Dns { dns_server, client: _ } => (dns_server.clone(), false),
69            Self::DnsSec { dns_server, client: _ } => (dns_server.clone(), true),
70        };
71
72        *self = Self::connect(dns_server.as_str(), dns_sec).await;
73    }
74
75    pub async fn query(
76        &mut self,
77        name: Name,
78        query_class: DNSClass,
79        query_type: RecordType,
80    ) -> Result<DnsResponse, ClientError> {
81        let ret = {
82            match self {
83                Self::Dns { client: c, .. } => {
84                    let mut c = c.lock().await;
85                    c.query(name.clone(), query_class, query_type).await
86                }
87                Self::DnsSec { client: c, .. } => {
88                    let mut c = c.lock().await;
89                    c.query(name.clone(), query_class, query_type).await
90                }
91            }
92        };
93
94        match ret {
95            Ok(a) => Ok(a),
96            Err(_) => {
97                self.reconnect().await;
98
99                match self {
100                    Self::Dns { client: c, .. } => {
101                        let mut c = c.lock().await;
102                        c.query(name, query_class, query_type).await
103                    }
104                    Self::DnsSec { client: c, .. } => {
105                        let mut c = c.lock().await;
106                        c.query(name, query_class, query_type).await
107                    }
108                }
109            }
110        }
111    }
112}
ate_comms/dns/client.rs

ate_comms/dns/
client.rs
