1use hyper::client::connect::dns::Name;
2use hyper::service::Service;
3
4use std::collections::HashMap;
5use std::future::Future;
6use std::net::SocketAddr;
7use std::pin::Pin;
8use std::sync::Arc;
9use std::task::{Context, Poll};
10
11use crate::error::BoxError;
12
13/// Alias for an `Iterator` trait object over `SocketAddr`.
14pub type Addrs = Box<dyn Iterator<Item = SocketAddr> + Send>;
15
16/// Alias for the `Future` type returned by a DNS resolver.
17pub type Resolving = Pin<Box<dyn Future<Output = Result<Addrs, BoxError>> + Send>>;
18
19/// Trait for customizing DNS resolution in reqwest.
20pub trait Resolve: Send + Sync {
21 /// Performs DNS resolution on a `Name`.
22 /// The return type is a future containing an iterator of `SocketAddr`.
23 ///
24 /// It differs from `tower_service::Service<Name>` in several ways:
25 /// * It is assumed that `resolve` will always be ready to poll.
26 /// * It does not need a mutable reference to `self`.
27 /// * Since trait objects cannot make use of associated types, it requires
28 /// wrapping the returned `Future` and its contained `Iterator` with `Box`.
29 fn resolve(&self, name: Name) -> Resolving;
30}
31
32#[derive(Clone)]
33pub(crate) struct DynResolver {
34 resolver: Arc<dyn Resolve>,
35}
36
37impl DynResolver {
38 pub(crate) fn new(resolver: Arc<dyn Resolve>) -> Self {
39 Self { resolver }
40 }
41}
42
43impl Service<Name> for DynResolver {
44 type Response = Addrs;
45 type Error = BoxError;
46 type Future = Resolving;
47
48 fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
49 Poll::Ready(Ok(()))
50 }
51
52 fn call(&mut self, name: Name) -> Self::Future {
53 self.resolver.resolve(name)
54 }
55}
56
57pub(crate) struct DnsResolverWithOverrides {
58 dns_resolver: Arc<dyn Resolve>,
59 overrides: Arc<HashMap<String, Vec<SocketAddr>>>,
60}
61
62impl DnsResolverWithOverrides {
63 pub(crate) fn new(
64 dns_resolver: Arc<dyn Resolve>,
65 overrides: HashMap<String, Vec<SocketAddr>>,
66 ) -> Self {
67 DnsResolverWithOverrides {
68 dns_resolver,
69 overrides: Arc::new(data:overrides),
70 }
71 }
72}
73
74impl Resolve for DnsResolverWithOverrides {
75 fn resolve(&self, name: Name) -> Resolving {
76 match self.overrides.get(name.as_str()) {
77 Some(dest: &Vec) => {
78 let addrs: Addrs = Box::new(dest.clone().into_iter());
79 Box::pin(futures_util::future::ready(Ok(addrs)))
80 }
81 None => self.dns_resolver.resolve(name),
82 }
83 }
84}
85