Add dns feature for DNS stuff (not just in starttls)

tokio-xmpp/src/connect.rs

@@ -1,7 +1,17 @@
 //! `ServerConnector` provides streams for XMPP clients
 
+#[cfg(feature = "dns")]
+use futures::{future::select_ok, FutureExt};
+#[cfg(feature = "dns")]
+use hickory_resolver::{
+    config::LookupIpStrategy, name_server::TokioConnectionProvider, IntoName, TokioAsyncResolver,
+};
+#[cfg(feature = "dns")]
+use log::debug;
 use sasl::common::ChannelBinding;
+use std::net::{IpAddr, SocketAddr};
 use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::net::TcpStream;
 use xmpp_parsers::jid::Jid;
 
 use crate::xmpp_stream::XMPPStream;
@@ -32,3 +42,78 @@ pub trait ServerConnector: Clone + core::fmt::Debug + Send + Unpin + 'static {
         Ok(ChannelBinding::None)
     }
 }
+
+/// A simple wrapper to build [`TcpStream`]
+pub struct Tcp;
+
+impl Tcp {
+    /// Connect directly to an IP/Port combo
+    pub async fn connect(ip: IpAddr, port: u16) -> Result<TcpStream, Error> {
+        Ok(TcpStream::connect(&SocketAddr::new(ip, port)).await?)
+    }
+
+    /// Connect over TCP, resolving A/AAAA records (happy eyeballs)
+    #[cfg(feature = "dns")]
+    pub async fn resolve(domain: &str, port: u16) -> Result<TcpStream, Error> {
+        let ascii_domain = idna::domain_to_ascii(&domain)?;
+
+        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 = TokioAsyncResolver::new(config, options, TokioConnectionProvider::default());
+
+        let ips = resolver.lookup_ip(ascii_domain).await?;
+
+        // Happy Eyeballs: connect to all records in parallel, return the
+        // first to succeed
+        select_ok(
+            ips.into_iter()
+                .map(|ip| TcpStream::connect(SocketAddr::new(ip, port)).boxed()),
+        )
+        .await
+        .map(|(result, _)| result)
+        .map_err(|_| Error::Disconnected)
+    }
+
+    /// Connect over TCP, resolving SRV records
+    #[cfg(feature = "dns")]
+    pub async fn resolve_with_srv(
+        domain: &str,
+        srv: &str,
+        fallback_port: u16,
+    ) -> Result<TcpStream, Error> {
+        let ascii_domain = idna::domain_to_ascii(&domain)?;
+
+        if let Ok(ip) = ascii_domain.parse() {
+            debug!("Attempting connection to {ip}:{fallback_port}");
+            return Ok(TcpStream::connect(&SocketAddr::new(ip, fallback_port)).await?);
+        }
+
+        let resolver = TokioAsyncResolver::tokio_from_system_conf()?;
+
+        let srv_domain = format!("{}.{}.", srv, ascii_domain).into_name()?;
+        let srv_records = resolver.srv_lookup(srv_domain.clone()).await.ok();
+
+        match srv_records {
+            Some(lookup) => {
+                // TODO: sort lookup records by priority/weight
+                for srv in lookup.iter() {
+                    debug!("Attempting connection to {srv_domain} {srv}");
+                    match Self::resolve(&srv.target().to_ascii(), srv.port()).await {
+                        Ok(stream) => return Ok(stream),
+                        Err(_) => {}
+                    }
+                }
+                Err(Error::Disconnected)
+            }
+            None => {
+                // SRV lookup error, retry with hostname
+                debug!("Attempting connection to {domain}:{fallback_port}");
+                Self::resolve(domain, fallback_port).await
+            }
+        }
+    }
+}