Set read/write timeouts for IMAP SOCKS5 streams

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,
-        timeout: Option<Duration>,
-    ) -> io::Result<Socks5Stream<TcpStream>> {
-        self.to_async_smtp_socks5_config()
-            .connect(target_addr, timeout)
+        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 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 {