actori_connect/
resolve.rs

1use std::future::Future;
2use std::marker::PhantomData;
3use std::net::SocketAddr;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6
7use actori_service::{Service, ServiceFactory};
8use futures::future::{ok, Either, Ready};
9use trust_dns_resolver::lookup_ip::LookupIpFuture;
10use trust_dns_resolver::{AsyncResolver, Background};
11
12use crate::connect::{Address, Connect};
13use crate::error::ConnectError;
14use crate::get_default_resolver;
15
16/// DNS Resolver Service factory
17pub struct ResolverFactory<T> {
18    resolver: Option<AsyncResolver>,
19    _t: PhantomData<T>,
20}
21
22impl<T> ResolverFactory<T> {
23    /// Create new resolver instance with custom configuration and options.
24    pub fn new(resolver: AsyncResolver) -> Self {
25        ResolverFactory {
26            resolver: Some(resolver),
27            _t: PhantomData,
28        }
29    }
30
31    pub fn service(&self) -> Resolver<T> {
32        Resolver {
33            resolver: self.resolver.clone(),
34            _t: PhantomData,
35        }
36    }
37}
38
39impl<T> Default for ResolverFactory<T> {
40    fn default() -> Self {
41        ResolverFactory {
42            resolver: None,
43            _t: PhantomData,
44        }
45    }
46}
47
48impl<T> Clone for ResolverFactory<T> {
49    fn clone(&self) -> Self {
50        ResolverFactory {
51            resolver: self.resolver.clone(),
52            _t: PhantomData,
53        }
54    }
55}
56
57impl<T: Address> ServiceFactory for ResolverFactory<T> {
58    type Request = Connect<T>;
59    type Response = Connect<T>;
60    type Error = ConnectError;
61    type Config = ();
62    type Service = Resolver<T>;
63    type InitError = ();
64    type Future = Ready<Result<Self::Service, Self::InitError>>;
65
66    fn new_service(&self, _: ()) -> Self::Future {
67        ok(self.service())
68    }
69}
70
71/// DNS Resolver Service
72pub struct Resolver<T> {
73    resolver: Option<AsyncResolver>,
74    _t: PhantomData<T>,
75}
76
77impl<T> Resolver<T> {
78    /// Create new resolver instance with custom configuration and options.
79    pub fn new(resolver: AsyncResolver) -> Self {
80        Resolver {
81            resolver: Some(resolver),
82            _t: PhantomData,
83        }
84    }
85}
86
87impl<T> Default for Resolver<T> {
88    fn default() -> Self {
89        Resolver {
90            resolver: None,
91            _t: PhantomData,
92        }
93    }
94}
95
96impl<T> Clone for Resolver<T> {
97    fn clone(&self) -> Self {
98        Resolver {
99            resolver: self.resolver.clone(),
100            _t: PhantomData,
101        }
102    }
103}
104
105impl<T: Address> Service for Resolver<T> {
106    type Request = Connect<T>;
107    type Response = Connect<T>;
108    type Error = ConnectError;
109    type Future = Either<ResolverFuture<T>, Ready<Result<Connect<T>, Self::Error>>>;
110
111    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
112        Poll::Ready(Ok(()))
113    }
114
115    fn call(&mut self, mut req: Connect<T>) -> Self::Future {
116        if req.addr.is_some() {
117            Either::Right(ok(req))
118        } else if let Ok(ip) = req.host().parse() {
119            req.addr = Some(either::Either::Left(SocketAddr::new(ip, req.port())));
120            Either::Right(ok(req))
121        } else {
122            trace!("DNS resolver: resolving host {:?}", req.host());
123            if self.resolver.is_none() {
124                self.resolver = Some(get_default_resolver());
125            }
126            Either::Left(ResolverFuture::new(req, self.resolver.as_ref().unwrap()))
127        }
128    }
129}
130
131#[doc(hidden)]
132/// Resolver future
133pub struct ResolverFuture<T: Address> {
134    req: Option<Connect<T>>,
135    lookup: Background<LookupIpFuture>,
136}
137
138impl<T: Address> ResolverFuture<T> {
139    pub fn new(req: Connect<T>, resolver: &AsyncResolver) -> Self {
140        let lookup = if let Some(host) = req.host().splitn(2, ':').next() {
141            resolver.lookup_ip(host)
142        } else {
143            resolver.lookup_ip(req.host())
144        };
145
146        ResolverFuture {
147            lookup,
148            req: Some(req),
149        }
150    }
151}
152
153impl<T: Address> Future for ResolverFuture<T> {
154    type Output = Result<Connect<T>, ConnectError>;
155
156    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
157        let this = self.get_mut();
158
159        match Pin::new(&mut this.lookup).poll(cx) {
160            Poll::Pending => Poll::Pending,
161            Poll::Ready(Ok(ips)) => {
162                let req = this.req.take().unwrap();
163                let port = req.port();
164                let req = req.set_addrs(ips.iter().map(|ip| SocketAddr::new(ip, port)));
165
166                trace!(
167                    "DNS resolver: host {:?} resolved to {:?}",
168                    req.host(),
169                    req.addrs()
170                );
171
172                if req.addr.is_none() {
173                    Poll::Ready(Err(ConnectError::NoRecords))
174                } else {
175                    Poll::Ready(Ok(req))
176                }
177            }
178            Poll::Ready(Err(e)) => {
179                trace!(
180                    "DNS resolver: failed to resolve host {:?} err: {}",
181                    this.req.as_ref().unwrap().host(),
182                    e
183                );
184                Poll::Ready(Err(e.into()))
185            }
186        }
187    }
188}