From bd8b1265c44f7a8f53b3c86971e1d495cdcde830 Mon Sep 17 00:00:00 2001
From: Alexis Mousset <contact@amousset.me>
Date: Sun, 6 Nov 2016 22:10:46 +0100
Subject: [PATCH] feat(transport-smtp): Add timeout suppor to SMTP transport

Fixes #106
---
 src/transport/smtp/client/mod.rs | 17 +++++++++++++++--
 src/transport/smtp/client/net.rs | 26 ++++++++++++++++++++++++++
 src/transport/smtp/mod.rs        | 13 +++++++++++++
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/src/transport/smtp/client/mod.rs b/src/transport/smtp/client/mod.rs
index 64f1cf1..65ac325 100644
--- a/src/transport/smtp/client/mod.rs
+++ b/src/transport/smtp/client/mod.rs
@@ -7,9 +7,10 @@ use std::io;
 use std::io::{BufRead, Read, Write};
 use std::net::ToSocketAddrs;
 use std::string::String;
+use std::time::Duration;
 use transport::smtp::{CRLF, MESSAGE_ENDING};
 use transport::smtp::authentication::Mechanism;
-use transport::smtp::client::net::{Connector, NetworkStream};
+use transport::smtp::client::net::{Connector, NetworkStream, Timeout};
 use transport::smtp::error::{Error, SmtpResult};
 use transport::smtp::response::ResponseParser;
 
@@ -64,7 +65,7 @@ impl<S: Write + Read> Client<S> {
     }
 }
 
-impl<S: Connector + Write + Read + Debug> Client<S> {
+impl<S: Connector + Timeout + Write + Read + Debug> Client<S> {
     /// Closes the SMTP transaction if possible
     pub fn close(&mut self) {
         let _ = self.quit();
@@ -92,6 +93,18 @@ impl<S: Connector + Write + Read + Debug> Client<S> {
         }
     }
 
+    /// Set timeout
+    pub fn set_timeout(&mut self, duration: Option<Duration>) -> io::Result<()> {
+        match self.stream {
+            Some(ref mut stream) => {
+                try!(stream.get_mut().set_read_timeout(duration));
+                try!(stream.get_mut().set_read_timeout(duration));
+                Ok(())
+            }
+            None => Ok(()),
+        }
+    }
+
     /// Connects to the configured server
     pub fn connect<A: ToSocketAddrs>(&mut self,
                                      addr: &A,
diff --git a/src/transport/smtp/client/net.rs b/src/transport/smtp/client/net.rs
index 2c91185..5e49e1c 100644
--- a/src/transport/smtp/client/net.rs
+++ b/src/transport/smtp/client/net.rs
@@ -6,6 +6,7 @@ use std::fmt::{Debug, Formatter};
 use std::io;
 use std::io::{ErrorKind, Read, Write};
 use std::net::{SocketAddr, TcpStream};
+use std::time::Duration;
 
 /// A trait for the concept of opening a stream
 pub trait Connector: Sized {
@@ -97,3 +98,28 @@ impl Write for NetworkStream {
         }
     }
 }
+
+/// A trait for read and write timeout support
+pub trait Timeout: Sized {
+    /// Set read timeout for IO calls
+    fn set_read_timeout(&mut self, duration: Option<Duration>) -> io::Result<()>;
+    /// Set write timeout for IO calls
+    fn set_write_timeout(&mut self, duration: Option<Duration>) -> io::Result<()>;
+}
+
+impl Timeout for NetworkStream {
+    fn set_read_timeout(&mut self, duration: Option<Duration>) -> io::Result<()> {
+        match *self {
+            NetworkStream::Plain(ref mut stream) => stream.set_read_timeout(duration),
+            NetworkStream::Ssl(ref mut stream) => stream.get_mut().set_read_timeout(duration),
+        }
+    }
+
+    /// Set write tiemout for IO calls
+    fn set_write_timeout(&mut self, duration: Option<Duration>) -> io::Result<()> {
+        match *self {
+            NetworkStream::Plain(ref mut stream) => stream.set_write_timeout(duration),
+            NetworkStream::Ssl(ref mut stream) => stream.get_mut().set_write_timeout(duration),
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/transport/smtp/mod.rs b/src/transport/smtp/mod.rs
index 9df2ea9..84c4321 100644
--- a/src/transport/smtp/mod.rs
+++ b/src/transport/smtp/mod.rs
@@ -4,6 +4,7 @@ use email::SendableEmail;
 use openssl::ssl::{SslContext, SslMethod};
 use std::net::{SocketAddr, ToSocketAddrs};
 use std::string::String;
+use std::time::Duration;
 use transport::EmailTransport;
 use transport::smtp::authentication::Mechanism;
 use transport::smtp::client::Client;
@@ -85,6 +86,9 @@ pub struct SmtpTransportBuilder {
     smtp_utf8: bool,
     /// Optional enforced authentication mechanism
     authentication_mechanism: Option<Mechanism>,
+    /// Define network timeout
+    /// It can be changed later for specific needs (like a different timeout for each SMTP command)
+    timeout: Option<Duration>,
 }
 
 /// Builder for the SMTP `SmtpTransport`
@@ -105,6 +109,7 @@ impl SmtpTransportBuilder {
                     connection_reuse: false,
                     hello_name: "localhost".to_string(),
                     authentication_mechanism: None,
+                    timeout: Some(Duration::new(60, 0)),
                 })
             }
             None => Err(From::from("Could nor resolve hostname")),
@@ -183,6 +188,12 @@ impl SmtpTransportBuilder {
         self
     }
 
+    /// Set the timeout duration
+    pub fn timeout(mut self, timeout: Option<Duration>) -> SmtpTransportBuilder {
+        self.timeout = timeout;
+        self
+    }
+
     /// Build the SMTP client
     ///
     /// It does not connect to the server, but only creates the `SmtpTransport`
@@ -297,6 +308,8 @@ impl EmailTransport<SmtpResult> for SmtpTransport {
                                          _ => None,
                                      }));
 
+            try!(self.client.set_timeout(self.client_info.timeout));
+
             // Log the connection
             info!("connection established to {}", self.client_info.server_addr);