Skip to main content

saluki_io/net/dns/
hyper.rs

1use std::{
2    future::Future,
3    net::{IpAddr, SocketAddr},
4    pin::Pin,
5    sync::Arc,
6    task::{Context, Poll},
7};
8
9use hickory_resolver::{net::NetError, TokioResolver};
10use hyper_util::client::legacy::connect::{dns::Name, HttpConnector};
11use metrics::Counter;
12use saluki_error::{ErrorContext as _, GenericError};
13use tower::Service;
14
15/// An [`HttpConnector`] that uses [`HickoryResolver`].
16pub type HickoryHttpConnector = HttpConnector<HickoryResolver>;
17
18/// A DNS resolver for `hyper` based on `hickory`.
19#[derive(Clone)]
20pub struct HickoryResolver {
21    resolver: Arc<TokioResolver>,
22    lookup_errors: Option<Counter>,
23}
24
25impl HickoryResolver {
26    /// Creates a new [`HickoryResolver`] based on the system's resolver configuration.
27    ///
28    /// # Panics
29    ///
30    /// This will panic if called outside the context of a Tokio runtime.
31    ///
32    /// # Errors
33    ///
34    /// If there is an issue loading the system configuration, an error is returned.
35    pub fn from_system_conf() -> Result<Self, GenericError> {
36        let resolver = TokioResolver::builder_tokio()
37            .error_context("Failed to load the system resolver configuration.")?
38            .build()
39            .error_context("Failed to build the resolver.")?;
40
41        Ok(Self {
42            resolver: Arc::new(resolver),
43            lookup_errors: None,
44        })
45    }
46
47    /// Sets a counter that's incremented when DNS lookup fails.
48    pub fn with_lookup_errors_counter(mut self, counter: Counter) -> Self {
49        self.lookup_errors = Some(counter);
50        self
51    }
52
53    /// Consumes `self` and creates a new [`HickoryHttpConnector`] with this resolver.
54    pub fn into_http_connector(self) -> HickoryHttpConnector {
55        HickoryHttpConnector::new_with_resolver(self)
56    }
57}
58
59impl Service<Name> for HickoryResolver {
60    type Response = SocketAddrs;
61    type Error = NetError;
62
63    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
64
65    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
66        Poll::Ready(Ok(()))
67    }
68
69    fn call(&mut self, name: Name) -> Self::Future {
70        let resolver = self.resolver.clone();
71        let lookup_errors = self.lookup_errors.clone();
72
73        Box::pin(async move {
74            let response = match resolver.lookup_ip(name.as_str()).await {
75                Ok(response) => response,
76                Err(error) => {
77                    if let Some(lookup_errors) = lookup_errors {
78                        lookup_errors.increment(1);
79                    }
80                    return Err(error);
81                }
82            };
83            Ok(response.iter().collect())
84        })
85    }
86}
87
88pub struct SocketAddrs(Vec<IpAddr>);
89
90impl Iterator for SocketAddrs {
91    type Item = SocketAddr;
92
93    fn next(&mut self) -> Option<Self::Item> {
94        self.0.pop().map(|ip| SocketAddr::new(ip, 0))
95    }
96}
97
98impl FromIterator<IpAddr> for SocketAddrs {
99    fn from_iter<I: IntoIterator<Item = IpAddr>>(iter: I) -> Self {
100        Self(iter.into_iter().collect())
101    }
102}