Set read/write timeouts for IMAP SOCKS5 streams

src/imap.rs

@@ -27,7 +27,7 @@ use crate::context::Context;
 use crate::events::EventType;
 use crate::headerdef::{HeaderDef, HeaderDefMap};
 use crate::job;
-use crate::login_param::{CertificateChecks, LoginParam, ServerAddress, ServerLoginParam};
+use crate::login_param::{CertificateChecks, LoginParam, ServerLoginParam};
 use crate::message::{self, Message, MessageState, MessengerMessage, MsgId, Viewtype};
 use crate::mimeparser;
 use crate::oauth2::get_oauth2_access_token;
@@ -305,14 +305,8 @@ impl Imap {
             let imap_port = config.lp.port;
 
             let connection = if let Some(socks5_config) = &config.socks5_config {
-                Client::connect_insecure_socks5(
+                Client::connect_insecure_socks5((imap_server, imap_port), socks5_config.clone())
-                    &ServerAddress {
+                    .await
-                        host: imap_server.to_string(),
-                        port: imap_port,
-                    },
-                    socks5_config.clone(),
-                )
-                .await
             } else {
                 Client::connect_insecure((imap_server, imap_port)).await
             };
@@ -334,17 +328,14 @@ impl Imap {
 
             if let Some(socks5_config) = &config.socks5_config {
                 Client::connect_secure_socks5(
-                    &ServerAddress {
+                    (imap_server, imap_port),
-                        host: imap_server.to_string(),
+                    imap_server,
-                        port: imap_port,
-                    },
                     config.strict_tls,
                     socks5_config.clone(),
                 )
                 .await
             } else {
-                Client::connect_secure((imap_server, imap_port), imap_server, config.strict_tls)
+                Client::connect_secure(imap_server, imap_port, config.strict_tls).await
-                    .await
             }
         };
 

src/imap/client.rs

@@ -8,7 +8,6 @@ use anyhow::{Context as _, Result};
 use async_imap::Client as ImapClient;
 use async_imap::Session as ImapSession;
 
-use async_smtp::ServerAddress;
 use tokio::net::{self, TcpStream};
 use tokio::time::timeout;
 use tokio_io_timeout::TimeoutStream;
@@ -93,12 +92,8 @@ impl Client {
         Ok(Session::new(session, capabilities))
     }
 
-    pub async fn connect_secure(
+    pub async fn connect_secure(hostname: &str, port: u16, strict_tls: bool) -> Result<Self> {
-        addr: impl net::ToSocketAddrs,
+        let tcp_stream = timeout(IMAP_TIMEOUT, TcpStream::connect((hostname, port))).await??;
-        domain: &str,
-        strict_tls: bool,
-    ) -> 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));
@@ -106,7 +101,7 @@ impl Client {
 
         let tls = build_tls(strict_tls);
         let tls_stream: Box<dyn SessionStream> =
-            Box::new(tls.connect(domain, timeout_stream).await?);
+            Box::new(tls.connect(hostname, timeout_stream).await?);
         let mut client = ImapClient::new(tls_stream);
 
         let _greeting = client
@@ -141,19 +136,17 @@ impl Client {
     }
 
     pub async fn connect_secure_socks5(
-        target_addr: &ServerAddress,
+        target_addr: impl net::ToSocketAddrs,
+        domain: &str,
         strict_tls: bool,
         socks5_config: Socks5Config,
     ) -> Result<Self> {
-        let socks5_stream: Box<dyn SessionStream> = Box::new(
+        let socks5_stream: Box<dyn SessionStream> =
-            socks5_config
+            Box::new(socks5_config.connect(target_addr, IMAP_TIMEOUT).await?);
-                .connect(target_addr, Some(IMAP_TIMEOUT))
-                .await?,
-        );
 
         let tls = build_tls(strict_tls);
         let tls_stream: Box<dyn SessionStream> =
-            Box::new(tls.connect(target_addr.host.clone(), socks5_stream).await?);
+            Box::new(tls.connect(domain, socks5_stream).await?);
         let mut client = ImapClient::new(tls_stream);
 
         let _greeting = client
@@ -168,14 +161,11 @@ impl Client {
     }
 
     pub async fn connect_insecure_socks5(
-        target_addr: &ServerAddress,
+        target_addr: impl net::ToSocketAddrs,
         socks5_config: Socks5Config,
     ) -> Result<Self> {
-        let socks5_stream: Box<dyn SessionStream> = Box::new(
+        let socks5_stream: Box<dyn SessionStream> =
-            socks5_config
+            Box::new(socks5_config.connect(target_addr, IMAP_TIMEOUT).await?);
-                .connect(target_addr, Some(IMAP_TIMEOUT))
-                .await?,
-        );
 
         let mut client = ImapClient::new(socks5_stream);
         let _greeting = client

src/imap/session.rs

@@ -48,10 +48,9 @@ impl SessionStream for Pin<Box<TimeoutStream<TcpStream>>> {
         self.as_mut().set_read_timeout_pinned(timeout);
     }
 }
-impl SessionStream for Socks5Stream<TcpStream> {
+impl SessionStream for Socks5Stream<Pin<Box<TimeoutStream<TcpStream>>>> {
-    fn set_read_timeout(&mut self, _timeout: Option<Duration>) {
+    fn set_read_timeout(&mut self, timeout: Option<Duration>) {
-        // FIXME: build SOCKS streams on top of TimeoutStream, not directly TcpStream,
+        self.get_socket_mut().set_read_timeout(timeout)
-        // so we can set a read timeout for them.
     }
 }
 

src/socks.rs

@@ -1,14 +1,18 @@
 //! # SOCKS5 support.
 
 use std::fmt;
+use std::pin::Pin;
 use std::time::Duration;
 
-use anyhow::Result;
+use anyhow::{Context as _, Result};
 pub use async_smtp::ServerAddress;
-use tokio::{io, net::TcpStream};
+use tokio::net::{self, TcpStream};
+use tokio::time::timeout;
+use tokio_io_timeout::TimeoutStream;
 
 use crate::context::Context;
-use fast_socks5::client::Socks5Stream;
+use fast_socks5::client::{Config, Socks5Stream};
+use fast_socks5::AuthenticationMethod;
 
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
 pub struct Socks5Config {
@@ -52,12 +56,32 @@ impl Socks5Config {
 
     pub async fn connect(
         &self,
-        target_addr: &ServerAddress,
+        target_addr: impl net::ToSocketAddrs,
-        timeout: Option<Duration>,
+        timeout_val: Duration,
-    ) -> io::Result<Socks5Stream<TcpStream>> {
+    ) -> Result<Socks5Stream<Pin<Box<TimeoutStream<TcpStream>>>>> {
-        self.to_async_smtp_socks5_config()
+        let tcp_stream = timeout(timeout_val, TcpStream::connect(target_addr))
-            .connect(target_addr, timeout)
             .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 authentication_method = if let Some((username, password)) = self.user_password.as_ref()
+        {
+            Some(AuthenticationMethod::Password {
+                username: username.into(),
+                password: password.into(),
+            })
+        } else {
+            None
+        };
+        let socks_stream =
+            Socks5Stream::use_stream(timeout_stream, authentication_method, Config::default())
+                .await?;
+
+        Ok(socks_stream)
     }
 
     pub fn to_async_smtp_socks5_config(&self) -> async_smtp::smtp::Socks5Config {