From c9bd7ed852a35907da0454eea3b399209d9c59cc Mon Sep 17 00:00:00 2001
From: Alexis Mousset <alexis.mousset@rudder.io>
Date: Wed, 1 May 2019 18:23:58 +0200
Subject: [PATCH] fix(Transport): Apply timeout to TCP connection

---
 lettre/src/smtp/client/mod.rs            | 3 ++-
 lettre/src/smtp/client/net.rs            | 8 ++++++--
 lettre/src/smtp/extension.rs             | 2 +-
 lettre/src/smtp/mod.rs                   | 1 +
 website/content/sending-messages/smtp.md | 2 +-
 5 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/lettre/src/smtp/client/mod.rs b/lettre/src/smtp/client/mod.rs
index 81a7417..43e2a39 100644
--- a/lettre/src/smtp/client/mod.rs
+++ b/lettre/src/smtp/client/mod.rs
@@ -145,6 +145,7 @@ impl<S: Connector + Write + Read + Timeout + Debug> InnerClient<S> {
     pub fn connect<A: ToSocketAddrs>(
         &mut self,
         addr: &A,
+        timeout: Option<Duration>,
         tls_parameters: Option<&ClientTlsParameters>,
     ) -> SmtpResult {
         // Connect should not be called when the client is already connected
@@ -162,7 +163,7 @@ impl<S: Connector + Write + Read + Timeout + Debug> InnerClient<S> {
         debug!("connecting to {}", server_addr);
 
         // Try to connect
-        self.set_stream(Connector::connect(&server_addr, tls_parameters)?);
+        self.set_stream(Connector::connect(&server_addr, timeout, tls_parameters)?);
 
         self.read_response()
     }
diff --git a/lettre/src/smtp/client/net.rs b/lettre/src/smtp/client/net.rs
index 3d4224e..d739722 100644
--- a/lettre/src/smtp/client/net.rs
+++ b/lettre/src/smtp/client/net.rs
@@ -92,7 +92,7 @@ impl Write for NetworkStream {
 /// A trait for the concept of opening a stream
 pub trait Connector: Sized {
     /// Opens a connection to the given IP socket
-    fn connect(addr: &SocketAddr, tls_parameters: Option<&ClientTlsParameters>)
+    fn connect(addr: &SocketAddr, timeout: Option<Duration>, tls_parameters: Option<&ClientTlsParameters>)
         -> io::Result<Self>;
     /// Upgrades to TLS connection
     fn upgrade_tls(&mut self, tls_parameters: &ClientTlsParameters) -> io::Result<()>;
@@ -103,9 +103,13 @@ pub trait Connector: Sized {
 impl Connector for NetworkStream {
     fn connect(
         addr: &SocketAddr,
+        timeout: Option<Duration>,
         tls_parameters: Option<&ClientTlsParameters>,
     ) -> io::Result<NetworkStream> {
-        let tcp_stream = TcpStream::connect(addr)?;
+        let tcp_stream = match timeout {
+            Some(duration) => TcpStream::connect_timeout(addr, duration)?,
+            None => TcpStream::connect(addr)?,
+        };
 
         match tls_parameters {
             Some(context) => context
diff --git a/lettre/src/smtp/extension.rs b/lettre/src/smtp/extension.rs
index a36d0a2..20c25e4 100644
--- a/lettre/src/smtp/extension.rs
+++ b/lettre/src/smtp/extension.rs
@@ -11,7 +11,7 @@ use std::net::{Ipv4Addr, Ipv6Addr};
 use std::result::Result;
 
 /// Default client id
-pub const DEFAULT_DOMAIN_CLIENT_ID: &str = "localhost";
+const DEFAULT_DOMAIN_CLIENT_ID: &str = "localhost";
 
 /// Client identifier, the parameter to `EHLO`
 #[derive(PartialEq, Eq, Clone, Debug)]
diff --git a/lettre/src/smtp/mod.rs b/lettre/src/smtp/mod.rs
index 3ff82d5..9f7f55e 100644
--- a/lettre/src/smtp/mod.rs
+++ b/lettre/src/smtp/mod.rs
@@ -268,6 +268,7 @@ impl<'a> SmtpTransport {
 
         self.client.connect(
             &self.client_info.server_addr,
+            self.client_info.timeout,
             match self.client_info.security {
                 ClientSecurity::Wrapper(ref tls_parameters) => Some(tls_parameters),
                 _ => None,
diff --git a/website/content/sending-messages/smtp.md b/website/content/sending-messages/smtp.md
index 7337e1b..cdca2b8 100644
--- a/website/content/sending-messages/smtp.md
+++ b/website/content/sending-messages/smtp.md
@@ -163,7 +163,7 @@ use lettre::smtp::commands::*;
 
 fn main() {
     let mut email_client: InnerClient<NetworkStream> = InnerClient::new();
-    let _ = email_client.connect(&("localhost", SMTP_PORT), None);
+    let _ = email_client.connect(&("localhost", SMTP_PORT), None, None);
     let _ = email_client.command(EhloCommand::new(ClientId::new("my_hostname".to_string())));
     let _ = email_client.command(
                 MailCommand::new(Some(EmailAddress::new("user@example.com".to_string()).unwrap()), vec![])