reqwest/dns/
resolve.rs

1use hyper_util::client::legacy::connect::dns::Name as HyperName;
2use tower_service::Service;
3
4use std::collections::HashMap;
5use std::future::Future;
6use std::net::SocketAddr;
7use std::pin::Pin;
8use std::str::FromStr;
9use std::sync::Arc;
10use std::task::{Context, Poll};
11
12use crate::error::BoxError;
13
14/// Alias for an `Iterator` trait object over `SocketAddr`.
15pub type Addrs = Box<dyn Iterator<Item = SocketAddr> + Send>;
16
17/// Alias for the `Future` type returned by a DNS resolver.
18pub type Resolving = Pin<Box<dyn Future<Output = Result<Addrs, BoxError>> + Send>>;
19
20/// Trait for customizing DNS resolution in reqwest.
21pub trait Resolve: Send + Sync {
22    /// Performs DNS resolution on a `Name`.
23    /// The return type is a future containing an iterator of `SocketAddr`.
24    ///
25    /// It differs from `tower_service::Service<Name>` in several ways:
26    ///  * It is assumed that `resolve` will always be ready to poll.
27    ///  * It does not need a mutable reference to `self`.
28    ///  * Since trait objects cannot make use of associated types, it requires
29    ///    wrapping the returned `Future` and its contained `Iterator` with `Box`.
30    fn resolve(&self, name: Name) -> Resolving;
31}
32
33/// A name that must be resolved to addresses.
34#[derive(Debug)]
35pub struct Name(pub(super) HyperName);
36
37impl Name {
38    /// View the name as a string.
39    pub fn as_str(&self) -> &str {
40        self.0.as_str()
41    }
42}
43
44impl FromStr for Name {
45    type Err = sealed::InvalidNameError;
46
47    fn from_str(host: &str) -> Result<Self, Self::Err> {
48        HyperName::from_str(host.into())
49            .map(Name)
50            .map_err(|_| sealed::InvalidNameError { _ext: () })
51    }
52}
53
54#[derive(Clone)]
55pub(crate) struct DynResolver {
56    resolver: Arc<dyn Resolve>,
57}
58
59impl DynResolver {
60    pub(crate) fn new(resolver: Arc<dyn Resolve>) -> Self {
61        Self { resolver }
62    }
63}
64
65impl Service<HyperName> for DynResolver {
66    type Response = Addrs;
67    type Error = BoxError;
68    type Future = Resolving;
69
70    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
71        Poll::Ready(Ok(()))
72    }
73
74    fn call(&mut self, name: HyperName) -> Self::Future {
75        self.resolver.resolve(Name(name))
76    }
77}
78
79pub(crate) struct DnsResolverWithOverrides {
80    dns_resolver: Arc<dyn Resolve>,
81    overrides: Arc<HashMap<String, Vec<SocketAddr>>>,
82}
83
84impl DnsResolverWithOverrides {
85    pub(crate) fn new(
86        dns_resolver: Arc<dyn Resolve>,
87        overrides: HashMap<String, Vec<SocketAddr>>,
88    ) -> Self {
89        DnsResolverWithOverrides {
90            dns_resolver,
91            overrides: Arc::new(overrides),
92        }
93    }
94}
95
96impl Resolve for DnsResolverWithOverrides {
97    fn resolve(&self, name: Name) -> Resolving {
98        match self.overrides.get(name.as_str()) {
99            Some(dest) => {
100                let addrs: Addrs = Box::new(dest.clone().into_iter());
101                Box::pin(futures_util::future::ready(Ok(addrs)))
102            }
103            None => self.dns_resolver.resolve(name),
104        }
105    }
106}
107
108mod sealed {
109    use std::fmt;
110
111    #[derive(Debug)]
112    pub struct InvalidNameError {
113        pub(super) _ext: (),
114    }
115
116    impl fmt::Display for InvalidNameError {
117        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118            f.write_str("invalid DNS name")
119        }
120    }
121
122    impl std::error::Error for InvalidNameError {}
123}