refactor: move DNS resolution into IMAP and SMTP connect code

src/smtp/connect.rs

@@ -1,13 +1,16 @@
 //! SMTP connection establishment.
 
-use anyhow::{bail, Context as _, Result};
+use std::net::SocketAddr;
+
+use anyhow::{bail, format_err, Context as _, Result};
 use async_smtp::{SmtpClient, SmtpTransport};
 use tokio::io::BufStream;
 
 use crate::context::Context;
+use crate::net::dns::{lookup_host_with_cache, update_connect_timestamp};
 use crate::net::session::SessionBufStream;
 use crate::net::tls::wrap_tls;
-use crate::net::{connect_starttls_smtp, connect_tcp, connect_tls};
+use crate::net::{connect_tcp_inner, connect_tls_inner};
 use crate::provider::Socket;
 use crate::socks::Socks5Config;
 
@@ -21,36 +24,55 @@ use crate::socks::Socks5Config;
 /// to unify the result regardless of whether TLS or STARTTLS is used.
 pub(crate) async fn connect_stream(
     context: &Context,
-    domain: &str,
+    host: &str,
     port: u16,
     strict_tls: bool,
     socks5_config: Option<Socks5Config>,
     security: Socket,
 ) -> Result<Box<dyn SessionBufStream>> {
-    let stream = if let Some(socks5_config) = socks5_config {
-        match security {
+    if let Some(socks5_config) = socks5_config {
+        let stream = match security {
             Socket::Automatic => bail!("SMTP port security is not configured"),
             Socket::Ssl => {
-                connect_secure_socks5(context, domain, port, strict_tls, socks5_config.clone())
+                connect_secure_socks5(context, host, port, strict_tls, socks5_config.clone())
                     .await?
             }
             Socket::Starttls => {
-                connect_starttls_socks5(context, domain, port, strict_tls, socks5_config.clone())
+                connect_starttls_socks5(context, host, port, strict_tls, socks5_config.clone())
                     .await?
             }
             Socket::Plain => {
-                connect_insecure_socks5(context, domain, port, socks5_config.clone()).await?
+                connect_insecure_socks5(context, host, port, socks5_config.clone()).await?
+            }
+        };
+        Ok(stream)
+    } else {
+        let mut first_error = None;
+        let load_cache = strict_tls && (security == Socket::Ssl || security == Socket::Starttls);
+
+        for resolved_addr in lookup_host_with_cache(context, host, port, load_cache).await? {
+            let res = match security {
+                Socket::Automatic => bail!("SMTP port security is not configured"),
+                Socket::Ssl => connect_secure(resolved_addr, host, strict_tls).await,
+                Socket::Starttls => connect_starttls(resolved_addr, host, strict_tls).await,
+                Socket::Plain => connect_insecure(resolved_addr).await,
+            };
+            match res {
+                Ok(stream) => {
+                    let ip_addr = resolved_addr.ip().to_string();
+                    if load_cache {
+                        update_connect_timestamp(context, host, &ip_addr).await?;
+                    }
+                    return Ok(stream);
+                }
+                Err(err) => {
+                    warn!(context, "Failed to connect to {resolved_addr}: {err:#}.");
+                    first_error.get_or_insert(err);
+                }
             }
         }
-    } else {
-        match security {
-            Socket::Automatic => bail!("SMTP port security is not configured"),
-            Socket::Ssl => connect_secure(context, domain, port, strict_tls).await?,
-            Socket::Starttls => connect_starttls(context, domain, port, strict_tls).await?,
-            Socket::Plain => connect_insecure(context, domain, port).await?,
-        }
-    };
-    Ok(stream)
+        Err(first_error.unwrap_or_else(|| format_err!("no DNS resolution results for {host}")))
+    }
 }
 
 /// Reads and ignores SMTP greeting.
@@ -132,12 +154,11 @@ async fn connect_insecure_socks5(
 }
 
 async fn connect_secure(
-    context: &Context,
+    addr: SocketAddr,
     hostname: &str,
-    port: u16,
     strict_tls: bool,
 ) -> Result<Box<dyn SessionBufStream>> {
-    let tls_stream = connect_tls(context, hostname, port, strict_tls, "smtp").await?;
+    let tls_stream = connect_tls_inner(addr, hostname, strict_tls, "smtp").await?;
     let mut buffered_stream = BufStream::new(tls_stream);
     skip_smtp_greeting(&mut buffered_stream).await?;
     let session_stream: Box<dyn SessionBufStream> = Box::new(buffered_stream);
@@ -145,24 +166,27 @@ async fn connect_secure(
 }
 
 async fn connect_starttls(
-    context: &Context,
-    hostname: &str,
-    port: u16,
+    addr: SocketAddr,
+    host: &str,
     strict_tls: bool,
 ) -> Result<Box<dyn SessionBufStream>> {
-    let tls_stream = connect_starttls_smtp(context, hostname, port, strict_tls).await?;
+    let tcp_stream = connect_tcp_inner(addr).await?;
+
+    // Run STARTTLS command and convert the client back into a stream.
+    let client = async_smtp::SmtpClient::new().smtp_utf8(true);
+    let transport = async_smtp::SmtpTransport::new(client, BufStream::new(tcp_stream)).await?;
+    let tcp_stream = transport.starttls().await?.into_inner();
+    let tls_stream = wrap_tls(strict_tls, host, "smtp", tcp_stream)
+        .await
+        .context("STARTTLS upgrade failed")?;
 
     let buffered_stream = BufStream::new(tls_stream);
     let session_stream: Box<dyn SessionBufStream> = Box::new(buffered_stream);
     Ok(session_stream)
 }
 
-async fn connect_insecure(
-    context: &Context,
-    hostname: &str,
-    port: u16,
-) -> Result<Box<dyn SessionBufStream>> {
-    let tcp_stream = connect_tcp(context, hostname, port, false).await?;
+async fn connect_insecure(addr: SocketAddr) -> Result<Box<dyn SessionBufStream>> {
+    let tcp_stream = connect_tcp_inner(addr).await?;
     let mut buffered_stream = BufStream::new(tcp_stream);
     skip_smtp_greeting(&mut buffered_stream).await?;
     let session_stream: Box<dyn SessionBufStream> = Box::new(buffered_stream);