Add a temporary IP filter

src/async_impl/client.rs

@@ -236,6 +236,8 @@ struct Config {
     #[cfg(feature = "cookies")]
     cookie_store: Option<Arc<dyn cookie::CookieStore>>,
     hickory_dns: bool,
+    #[cfg(feature = "hickory-dns")]
+    ip_filter: fn(std::net::IpAddr) -> bool,
     error: Option<crate::Error>,
     https_only: bool,
     #[cfg(feature = "http3")]
@@ -361,6 +363,8 @@ impl ClientBuilder {
                 interface: None,
                 nodelay: true,
                 hickory_dns: cfg!(feature = "hickory-dns"),
+                #[cfg(feature = "hickory-dns")]
+                ip_filter: |_| true,
                 #[cfg(feature = "cookies")]
                 cookie_store: None,
                 https_only: false,
@@ -417,7 +421,7 @@ impl ClientBuilder {
             let mut resolver: Arc<dyn Resolve> = match config.hickory_dns {
                 false => Arc::new(GaiResolver::new()),
                 #[cfg(feature = "hickory-dns")]
-                true => Arc::new(HickoryDnsResolver::default()),
+                true => Arc::new(HickoryDnsResolver::new(config.ip_filter)),
                 #[cfg(not(feature = "hickory-dns"))]
                 true => unreachable!("hickory-dns shouldn't be enabled unless the feature is"),
             };
@@ -999,6 +1003,9 @@ impl ClientBuilder {
         };
 
         let redirect_policy = {
+            #[cfg(feature = "hickory-dns")]
+            let mut p = TowerRedirectPolicy::new(config.redirect_policy, config.ip_filter);
+            #[cfg(not(feature = "hickory-dns"))]
             let mut p = TowerRedirectPolicy::new(config.redirect_policy);
             p.with_referer(config.referer)
                 .with_https_only(config.https_only);
@@ -1044,6 +1051,8 @@ impl ClientBuilder {
                 proxies_maybe_http_custom_headers,
                 https_only: config.https_only,
                 redirect_policy_desc,
+                #[cfg(feature = "hickory-dns")]
+                ip_filter: config.ip_filter,
             }),
         })
     }
@@ -2189,6 +2198,17 @@ impl ClientBuilder {
         }
     }
 
+    /// Adds a filter for valid IP addresses during DNS lookup.
+    ///
+    /// # Optional
+    ///
+    /// This requires the optional `hickory-dns` feature to be enabled.
+    #[cfg(feature = "hickory-dns")]
+    pub fn ip_filter(mut self, filter: fn(std::net::IpAddr) -> bool) -> ClientBuilder {
+        self.config.ip_filter = filter;
+        self
+    }
+
     /// Override DNS resolution for specific domains to a particular IP address.
     ///
     /// Set the port to `0` to use the conventional port for the given scheme (e.g. 80 for http).
@@ -2529,6 +2549,10 @@ impl Client {
             }
         }
 
