reqwest/dns/
resolve.rs

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(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) => {
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}