use std::future::Future; use std::marker::PhantomData; use std::net::SocketAddr; use std::pin::Pin; use std::task::{Context, Poll}; use actix_service::{Service, ServiceFactory}; use futures_util::future::{ok, Either, Ready}; use trust_dns_resolver::TokioAsyncResolver as AsyncResolver; use trust_dns_resolver::{error::ResolveError, lookup_ip::LookupIp}; use crate::connect::{Address, Connect}; use crate::error::ConnectError; use crate::get_default_resolver; /// DNS Resolver Service factory pub struct ResolverFactory { resolver: Option, _t: PhantomData, } impl ResolverFactory { /// Create new resolver instance with custom configuration and options. pub fn new(resolver: AsyncResolver) -> Self { ResolverFactory { resolver: Some(resolver), _t: PhantomData, } } pub fn service(&self) -> Resolver { Resolver { resolver: self.resolver.clone(), _t: PhantomData, } } } impl Default for ResolverFactory { fn default() -> Self { ResolverFactory { resolver: None, _t: PhantomData, } } } impl Clone for ResolverFactory { fn clone(&self) -> Self { ResolverFactory { resolver: self.resolver.clone(), _t: PhantomData, } } } impl ServiceFactory> for ResolverFactory { type Response = Connect; type Error = ConnectError; type Config = (); type Service = Resolver; type InitError = (); type Future = Ready>; fn new_service(&self, _: ()) -> Self::Future { ok(self.service()) } } /// DNS Resolver Service pub struct Resolver { resolver: Option, _t: PhantomData, } impl Resolver { /// Create new resolver instance with custom configuration and options. pub fn new(resolver: AsyncResolver) -> Self { Resolver { resolver: Some(resolver), _t: PhantomData, } } } impl Default for Resolver { fn default() -> Self { Resolver { resolver: None, _t: PhantomData, } } } impl Clone for Resolver { fn clone(&self) -> Self { Resolver { resolver: self.resolver.clone(), _t: PhantomData, } } } impl Service> for Resolver { type Response = Connect; type Error = ConnectError; #[allow(clippy::type_complexity)] type Future = Either< Pin>>>, Ready, Self::Error>>, >; actix_service::always_ready!(); fn call(&mut self, mut req: Connect) -> Self::Future { if req.addr.is_some() { Either::Right(ok(req)) } else if let Ok(ip) = req.host().parse() { req.addr = Some(either::Either::Left(SocketAddr::new(ip, req.port()))); Either::Right(ok(req)) } else { let resolver = self.resolver.as_ref().map(AsyncResolver::clone); Either::Left(Box::pin(async move { trace!("DNS resolver: resolving host {:?}", req.host()); let resolver = if let Some(resolver) = resolver { resolver } else { get_default_resolver() .await .expect("Failed to get default resolver") }; ResolverFuture::new(req, &resolver).await })) } } } type LookupIpFuture = Pin>>>; #[doc(hidden)] /// Resolver future pub struct ResolverFuture { req: Option>, lookup: LookupIpFuture, } impl ResolverFuture { pub fn new(req: Connect, resolver: &AsyncResolver) -> Self { let host = if let Some(host) = req.host().splitn(2, ':').next() { host } else { req.host() }; // Clone data to be moved to the lookup future let host_clone = host.to_owned(); let resolver_clone = resolver.clone(); ResolverFuture { lookup: Box::pin(async move { let resolver = resolver_clone; resolver.lookup_ip(host_clone).await }), req: Some(req), } } } impl Future for ResolverFuture { type Output = Result, ConnectError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); match Pin::new(&mut this.lookup).poll(cx) { Poll::Pending => Poll::Pending, Poll::Ready(Ok(ips)) => { let req = this.req.take().unwrap(); let port = req.port(); let req = req.set_addrs(ips.iter().map(|ip| SocketAddr::new(ip, port))); trace!( "DNS resolver: host {:?} resolved to {:?}", req.host(), req.addrs() ); if req.addr.is_none() { Poll::Ready(Err(ConnectError::NoRecords)) } else { Poll::Ready(Ok(req)) } } Poll::Ready(Err(e)) => { trace!( "DNS resolver: failed to resolve host {:?} err: {}", this.req.as_ref().unwrap().host(), e ); Poll::Ready(Err(e.into())) } } } }