From 5d4551557846a8d9cbd740eaf95cfdc16bb5c3c7 Mon Sep 17 00:00:00 2001
From: Aaron Weiss <aaronweiss74@gmail.com>
Date: Sat, 20 Dec 2014 17:53:27 -0500
Subject: [PATCH] Added set_timeout to Connection.

---
 src/conn.rs | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/conn.rs b/src/conn.rs
index 4a96075..c71ec83 100644
--- a/src/conn.rs
+++ b/src/conn.rs
@@ -88,10 +88,21 @@ impl Connection<BufferedReader<NetStream>, BufferedWriter<NetStream>> {
     /// Sets the keepalive for the network stream.
     #[experimental]
     pub fn set_keepalive(&self, delay_in_seconds: Option<uint>) -> IoResult<()> {
+        self.mod_stream(|tcp| tcp.set_keepalive(delay_in_seconds))
+    }
+    
+    /// Sets the timeout for the network stream.
+    #[experimental]
+    pub fn set_timeout(&self, timeout_ms: Option<u64>) {
+        self.mod_stream(|tcp| Ok(tcp.set_timeout(timeout_ms))).unwrap(); // this cannot fail.
+    }
+
+    /// Modifies the internal TcpStream using a function.
+    fn mod_stream(&self, f: |&mut TcpStream| -> IoResult<()>) -> IoResult<()> {
         match self.reader.lock().get_mut() {
-            &NetStream::UnsecuredTcpStream(ref mut tcp) => tcp.set_keepalive(delay_in_seconds),
+            &NetStream::UnsecuredTcpStream(ref mut tcp) => f(tcp),
             #[cfg(feature = "ssl")]
-            &NetStream::SslTcpStream(ref mut ssl) => ssl.get_mut().set_keepalive(delay_in_seconds),
+            &NetStream::SslTcpStream(ref mut ssl) => f(ssl.get_mut()),
         }
     }
 }