From 89b79fa2cf95801a3787f2dda2f16c1d857b8810 Mon Sep 17 00:00:00 2001 From: ElJeffe <161904+eljeffeg@users.noreply.github.com> Date: Wed, 29 Jul 2026 12:14:53 -0400 Subject: [PATCH] Add custom client resolver support --- src/client.rs | 174 +++++++++++++++++++++++++++++++++++++++++++------- src/lib.rs | 2 +- 2 files changed, 152 insertions(+), 24 deletions(-) diff --git a/src/client.rs b/src/client.rs index 8307811..f458a0e 100644 --- a/src/client.rs +++ b/src/client.rs @@ -44,15 +44,6 @@ use rustls_pki_types::ServerName; use rustls_platform_verifier::ConfigVerifierExt; use std::io::{BufReader, BufWriter, Error, ErrorKind, Result}; use std::net::{SocketAddr, TcpStream, ToSocketAddrs}; -#[cfg(all( - any( - feature = "rustls-aws-lc-webpki", - feature = "rustls-ring-webpki", - feature = "rustls-aws-lc-native", - feature = "rustls-ring-native" - ), - not(feature = "native-tls") -))] use std::sync::Arc; #[cfg(any( feature = "rustls-aws-lc-webpki", @@ -72,6 +63,25 @@ use url::Url; ))] use webpki_roots::TLS_SERVER_ROOTS; +/// Resolves a request URI to the socket addresses the client is allowed to +/// connect to. +/// +/// Custom resolvers may implement application-specific DNS resolution or +/// destination policies. The URI hostname is still used for the HTTP `Host` +/// header and TLS certificate verification. +pub trait Resolver: Send + Sync + 'static { + fn resolve(&self, uri: &Uri, default_port: u16) -> Result>; +} + +impl Resolver for F +where + F: Fn(&Uri, u16) -> Result> + Send + Sync + 'static, +{ + fn resolve(&self, uri: &Uri, default_port: u16) -> Result> { + self(uri, default_port) + } +} + /// An HTTP client. /// /// It aims at following the basic concepts of the [Web Fetch standard](https://fetch.spec.whatwg.org/) without the bits specific to web browsers (context, CORS...). @@ -119,6 +129,7 @@ pub struct Client { timeout: Option, user_agent: Option, redirection_limit: usize, + resolver: Option>, } impl Client { @@ -152,6 +163,16 @@ impl Client { self } + /// Sets the resolver used for the initial request and every redirect. + /// + /// The returned addresses are passed directly to the connection attempt, + /// while the URI hostname remains unchanged for HTTP and TLS. + #[inline] + pub fn with_resolver(mut self, resolver: impl Resolver) -> Self { + self.resolver = Some(Arc::new(resolver)); + self + } + pub fn request(&self, request: Request>) -> Result> { let mut request = request.map(Into::into); // Loops the number of allowed redirections + 1 @@ -232,7 +253,7 @@ impl Client { })?; if *scheme == Scheme::HTTP { - let addresses = get_and_validate_socket_addresses(request.uri(), 80)?; + let addresses = self.resolve_socket_addresses(request.uri(), 80)?; let stream = self.connect(&addresses)?; let stream = encode_request(request, BufWriter::with_capacity(BUFFER_CAPACITY, stream))? @@ -245,7 +266,7 @@ impl Client { if *scheme == Scheme::HTTPS { static TLS_CONNECTOR: OnceLock = OnceLock::new(); - let addresses = get_and_validate_socket_addresses(request.uri(), 443)?; + let addresses = self.resolve_socket_addresses(request.uri(), 443)?; let stream = self.connect(&addresses)?; let stream = TLS_CONNECTOR .get_or_init(|| match TlsConnector::new() { @@ -294,7 +315,7 @@ impl Client { .with_no_client_auth(), ) }); - let addresses = get_and_validate_socket_addresses(request.uri(), 443)?; + let addresses = self.resolve_socket_addresses(request.uri(), 443)?; let dns_name = ServerName::try_from(host) .map_err(invalid_input_error)? .to_owned(); @@ -324,6 +345,28 @@ impl Client { ))) } + fn resolve_socket_addresses(&self, uri: &Uri, default_port: u16) -> Result> { + let addresses = if let Some(resolver) = &self.resolver { + resolver.resolve(uri, default_port)? + } else { + default_resolve(uri, default_port)? + }; + if addresses.is_empty() { + return Err(invalid_input_error(format!( + "No socket addresses resolved for request URL {uri}" + ))); + } + for address in &addresses { + if BAD_PORTS.binary_search(&address.port()).is_ok() { + return Err(invalid_input_error(format!( + "The port {} is not allowed for HTTP(S) because it is dedicated to an other use", + address.port() + ))); + } + } + Ok(addresses) + } + fn connect(&self, addresses: &[SocketAddr]) -> Result { let stream = if let Some(timeout) = self.timeout { Self::connect_timeout(addresses, timeout) @@ -361,21 +404,12 @@ const BAD_PORTS: [u16; 80] = [ 6697, 10080, ]; -fn get_and_validate_socket_addresses(uri: &Uri, default_port: u16) -> Result> { +fn default_resolve(uri: &Uri, default_port: u16) -> Result> { let host = uri .host() .ok_or_else(|| invalid_input_error(format!("No host in request URL {uri}")))?; let port = uri.port_u16().unwrap_or(default_port); - let addresses = (host, port).to_socket_addrs()?.collect::>(); - for address in &addresses { - if BAD_PORTS.binary_search(&address.port()).is_ok() { - return Err(invalid_input_error(format!( - "The port {} is not allowed for HTTP(S) because it is dedicated to an other use", - address.port() - ))); - } - } - Ok(addresses) + Ok((host, port).to_socket_addrs()?.collect()) } fn join_urls(base: &Uri, relative: &str) -> Result { @@ -469,6 +503,100 @@ mod tests { .is_err()); } + #[test] + fn test_custom_resolver() -> Result<()> { + let expected = "192.0.2.1:8080".parse().unwrap(); + let client = Client::new().with_resolver(move |uri: &Uri, default_port| { + assert_eq!(uri, &"http://example.test/path".parse::().unwrap()); + assert_eq!(default_port, 80); + Ok(vec![expected]) + }); + assert_eq!( + client.resolve_socket_addresses( + &"http://example.test/path".parse::().unwrap(), + 80 + )?, + vec![expected] + ); + Ok(()) + } + + #[test] + fn test_custom_resolver_empty_result() { + let client = Client::new().with_resolver(|_: &Uri, _| Ok(Vec::new())); + assert!(client + .resolve_socket_addresses(&"http://example.test".parse::().unwrap(), 80) + .is_err()); + } + + #[test] + fn test_custom_resolver_cannot_bypass_bad_port_check() { + let client = + Client::new().with_resolver(|_: &Uri, _| Ok(vec!["192.0.2.1:22".parse().unwrap()])); + assert!(client + .resolve_socket_addresses(&"http://example.test".parse::().unwrap(), 80) + .is_err()); + } + + #[cfg(feature = "server")] + #[test] + fn test_custom_resolver_is_used_for_redirects() -> Result<()> { + use crate::model::header::LOCATION; + use crate::Server; + use std::net::{IpAddr, Ipv4Addr, TcpListener}; + use std::sync::Mutex; + + fn unused_local_port() -> Result { + Ok(TcpListener::bind((Ipv4Addr::LOCALHOST, 0))? + .local_addr()? + .port()) + } + + let redirect_port = unused_local_port()?; + let destination_port = unused_local_port()?; + let destination_url = format!("http://destination.example:{destination_port}/"); + let _redirect_server = Server::new(move |_| { + Response::builder() + .status(StatusCode::FOUND) + .header(LOCATION, &destination_url) + .body(Body::empty()) + .unwrap() + }) + .bind((Ipv4Addr::LOCALHOST, redirect_port)) + .spawn()?; + let _destination_server = + Server::new(|_| Response::builder().body(Body::from("redirected")).unwrap()) + .bind((Ipv4Addr::LOCALHOST, destination_port)) + .spawn()?; + + let resolved_hosts = Arc::new(Mutex::new(Vec::new())); + let resolver_hosts = Arc::clone(&resolved_hosts); + let client = Client::new().with_redirection_limit(1).with_resolver( + move |uri: &Uri, default_port| { + resolver_hosts + .lock() + .unwrap() + .push(uri.host().unwrap().to_owned()); + Ok(vec![SocketAddr::new( + IpAddr::V4(Ipv4Addr::LOCALHOST), + uri.port_u16().unwrap_or(default_port), + )]) + }, + ); + let response = client.request( + Request::builder() + .uri(format!("http://source.example:{redirect_port}/")) + .body(()) + .unwrap(), + )?; + assert_eq!(response.into_body().to_string()?, "redirected"); + assert_eq!( + *resolved_hosts.lock().unwrap(), + ["source.example", "destination.example"] + ); + Ok(()) + } + #[cfg(any( feature = "rustls-aws-lc-webpki", feature = "rustls-ring-webpki", diff --git a/src/lib.rs b/src/lib.rs index b025ee2..a260bb5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,6 +20,6 @@ mod server; mod utils; #[cfg(feature = "client")] -pub use client::Client; +pub use client::{Client, Resolver}; #[cfg(feature = "server")] pub use server::{ListeningServer, Server};