+        #[cfg(feature = "hickory-dns")]
+        if let Err(err) = redirect::validate_url(self.inner.ip_filter, &url) {
+            return Pending::new_err(err);
+        }
         let uri = match try_uri(&url) {
             Ok(uri) => uri,
             _ => return Pending::new_err(error::url_invalid_uri(url)),
@@ -2832,6 +2856,8 @@ struct ClientRef {
     proxies_maybe_http_custom_headers: bool,
     https_only: bool,
     redirect_policy_desc: Option<String>,
+    #[cfg(feature = "hickory-dns")]
+    ip_filter: fn(IpAddr) -> bool,
 }
 
 impl ClientRef {

src/dns/hickory.rs

@@ -12,16 +12,27 @@ use std::sync::Arc;
 use super::{Addrs, Name, Resolve, Resolving};
 
 /// Wrapper around an `AsyncResolver`, which implements the `Resolve` trait.
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Clone)]
 pub(crate) struct HickoryDnsResolver {
     /// Since we might not have been called in the context of a
     /// Tokio Runtime in initialization, so we must delay the actual
     /// construction of the resolver.
     state: Arc<OnceCell<TokioResolver>>,
+    filter: fn(std::net::IpAddr) -> bool,
 }
 
 struct SocketAddrs {
     iter: LookupIpIntoIter,
+    filter: fn(std::net::IpAddr) -> bool,
+}
+
+impl HickoryDnsResolver {
+    pub fn new(filter: fn(std::net::IpAddr) -> bool) -> Self {
+        Self {
+            state: Default::default(),
+            filter,
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -31,11 +42,18 @@ impl Resolve for HickoryDnsResolver {
     fn resolve(&self, name: Name) -> Resolving {
         let resolver = self.clone();
         Box::pin(async move {
+            let filter = resolver.filter;
             let resolver = resolver.state.get_or_try_init(new_resolver)?;
 
             let lookup = resolver.lookup_ip(name.as_str()).await?;
+            if !lookup.iter().any(filter) {
+                let e = hickory_resolver::ResolveError::from("destination is restricted");
+                return Err(e.into());
+            }
+
             let addrs: Addrs = Box::new(SocketAddrs {
                 iter: lookup.into_iter(),
+                filter,
             });
             Ok(addrs)
         })
@@ -46,7 +64,12 @@ impl Iterator for SocketAddrs {
     type Item = SocketAddr;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.iter.next().map(|ip_addr| SocketAddr::new(ip_addr, 0))
+        loop {
+            let ip_addr = self.iter.next()?;
+            if (self.filter)(ip_addr) {
+                return Some(SocketAddr::new(ip_addr, 0));
+            }
+        }
     }
 }
 

src/error.rs

@@ -152,6 +152,14 @@ impl Error {
                 if hyper_err.is_connect() {
                     return true;
                 }
+            } else {
+                #[cfg(feature = "hickory-dns")]
+                if err
+                    .downcast_ref::<hickory_resolver::ResolveError>()
+                    .is_some()
+                {
+                    return true;
+                }
             }
 
             source = err.source();

src/redirect.rs

@@ -5,6 +5,8 @@
 //! `redirect::Policy` can be used with a `ClientBuilder`.
 
 use std::fmt;
+#[cfg(feature = "hickory-dns")]
+use std::net::IpAddr;
 use std::{error::Error as StdError, sync::Arc};
 
 use crate::header::{AUTHORIZATION, COOKIE, PROXY_AUTHORIZATION, REFERER, WWW_AUTHENTICATE};
@@ -267,9 +269,23 @@ pub(crate) struct TowerRedirectPolicy {
     referer: bool,
     urls: Vec<Url>,
     https_only: bool,
+    #[cfg(feature = "hickory-dns")]
+    filter: fn(std::net::IpAddr) -> bool,
 }
 
 impl TowerRedirectPolicy {
+    #[cfg(feature = "hickory-dns")]
+    pub(crate) fn new(policy: Policy, filter: fn(std::net::IpAddr) -> bool) -> Self {
+        Self {
+            policy: Arc::new(policy),
+            referer: false,
+            urls: Vec::new(),
+            https_only: false,
+            filter,
+        }
+    }
+
+    #[cfg(not(feature = "hickory-dns"))]
     pub(crate) fn new(policy: Policy) -> Self {
         Self {
             policy: Arc::new(policy),
@@ -302,6 +318,21 @@ fn make_referer(next: &Url, previous: &Url) -> Option<HeaderValue> {
     referer.as_str().parse().ok()
 }
 
+#[cfg(feature = "hickory-dns")]
+pub(crate) fn validate_url(ip_filter: fn(IpAddr) -> bool, url: &Url) -> Result<(), crate::Error> {
+    let is_valid_ip = match url.host() {
+        Some(url::Host::Ipv4(ip)) => (ip_filter)(IpAddr::V4(ip)),
+        Some(url::Host::Ipv6(ip)) => (ip_filter)(IpAddr::V6(ip)),
+        _ => true,
+    };
+
+    if !is_valid_ip {
+        let e = hickory_resolver::ResolveError::from("destination is restricted");
+        return Err(crate::Error::new(crate::error::Kind::Request, Some(e)));
+    }
+    Ok(())
+}
+
 impl TowerPolicy<async_impl::body::Body, crate::Error> for TowerRedirectPolicy {
     fn redirect(&mut self, attempt: &TowerAttempt<'_>) -> Result<TowerAction, crate::Error> {
         let previous_url =
@@ -313,6 +344,8 @@ impl TowerPolicy<async_impl::body::Body, crate::Error> for TowerRedirectPolicy {
         };
 
         self.urls.push(previous_url.clone());
+        #[cfg(feature = "hickory-dns")]
+        validate_url(self.filter, &next_url)?;
 
         match self.policy.check(attempt.status(), &next_url, &self.urls) {
             ActionKind::Follow => {