Buffer IMAP client writes

async-imap does not do its own buffering, but calls flush() after
sending each command. Using BufWriter reduces the number of write()
system calls used to send a single command.

Note that BufWriter is set up on top of TLS streams, because
we can't guarantee that TLS libraries flush the stream before
waiting for response.

CHANGELOG.md

@@ -9,6 +9,7 @@
 - cargo: bump quick-xml from 0.23.0 to 0.26.0 #3722
 - Add fuzzing tests #3853
 - Add mappings for some file types to Viewtype / MIME type #3881
+- Buffer IMAP client writes #3888
 
 ### API-Changes
 - jsonrpc: add python API for webxdc updates #3872

src/imap.rs

@@ -304,22 +304,23 @@ impl Imap {
             let imap_server: &str = config.lp.server.as_ref();
             let imap_port = config.lp.port;
 
-            let connection = if let Some(socks5_config) = &config.socks5_config {
-                Client::connect_insecure_socks5((imap_server, imap_port), socks5_config.clone())
+            if let Some(socks5_config) = &config.socks5_config {
+                if config.lp.security == Socket::Starttls {
+                    Client::connect_starttls_socks5(
+                        imap_server,
+                        imap_port,
+                        socks5_config.clone(),
+                        config.strict_tls,
+                    )
                     .await
+                } else {
+                    Client::connect_insecure_socks5((imap_server, imap_port), socks5_config.clone())
+                        .await
+                }
+            } else if config.lp.security == Socket::Starttls {
+                Client::connect_starttls(imap_server, imap_port, config.strict_tls).await
             } else {
                 Client::connect_insecure((imap_server, imap_port)).await
-            };
-
-            match connection {
-                Ok(client) => {
-                    if config.lp.security == Socket::Starttls {
-                        client.secure(imap_server, config.strict_tls).await
-                    } else {
-                        Ok(client)
-                    }
-                }
-                Err(err) => Err(err),
             }
         } else {
             let config = &self.config;
@@ -328,8 +329,8 @@ impl Imap {
 
             if let Some(socks5_config) = &config.socks5_config {
                 Client::connect_secure_socks5(
-                    (imap_server, imap_port),
                     imap_server,
+                    imap_port,
                     config.strict_tls,
                     socks5_config.clone(),
                 )

src/imap/client.rs

@@ -8,13 +8,13 @@ use anyhow::{Context as _, Result};
 use async_imap::Client as ImapClient;
 use async_imap::Session as ImapSession;
 
-use tokio::net::{self, TcpStream};
-use tokio::time::timeout;
-use tokio_io_timeout::TimeoutStream;
+use tokio::io::BufWriter;
+use tokio::net::ToSocketAddrs;
 
 use super::capabilities::Capabilities;
 use super::session::Session;
 use crate::login_param::build_tls;
+use crate::net::connect_tcp;
 use crate::socks::Socks5Config;
 
 use super::session::SessionStream;
@@ -24,7 +24,6 @@ pub(crate) const IMAP_TIMEOUT: Duration = Duration::from_secs(30);
 
 #[derive(Debug)]
 pub(crate) struct Client {
-    is_secure: bool,
     inner: ImapClient<Box<dyn SessionStream>>,
 }
 
@@ -93,108 +92,104 @@ impl Client {
     }
 
     pub async fn connect_secure(hostname: &str, port: u16, strict_tls: bool) -> Result<Self> {
-        let tcp_stream = timeout(IMAP_TIMEOUT, TcpStream::connect((hostname, port))).await??;
-        let mut timeout_stream = TimeoutStream::new(tcp_stream);
-        timeout_stream.set_write_timeout(Some(IMAP_TIMEOUT));
-        timeout_stream.set_read_timeout(Some(IMAP_TIMEOUT));
-        let timeout_stream = Box::pin(timeout_stream);
-
+        let tcp_stream = connect_tcp((hostname, port), IMAP_TIMEOUT).await?;
         let tls = build_tls(strict_tls);
-        let tls_stream: Box<dyn SessionStream> =
-            Box::new(tls.connect(hostname, timeout_stream).await?);
-        let mut client = ImapClient::new(tls_stream);
+        let tls_stream = tls.connect(hostname, tcp_stream).await?;
+        let buffered_stream = BufWriter::new(tls_stream);
+        let session_stream: Box<dyn SessionStream> = Box::new(buffered_stream);
+        let mut client = ImapClient::new(session_stream);
 
         let _greeting = client
             .read_response()
             .await
             .context("failed to read greeting")?;
 
-        Ok(Client {
-            is_secure: true,
-            inner: client,
-        })
+        Ok(Client { inner: client })
     }
 
-    pub async fn connect_insecure(addr: impl net::ToSocketAddrs) -> Result<Self> {
-        let tcp_stream = timeout(IMAP_TIMEOUT, TcpStream::connect(addr)).await??;
-        let mut timeout_stream = TimeoutStream::new(tcp_stream);
-        timeout_stream.set_write_timeout(Some(IMAP_TIMEOUT));
-        timeout_stream.set_read_timeout(Some(IMAP_TIMEOUT));
-        let timeout_stream = Box::pin(timeout_stream);
-        let stream: Box<dyn SessionStream> = Box::new(timeout_stream);
-
-        let mut client = ImapClient::new(stream);
+    pub async fn connect_insecure(addr: impl ToSocketAddrs) -> Result<Self> {
+        let tcp_stream = connect_tcp(addr, IMAP_TIMEOUT).await?;
+        let buffered_stream = BufWriter::new(tcp_stream);
+        let session_stream: Box<dyn SessionStream> = Box::new(buffered_stream);
+        let mut client = ImapClient::new(session_stream);
         let _greeting = client
             .read_response()
             .await
             .context("failed to read greeting")?;
 
-        Ok(Client {
-            is_secure: false,
-            inner: client,
-        })
+        Ok(Client { inner: client })
+    }
+
+    pub async fn connect_starttls(hostname: &str, port: u16, strict_tls: bool) -> Result<Self> {
+        let tcp_stream = connect_tcp((hostname, port), IMAP_TIMEOUT).await?;
+        let tls = build_tls(strict_tls);
+        let tls_stream = tls.connect(hostname, tcp_stream).await?;
+        let buffered_stream = BufWriter::new(tls_stream);
+        let session_stream: Box<dyn SessionStream> = Box::new(buffered_stream);
+        let mut client = ImapClient::new(session_stream);
+        let _greeting = client
+            .read_response()
+            .await
+            .context("failed to read greeting")?;
+
+        Ok(Client { inner: client })
     }
 
     pub async fn connect_secure_socks5(
-        target_addr: impl net::ToSocketAddrs,
         domain: &str,
+        port: u16,
         strict_tls: bool,
         socks5_config: Socks5Config,
     ) -> Result<Self> {
-        let socks5_stream: Box<dyn SessionStream> =
-            Box::new(socks5_config.connect(target_addr, IMAP_TIMEOUT).await?);
-
+        let socks5_stream = socks5_config.connect((domain, port), IMAP_TIMEOUT).await?;
         let tls = build_tls(strict_tls);
-        let tls_stream: Box<dyn SessionStream> =
-            Box::new(tls.connect(domain, socks5_stream).await?);
-        let mut client = ImapClient::new(tls_stream);
-
+        let tls_stream = tls.connect(domain, socks5_stream).await?;
+        let buffered_stream = BufWriter::new(tls_stream);
+        let session_stream: Box<dyn SessionStream> = Box::new(buffered_stream);
+        let mut client = ImapClient::new(session_stream);
         let _greeting = client
             .read_response()
             .await
             .context("failed to read greeting")?;
 
-        Ok(Client {
-            is_secure: true,
-            inner: client,
-        })
+        Ok(Client { inner: client })
     }
 
     pub async fn connect_insecure_socks5(
-        target_addr: impl net::ToSocketAddrs,
+        target_addr: impl ToSocketAddrs,
         socks5_config: Socks5Config,
     ) -> Result<Self> {
-        let socks5_stream: Box<dyn SessionStream> =
-            Box::new(socks5_config.connect(target_addr, IMAP_TIMEOUT).await?);
-
-        let mut client = ImapClient::new(socks5_stream);
+        let socks5_stream = socks5_config.connect(target_addr, IMAP_TIMEOUT).await?;
+        let buffered_stream = BufWriter::new(socks5_stream);
+        let session_stream: Box<dyn SessionStream> = Box::new(buffered_stream);
+        let mut client = ImapClient::new(session_stream);
         let _greeting = client
             .read_response()
             .await
             .context("failed to read greeting")?;
 
-        Ok(Client {
-            is_secure: false,
-            inner: client,
-        })
+        Ok(Client { inner: client })
     }
 
-    pub async fn secure(self, domain: &str, strict_tls: bool) -> Result<Self> {
-        if self.is_secure {
-            Ok(self)
-        } else {
-            let Client { mut inner, .. } = self;
-            let tls = build_tls(strict_tls);
-            inner.run_command_and_check_ok("STARTTLS", None).await?;
+    pub async fn connect_starttls_socks5(
+        hostname: &str,
+        port: u16,
+        socks5_config: Socks5Config,
+        strict_tls: bool,
+    ) -> Result<Self> {
+        let socks5_stream = socks5_config
+            .connect((hostname, port), IMAP_TIMEOUT)
+            .await?;
+        let tls = build_tls(strict_tls);
+        let tls_stream = tls.connect(hostname, socks5_stream).await?;
+        let buffered_stream = BufWriter::new(tls_stream);
+        let session_stream: Box<dyn SessionStream> = Box::new(buffered_stream);
+        let mut client = ImapClient::new(session_stream);
+        let _greeting = client
+            .read_response()
+            .await
+            .context("failed to read greeting")?;
 
-            let stream = inner.into_inner();
-            let ssl_stream = tls.connect(domain, stream).await?;
-            let boxed: Box<dyn SessionStream> = Box::new(ssl_stream);
-
-            Ok(Client {
-                is_secure: true,
-                inner: ImapClient::new(boxed),
-            })
-        }
+        Ok(Client { inner: client })
     }
 }

src/imap/session.rs

@@ -6,6 +6,7 @@ use async_imap::types::Mailbox;
 use async_imap::Session as ImapSession;
 use async_native_tls::TlsStream;
 use fast_socks5::client::Socks5Stream;
+use tokio::io::BufWriter;
 use tokio::net::TcpStream;
 use tokio_io_timeout::TimeoutStream;
 
@@ -33,12 +34,17 @@ pub(crate) trait SessionStream:
     fn set_read_timeout(&mut self, timeout: Option<Duration>);
 }
 
-impl SessionStream for TlsStream<Box<dyn SessionStream>> {
+impl SessionStream for Box<dyn SessionStream> {
+    fn set_read_timeout(&mut self, timeout: Option<Duration>) {
+        self.as_mut().set_read_timeout(timeout);
+    }
+}
+impl<T: SessionStream> SessionStream for TlsStream<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_mut().set_read_timeout(timeout);
     }
 }
-impl SessionStream for TlsStream<Pin<Box<TimeoutStream<TcpStream>>>> {
+impl<T: SessionStream> SessionStream for BufWriter<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_mut().set_read_timeout(timeout);
     }
@@ -48,7 +54,7 @@ impl SessionStream for Pin<Box<TimeoutStream<TcpStream>>> {
         self.as_mut().set_read_timeout_pinned(timeout);
     }
 }
-impl SessionStream for Socks5Stream<Pin<Box<TimeoutStream<TcpStream>>>> {
+impl<T: SessionStream> SessionStream for Socks5Stream<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_socket_mut().set_read_timeout(timeout)
     }

src/lib.rs

@@ -100,6 +100,7 @@ mod dehtml;
 mod authres;
 mod color;
 pub mod html;
+mod net;
 pub mod plaintext;
 mod ratelimit;
 pub mod summary;

src/net.rs

@@ -0,0 +1,26 @@
+///! # Common network utilities.
+use std::pin::Pin;
+use std::time::Duration;
+
+use anyhow::{Context as _, Result};
+use tokio::net::{TcpStream, ToSocketAddrs};
+use tokio::time::timeout;
+use tokio_io_timeout::TimeoutStream;
+
+/// Returns a TCP connection with read/write timeouts set.
+pub(crate) async fn connect_tcp(
+    addr: impl ToSocketAddrs,
+    timeout_val: Duration,
+) -> Result<Pin<Box<TimeoutStream<TcpStream>>>> {
+    let tcp_stream = timeout(timeout_val, TcpStream::connect(addr))
+        .await
+        .context("connection timeout")?
+        .context("connection failure")?;
+
+    let mut timeout_stream = TimeoutStream::new(tcp_stream);
+    timeout_stream.set_write_timeout(Some(timeout_val));
+    timeout_stream.set_read_timeout(Some(timeout_val));
+    let pinned_stream = Box::pin(timeout_stream);
+
+    Ok(pinned_stream)
+}

src/socks.rs

@@ -4,10 +4,10 @@ use std::fmt;
 use std::pin::Pin;
 use std::time::Duration;
 
-use anyhow::{Context as _, Result};
+use crate::net::connect_tcp;
+use anyhow::Result;
 pub use async_smtp::ServerAddress;
 use tokio::net::{self, TcpStream};
-use tokio::time::timeout;
 use tokio_io_timeout::TimeoutStream;
 
 use crate::context::Context;
@@ -59,14 +59,7 @@ impl Socks5Config {
         target_addr: impl net::ToSocketAddrs,
         timeout_val: Duration,
     ) -> Result<Socks5Stream<Pin<Box<TimeoutStream<TcpStream>>>>> {
-        let tcp_stream = timeout(timeout_val, TcpStream::connect(target_addr))
-            .await
-            .context("connection timeout")?
-            .context("connection failure")?;
-        let mut timeout_stream = TimeoutStream::new(tcp_stream);
-        timeout_stream.set_write_timeout(Some(timeout_val));
-        timeout_stream.set_read_timeout(Some(timeout_val));
-        let timeout_stream = Box::pin(timeout_stream);
+        let tcp_stream = connect_tcp(target_addr, timeout_val).await?;
 
         let authentication_method = if let Some((username, password)) = self.user_password.as_ref()
         {
@@ -78,8 +71,7 @@ impl Socks5Config {
             None
         };
         let socks_stream =
-            Socks5Stream::use_stream(timeout_stream, authentication_method, Config::default())
-                .await?;
+            Socks5Stream::use_stream(tcp_stream, authentication_method, Config::default()).await?;
 
         Ok(socks_stream)
     }