actori_connect/
resolve.rs1use 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
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 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)]
132pub 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}