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

use std::collections::VecDeque;
use std::marker::PhantomData;
use std::net::IpAddr;

use futures::{Async, Future, Poll};

use tokio_current_thread::spawn;
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};

use super::service::Service;

/// 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);
        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: Self::Request) -> Self::Future {
        ResolverFuture::new(req, &self.resolver)
    }
}

#[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) -> Self {
        // we need to do dns resolution
        let lookup = Some(resolver.lookup_ip(addr.host()));
        ResolverFuture {
            lookup,
            req: Some(addr),
            addrs: None,
        }
    }
}

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),
            }
        }
    }
}