use std::collections::VecDeque; use std::marker::PhantomData; use std::net::IpAddr; use actix_service::Service; use futures::{Async, Future, Poll}; use trust_dns_resolver::config::{ResolverConfig, ResolverOpts}; pub use trust_dns_resolver::error::ResolveError; use trust_dns_resolver::lookup_ip::LookupIpFuture; use trust_dns_resolver::system_conf::read_system_conf; use trust_dns_resolver::{AsyncResolver, Background}; /// Host name of the request pub trait RequestHost { fn host(&self) -> &str; } impl RequestHost for String { fn host(&self) -> &str { self.as_ref() } } pub struct Resolver { resolver: AsyncResolver, req: PhantomData, } impl Default for Resolver { fn default() -> Self { let (cfg, opts) = if let Ok((cfg, opts)) = read_system_conf() { (cfg, opts) } else { (ResolverConfig::default(), ResolverOpts::default()) }; Resolver::new(cfg, opts) } } impl Resolver { /// Create new resolver instance with custom configuration and options. pub fn new(cfg: ResolverConfig, opts: ResolverOpts) -> Self { let (resolver, bg) = AsyncResolver::new(cfg, opts); tokio_current_thread::spawn(bg); Resolver { resolver, req: PhantomData, } } /// Change type of resolver request. pub fn into_request(&self) -> Resolver { Resolver { resolver: self.resolver.clone(), req: PhantomData, } } } impl Clone for Resolver { fn clone(&self) -> Self { Resolver { resolver: self.resolver.clone(), req: PhantomData, } } } impl Service for Resolver { type Request = T; type Response = (T, VecDeque); type Error = ResolveError; type Future = ResolverFuture; fn poll_ready(&mut self) -> Poll<(), Self::Error> { Ok(Async::Ready(())) } fn call(&mut self, req: T) -> Self::Future { if let Ok(ip) = req.host().parse() { let mut addrs = VecDeque::new(); addrs.push_back(ip); ResolverFuture::new(req, &self.resolver, Some(addrs)) } else { ResolverFuture::new(req, &self.resolver, None) } } } #[doc(hidden)] /// Resolver future pub struct ResolverFuture { req: Option, lookup: Option>, addrs: Option>, } impl ResolverFuture { pub fn new(addr: T, resolver: &AsyncResolver, addrs: Option>) -> Self { // we need to do dns resolution let lookup = Some(resolver.lookup_ip(addr.host())); ResolverFuture { lookup, addrs, req: Some(addr), } } } impl Future for ResolverFuture { type Item = (T, VecDeque); type Error = ResolveError; fn poll(&mut self) -> Poll { if let Some(addrs) = self.addrs.take() { Ok(Async::Ready((self.req.take().unwrap(), addrs))) } else { match self.lookup.as_mut().unwrap().poll() { Ok(Async::NotReady) => Ok(Async::NotReady), Ok(Async::Ready(ips)) => Ok(Async::Ready(( self.req.take().unwrap(), ips.iter().collect(), ))), Err(err) => Err(err), } } } }