From ebab05a6b8d0a1c90b0e287796e8be59cbcc3752 Mon Sep 17 00:00:00 2001
From: Aaron Weiss <aaronweiss74@gmail.com>
Date: Sat, 29 Nov 2014 04:52:50 -0500
Subject: [PATCH] Added the ability to set timeouts for live connections.

---
 Cargo.toml        |  2 +-
 src/conn.rs       | 42 +++++++++++++++++++++++++++++++++++++-----
 src/server/mod.rs | 34 +++++++++++++++++++++++++---------
 3 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 5588fef..85db546 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 
 name = "irc"
-version = "0.1.2"
+version = "0.2.0"
 description = "A simple, thread-safe IRC client library."
 authors = ["Aaron Weiss <aaronweiss74@gmail.com>"]
 license = "Unlicense"
diff --git a/src/conn.rs b/src/conn.rs
index e5259c7..bfb1aa7 100644
--- a/src/conn.rs
+++ b/src/conn.rs
@@ -18,7 +18,23 @@ impl Connection<BufferedStream<TcpStream>> {
     /// Creates a thread-safe TCP connection to the specified server.
     #[experimental]
     pub fn connect(host: &str, port: u16) -> IoResult<Connection<BufferedStream<NetStream>>> {
-        let socket = try!(TcpStream::connect(format!("{}:{}", host, port)[]));
+        Connection::connect_internal(host, port, None)
+    }
+
+    /// Creates a thread-safe TCP connection to the specified server with a given timeout in 
+    /// milliseconds.
+    #[experimental]
+    pub fn connect_with_timeout(host: &str, port: u16, timeout_ms: u64) 
+        -> IoResult<Connection<BufferedStream<NetStream>>> {   
+        Connection::connect_internal(host, port, Some(timeout_ms))
+    }
+
+    /// Creates a thread-safe TCP connection with an optional timeout.
+    #[experimental]
+    fn connect_internal(host: &str, port: u16, timeout_ms: Option<u64>) 
+    -> IoResult<Connection<BufferedStream<NetStream>>> {  
+        let mut socket = try!(TcpStream::connect(format!("{}:{}", host, port)[]));
+        socket.set_timeout(timeout_ms);
         Ok(Connection::new(BufferedStream::new(NetStream::UnsecuredTcpStream(socket))))
     }
 
@@ -27,10 +43,14 @@ impl Connection<BufferedStream<TcpStream>> {
     #[experimental]
     #[cfg(feature = "ssl")]
     pub fn connect_ssl(host: &str, port: u16) -> IoResult<Connection<BufferedStream<NetStream>>> {
-        let socket = try!(TcpStream::connect(format!("{}:{}", host, port)[]));
-        let ssl = try!(ssl_to_io(SslContext::new(Tlsv1)));
-        let ssl_socket = try!(ssl_to_io(SslStream::new(&ssl, socket)));
-        Ok(Connection::new(BufferedStream::new(NetStream::SslTcpStream(ssl_socket))))
+        Connection::connect_ssl_internal(host, port, None)
+    }
+
+    #[experimental]
+    #[cfg(feature = "ssl")]
+    pub fn connect_ssl_with_timeout(host: &str, port: u16, timeout_ms: u64)
+        -> IoResult<Connection<BufferedStream<NetStream>>> {
+        Connection::connect_ssl_internal(host, port, Some(timeout_ms))
     }
 
     /// Creates a thread-safe TCP connection to the specified server over SSL.
@@ -40,6 +60,18 @@ impl Connection<BufferedStream<TcpStream>> {
     pub fn connect_ssl(host: &str, port: u16) -> IoResult<Connection<BufferedStream<NetStream>>> {
         panic!("Cannot connect to {}:{} over SSL without compiling with SSL support.", host, port)
     }
+
+    /// Creates a thread-safe TCP connection over SSL with an optional timeout.
+    #[experimental]
+    #[cfg(feature = "ssl")]
+    fn connect_ssl_internal(host: &str, port: u16, timeout_ms: Option<u64>)
+    -> IoResult<Connection<BufferedStream<NetStream>>> {
+        let mut socket = try!(TcpStream::connect(format!("{}:{}", host, port)[]));
+        socket.set_timeout(timeout_ms);
+        let ssl = try!(ssl_to_io(SslContext::new(Tlsv1)));
+        let ssl_socket = try!(ssl_to_io(SslStream::new(&ssl, socket)));
+        Ok(Connection::new(BufferedStream::new(NetStream::SslTcpStream(ssl_socket))))
+    }
 }
 
 /// Converts a Result<T, SslError> into an IoResult<T>.
diff --git a/src/server/mod.rs b/src/server/mod.rs
index a27a2dc..bb34990 100644
--- a/src/server/mod.rs
+++ b/src/server/mod.rs
@@ -35,17 +35,19 @@ pub struct IrcServer<T> where T: IrcStream {
 }
 
 impl IrcServer<BufferedStream<NetStream>> {
-    /// Creates a new IRC Server connection from the configuration at the specified path, connecting
-    /// immediately.
+    /// Creates a new IRC Server connection from the configuration at the specified path, 
+    /// connecting immediately.
     #[experimental]
     pub fn new(config: &str) -> IoResult<IrcServer<BufferedStream<NetStream>>> {
-        let config = try!(Config::load_utf8(config));
-        let conn = try!(if config.use_ssl {
-            Connection::connect_ssl(config.server[], config.port)
-        } else {
-            Connection::connect(config.server[], config.port)
-        });
-        Ok(IrcServer { config: config, conn: conn, chanlists: Mutex::new(HashMap::new()) })
+        IrcServer::from_config(try!(Config::load_utf8(config)))
+    }
+
+    /// Creates a new IRC server connection from the configuration at the specified path with the
+    /// specified timeout in milliseconds, connecting immediately.
+    #[experimental]
+    pub fn with_timeout(config: &str, timeout_ms: u64) 
+        -> IoResult<IrcServer<BufferedStream<NetStream>>> {
+        IrcServer::from_config_with_timeout(try!(Config::load_utf8(config)), timeout_ms)    
     }
 
     /// Creates a new IRC server connection from the specified configuration, connecting
@@ -59,6 +61,20 @@ impl IrcServer<BufferedStream<NetStream>> {
         });
         Ok(IrcServer { config: config, conn: conn, chanlists: Mutex::new(HashMap::new()) })
     }
+
+    /// Creates a new IRC server connection from the specified configuration with the specified 
+    /// timeout in milliseconds, connecting immediately.
+    #[experimental]
+    pub fn from_config_with_timeout(config: Config, timeout_ms: u64) 
+        -> IoResult<IrcServer<BufferedStream<NetStream>>> {
+        let conn = try!(if config.use_ssl {
+            Connection::connect_ssl_with_timeout(config.server[], config.port, timeout_ms)
+        } else {
+            Connection::connect_with_timeout(config.server[], config.port, timeout_ms)
+        });
+        Ok(IrcServer { config: config, conn: conn, chanlists: Mutex::new(HashMap::new()) })
+    }
+
 }
 
 impl<'a, T> Server<'a, T> for IrcServer<T> where T: IrcStream {