From 559159d4585c9a82e359d2b2305bac0b44efc2f4 Mon Sep 17 00:00:00 2001 From: famfo Date: Thu, 21 May 2026 13:57:26 +0200 Subject: [PATCH] tokio-xmpp: add option to set custom hickory resolver --- tokio-xmpp/ChangeLog | 2 + tokio-xmpp/examples/keep_connection.rs | 1 + tokio-xmpp/src/connect/dns.rs | 79 ++++++++++++++++++++------ 3 files changed, 64 insertions(+), 18 deletions(-) diff --git a/tokio-xmpp/ChangeLog b/tokio-xmpp/ChangeLog index d9316581..11ebd407 100644 --- a/tokio-xmpp/ChangeLog +++ b/tokio-xmpp/ChangeLog @@ -12,6 +12,8 @@ Version NEXT: - Expose `client_auth` method to allow manual stream setups for advanced use cases. - Add feature flag to disable validation in xmpp_parsers + - Added option to set custom DNS resolver to + `tokio_xmpp::connect::DnsConfig` * Fixed: - Ignore missing "version" stream attribute for 0114 components. - Gate `AsRawFd` behind `ktls` feature to make Windows build work again. diff --git a/tokio-xmpp/examples/keep_connection.rs b/tokio-xmpp/examples/keep_connection.rs index c8b2b29a..2cbe63e8 100644 --- a/tokio-xmpp/examples/keep_connection.rs +++ b/tokio-xmpp/examples/keep_connection.rs @@ -66,6 +66,7 @@ async fn main() { host: jid.domain().as_str().to_owned(), srv: "_xmpp-client._tcp".to_owned(), fallback_port: 5222, + resolver: None, }), jid.clone().into(), password.clone(), diff --git a/tokio-xmpp/src/connect/dns.rs b/tokio-xmpp/src/connect/dns.rs index 716ce714..dbe0ff75 100644 --- a/tokio-xmpp/src/connect/dns.rs +++ b/tokio-xmpp/src/connect/dns.rs @@ -23,6 +23,8 @@ pub enum DnsConfig { srv: String, /// When SRV resolution fails what port to use fallback_port: u16, + /// Pre-configured DNS resolver + resolver: Option, }, /// Manually define server host and port @@ -33,6 +35,8 @@ pub enum DnsConfig { host: String, /// Server port port: u16, + /// Pre-configured DNS resolver + resolver: Option, }, /// Manually define IP: port (TODO: socket) @@ -49,7 +53,7 @@ impl fmt::Display for DnsConfig { #[cfg(feature = "dns")] Self::UseSrv { host, .. } => write!(f, "{}", host), #[cfg(feature = "dns")] - Self::NoSrv { host, port } => write!(f, "{}:{}", host, port), + Self::NoSrv { host, port, .. } => write!(f, "{}:{}", host, port), Self::Addr { addr } => write!(f, "{}", addr), } } @@ -63,6 +67,7 @@ impl DnsConfig { host: host.to_string(), srv: srv.to_string(), fallback_port, + resolver: None, } } @@ -73,6 +78,7 @@ impl DnsConfig { host: host.to_string(), srv: "_xmpp-client._tcp".to_string(), fallback_port: 5222, + resolver: None, } } @@ -83,6 +89,7 @@ impl DnsConfig { host: host.to_string(), srv: "_xmpps-client._tcp".to_string(), fallback_port: 5223, + resolver: None, } } @@ -92,6 +99,7 @@ impl DnsConfig { Self::NoSrv { host: host.to_string(), port, + resolver: None, } } @@ -102,6 +110,20 @@ impl DnsConfig { } } + /// Set pre-configured DNS resolver + #[cfg(feature = "dns")] + pub fn with_resolver(&mut self, custom_resolver: TokioResolver) { + match self { + Self::UseSrv { + ref mut resolver, .. + } => *resolver = Some(custom_resolver), + Self::NoSrv { + ref mut resolver, .. + } => *resolver = Some(custom_resolver), + Self::Addr { .. } => {} + } + } + /// Try resolve the DnsConfig to a TcpStream pub async fn resolve(&self) -> Result { match self { @@ -110,9 +132,14 @@ impl DnsConfig { host, srv, fallback_port, - } => Self::resolve_srv(host, srv, *fallback_port).await, + resolver, + } => Self::resolve_srv(host, srv, *fallback_port, resolver).await, #[cfg(feature = "dns")] - Self::NoSrv { host, port } => Self::resolve_no_srv(host, *port).await, + Self::NoSrv { + host, + port, + resolver, + } => Self::resolve_no_srv(host, *port, resolver).await, Self::Addr { addr } => { // TODO: Unix domain socket let addr: SocketAddr = addr.parse()?; @@ -122,7 +149,12 @@ impl DnsConfig { } #[cfg(feature = "dns")] - async fn resolve_srv(host: &str, srv: &str, fallback_port: u16) -> Result { + async fn resolve_srv( + host: &str, + srv: &str, + fallback_port: u16, + resolver: &Option, + ) -> Result { let ascii_domain = idna::domain_to_ascii(host)?; if let Ok(ip) = ascii_domain.parse() { @@ -130,14 +162,11 @@ impl DnsConfig { return Ok(TcpStream::connect(&SocketAddr::new(ip, fallback_port)).await?); } - let (_config, options) = hickory_resolver::system_conf::read_system_conf()?; - let resolver = TokioResolver::builder_tokio()? - .with_options(options) - .build()?; - + let resolver = Self::new_resolver(resolver)?; let srv_domain = format!("{}.{}.", srv, ascii_domain).into_name()?; let srv_records = resolver.srv_lookup(srv_domain.clone()).await.ok(); + let resolver_ref = Some(resolver); match srv_records { Some(lookup) => { // TODO: sort lookup records by priority/weight @@ -147,7 +176,8 @@ impl DnsConfig { }; debug!("Attempting connection to {srv_domain} {srv}"); - if let Ok(stream) = Self::resolve_no_srv(&srv.target.to_ascii(), srv.port).await + if let Ok(stream) = + Self::resolve_no_srv(&srv.target.to_ascii(), srv.port, &resolver_ref).await { return Ok(stream); } @@ -157,25 +187,24 @@ impl DnsConfig { None => { // SRV lookup error, retry with hostname debug!("Attempting connection to {host}:{fallback_port}"); - Self::resolve_no_srv(host, fallback_port).await + Self::resolve_no_srv(host, fallback_port, &resolver_ref).await } } } #[cfg(feature = "dns")] - async fn resolve_no_srv(host: &str, port: u16) -> Result { + async fn resolve_no_srv( + host: &str, + port: u16, + resolver: &Option, + ) -> Result { let ascii_domain = idna::domain_to_ascii(host)?; if let Ok(ip) = ascii_domain.parse() { return Ok(TcpStream::connect(&SocketAddr::new(ip, port)).await?); } - let (_config, mut options) = hickory_resolver::system_conf::read_system_conf()?; - options.ip_strategy = LookupIpStrategy::Ipv4AndIpv6; - let resolver = TokioResolver::builder_tokio()? - .with_options(options) - .build()?; - + let resolver = Self::new_resolver(resolver)?; let ips = resolver.lookup_ip(ascii_domain).await?; // Happy Eyeballs: connect to all records in parallel, return the @@ -188,4 +217,18 @@ impl DnsConfig { .map(|(result, _)| result) .map_err(|_| Error::Disconnected) } + + #[cfg(feature = "dns")] + fn new_resolver(resolver: &Option) -> Result { + if let Some(resolver) = resolver { + return Ok(resolver.clone()); + } + + let (_config, mut options) = hickory_resolver::system_conf::read_system_conf()?; + options.ip_strategy = LookupIpStrategy::Ipv4AndIpv6; + + Ok(TokioResolver::builder_tokio()? + .with_options(options) + .build()?) + } }