use std::future::Future; use std::io; use std::net::SocketAddr; use std::pin::Pin; use std::sync::Arc; use std::task::{self, Poll}; use hyper::client::connect::dns as hyper_dns; use hyper::service::Service; use once_cell::sync::Lazy; use tokio::sync::Mutex; use trust_dns_resolver::{ config::{ResolverConfig, ResolverOpts}, lookup_ip::LookupIpIntoIter, system_conf, AsyncResolver, TokioConnection, TokioConnectionProvider, TokioHandle, }; use crate::error::BoxError; type SharedResolver = Arc>; static SYSTEM_CONF: Lazy> = Lazy::new(|| system_conf::read_system_conf().map_err(io::Error::from)); #[derive(Clone)] pub(crate) struct TrustDnsResolver { state: Arc>, } pub(crate) struct SocketAddrs { iter: LookupIpIntoIter, } enum State { Init, Ready(SharedResolver), } impl TrustDnsResolver { pub(crate) fn new() -> io::Result { SYSTEM_CONF.as_ref().map_err(|e| { io::Error::new(e.kind(), format!("error reading DNS system conf: {}", e)) })?; // At this stage, we might not have been called in the context of a // Tokio Runtime, so we must delay the actual construction of the // resolver. Ok(TrustDnsResolver { state: Arc::new(Mutex::new(State::Init)), }) } } impl Service for TrustDnsResolver { type Response = SocketAddrs; type Error = BoxError; type Future = Pin> + Send>>; fn poll_ready(&mut self, _: &mut task::Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn call(&mut self, name: hyper_dns::Name) -> Self::Future { let resolver = self.clone(); Box::pin(async move { let mut lock = resolver.state.lock().await; let resolver = match &*lock { State::Init => { let resolver = new_resolver().await?; *lock = State::Ready(resolver.clone()); resolver } State::Ready(resolver) => resolver.clone(), }; // Don't keep lock once the resolver is constructed, otherwise // only one lookup could be done at a time. drop(lock); let lookup = resolver.lookup_ip(name.as_str()).await?; Ok(SocketAddrs { iter: lookup.into_iter(), }) }) } } impl Iterator for SocketAddrs { type Item = SocketAddr; fn next(&mut self) -> Option { self.iter.next().map(|ip_addr| SocketAddr::new(ip_addr, 0)) } } async fn new_resolver() -> Result { let (config, opts) = SYSTEM_CONF .as_ref() .expect("can't construct TrustDnsResolver if SYSTEM_CONF is error") .clone(); let resolver = AsyncResolver::new(config, opts, TokioHandle)?; Ok(Arc::new(resolver)) }