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
16pub struct ResolverFactory<T> {
18 resolver: Option<AsyncResolver>,
19 _t: PhantomData<T>,
20}
21
22impl<T> ResolverFactory<T> {
23 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
71pub struct Resolver<T> {
73 resolver: Option<AsyncResolver>,
74 _t: PhantomData<T>,
75}
76
77impl<T> Resolver<T> {
78 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)]
145pub 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 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}