actix_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 actix_service::{Service, ServiceFactory};
8use futures_util::future::{ok, Either, Ready};
9use trust_dns_resolver::TokioAsyncResolver as AsyncResolver;
10use trust_dns_resolver::{error::ResolveError, lookup_ip::LookupIp};
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    #[allow(clippy::type_complexity)]
110    type Future = Either<
111        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>,
112        Ready<Result<Connect<T>, Self::Error>>,
113    >;
114
115    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
116        Poll::Ready(Ok(()))
117    }
118
119    fn call(&mut self, mut req: Connect<T>) -> Self::Future {
120        if req.addr.is_some() {
121            Either::Right(ok(req))
122        } else if let Ok(ip) = req.host().parse() {
123            req.addr = Some(either::Either::Left(SocketAddr::new(ip, req.port())));
124            Either::Right(ok(req))
125        } else {
126            let resolver = self.resolver.as_ref().map(AsyncResolver::clone);
127            Either::Left(Box::pin(async move {
128                trace!("DNS resolver: resolving host {:?}", req.host());
129                let resolver = if let Some(resolver) = resolver {
130                    resolver
131                } else {
132                    get_default_resolver()
133                        .await
134                        .expect("Failed to get default resolver")
135                };
136                ResolverFuture::new(req, &resolver).await
137            }))
138        }
139    }
140}
141
142type LookupIpFuture = Pin<Box<dyn Future<Output = Result<LookupIp, ResolveError>>>>;
143
144#[doc(hidden)]
145/// Resolver future
146pub struct ResolverFuture<T: Address> {
147    req: Option<Connect<T>>,
148    lookup: LookupIpFuture,
149}
150
151impl<T: Address> ResolverFuture<T> {
152    pub fn new(req: Connect<T>, resolver: &AsyncResolver) -> Self {
153        let host = if let Some(host) = req.host().splitn(2, ':').next() {
154            host
155        } else {
156            req.host()
157        };
158
159        // Clone data to be moved to the lookup future
160        let host_clone = host.to_owned();
161        let resolver_clone = resolver.clone();
162
163        ResolverFuture {
164            lookup: Box::pin(async move {
165                let resolver = resolver_clone;
166                resolver.lookup_ip(host_clone).await
167            }),
168            req: Some(req),
169        }
170    }
171}
172
173impl<T: Address> Future for ResolverFuture<T> {
174    type Output = Result<Connect<T>, ConnectError>;
175
176    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
177        let this = self.get_mut();
178
179        match Pin::new(&mut this.lookup).poll(cx) {
180            Poll::Pending => Poll::Pending,
181            Poll::Ready(Ok(ips)) => {
182                let req = this.req.take().unwrap();
183                let port = req.port();
184                let req = req.set_addrs(ips.iter().map(|ip| SocketAddr::new(ip, port)));
185
186                trace!(
187                    "DNS resolver: host {:?} resolved to {:?}",
188                    req.host(),
189                    req.addrs()
190                );
191
192                if req.addr.is_none() {
193                    Poll::Ready(Err(ConnectError::NoRecords))
194                } else {
195                    Poll::Ready(Ok(req))
196                }
197            }
198            Poll::Ready(Err(e)) => {
199                trace!(
200                    "DNS resolver: failed to resolve host {:?} err: {}",
201                    this.req.as_ref().unwrap().host(),
202                    e
203                );
204                Poll::Ready(Err(e.into()))
205            }
206        }
207    }
208}