From 8d1de6c497baf8a84ee830c3e080d65383e9f498 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Fri, 9 Mar 2018 13:12:14 -0800
Subject: [PATCH] ws client timeouts

---
 src/ws/client.rs | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/ws/client.rs b/src/ws/client.rs
index 80169b32..cb406dbe 100644
--- a/src/ws/client.rs
+++ b/src/ws/client.rs
@@ -2,17 +2,18 @@
 use std::{fmt, io, str};
 use std::rc::Rc;
 use std::cell::UnsafeCell;
+use std::time::Duration;
 
 use base64;
 use rand;
 use bytes::Bytes;
 use cookie::Cookie;
+use byteorder::{ByteOrder, NetworkEndian};
 use http::{HttpTryFrom, StatusCode, Error as HttpError};
 use http::header::{self, HeaderName, HeaderValue};
 use sha1::Sha1;
 use futures::{Async, Future, Poll, Stream};
 use futures::unsync::mpsc::{unbounded, UnboundedSender};
-use byteorder::{ByteOrder, NetworkEndian};
 
 use actix::prelude::*;
 
@@ -299,9 +300,32 @@ impl ClientHandshake {
             request: None,
             tx: None,
             error: Some(err),
-            max_size: 0
+            max_size: 0,
         }
     }
+
+    /// Set handshake timeout
+    ///
+    /// Handshake timeout is a total time before handshake should be completed.
+    /// Default value is 5 seconds.
+    pub fn timeout(mut self, timeout: Duration) -> Self {
+        if let Some(request) = self.request.take() {
+            self.request = Some(request.timeout(timeout));
+        }
+        self
+    }
+
+    /// Set connection timeout
+    ///
+    /// Connection timeout includes resolving hostname and actual connection to
+    /// the host.
+    /// Default value is 1 second.
+    pub fn conn_timeout(mut self, timeout: Duration) -> Self {
+        if let Some(request) = self.request.take() {
+            self.request = Some(request.conn_timeout(timeout));
+        }
+        self
+    }
 }
 
 impl Future for ClientHandshake {