1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

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<T = String> {
    resolver: AsyncResolver,
    req: PhantomData<T>,
}

impl<T: RequestHost> Default for Resolver<T> {
    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<T: RequestHost> Resolver<T> {
    /// 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<T2: RequestHost>(&self) -> Resolver<T2> {
        Resolver {
            resolver: self.resolver.clone(),
            req: PhantomData,
        }
    }
}

impl<T> Clone for Resolver<T> {
    fn clone(&self) -> Self {
        Resolver {
            resolver: self.resolver.clone(),
            req: PhantomData,
        }
    }
}

impl<T: RequestHost> Service for Resolver<T> {
    type Request = T;
    type Response = (T, VecDeque<IpAddr>);
    type Error = ResolveError;
    type Future = ResolverFuture<T>;

    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<T> {
    req: Option<T>,
    lookup: Option<Background<LookupIpFuture>>,
    addrs: Option<VecDeque<IpAddr>>,
}

impl<T: RequestHost> ResolverFuture<T> {
    pub fn new(addr: T, resolver: &AsyncResolver, addrs: Option<VecDeque<IpAddr>>) -> Self {
        // we need to do dns resolution
        let lookup = Some(resolver.lookup_ip(addr.host()));
        ResolverFuture {
            lookup,
            addrs,
            req: Some(addr),
        }
    }
}

impl<T: RequestHost> Future for ResolverFuture<T> {
    type Item = (T, VecDeque<IpAddr>);
    type Error = ResolveError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        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),
            }
        }
    }
}