Fix `WouldBlock` error handling for the handshakes

src/handshake/client.rs

@@ -9,6 +9,9 @@ use url::Url;
 
 use input_buffer::{InputBuffer, MIN_READ};
 use error::{Error, Result};
+use protocol::{
+    WebSocket, Role,
+};
 use super::{
     Headers,
     Httparse, FromHttparse,
@@ -16,9 +19,7 @@ use super::{
     convert_key,
     MAX_HEADERS,
 };
-use protocol::{
+use util::NonBlockingResult;
-    WebSocket, Role,
-};
 
 /// Client request.
 pub struct Request {
@@ -87,7 +88,7 @@ impl<Stream: Read + Write> Handshake for ClientHandshake<Stream> {
         debug!("Performing client handshake...");
         match self.state {
             HandshakeState::SendingRequest(mut req) => {
-                let size = self.stream.write(Buf::bytes(&req))?;
+                let size = self.stream.write(Buf::bytes(&req)).no_block()?.unwrap_or(0);
                 Buf::advance(&mut req, size);
                 let state = if req.has_remaining() {
                     HandshakeState::SendingRequest(req)
@@ -102,7 +103,7 @@ impl<Stream: Read + Write> Handshake for ClientHandshake<Stream> {
             HandshakeState::ReceivingResponse(mut resp_buf) => {
                 resp_buf.reserve(MIN_READ, usize::max_value())
                     .map_err(|_| Error::Capacity("Header too long".into()))?;
-                resp_buf.read_from(&mut self.stream)?;
+                resp_buf.read_from(&mut self.stream).no_block()?;
                 if let Some(resp) = Response::parse(&mut resp_buf)? {
                     self.verify_data.verify_response(&resp)?;
                     let ws = WebSocket::from_partially_read(self.stream,

src/handshake/server.rs

@@ -5,6 +5,7 @@ use httparse::Status;
 
 use input_buffer::{InputBuffer, MIN_READ};
 use error::{Error, Result};
+use protocol::{WebSocket, Role};
 use super::{
     Handshake,
     HandshakeResult,
@@ -14,7 +15,7 @@ use super::{
     convert_key,
     MAX_HEADERS
 };
-use protocol::{WebSocket, Role};
+use util::NonBlockingResult;
 
 /// Request from the client.
 pub struct Request {
@@ -91,7 +92,7 @@ impl<Stream: Read + Write> Handshake for ServerHandshake<Stream> {
             HandshakeState::ReceivingRequest(mut req_buf) => {
                 req_buf.reserve(MIN_READ, usize::max_value())
                     .map_err(|_| Error::Capacity("Header too long".into()))?;
-                req_buf.read_from(&mut self.stream)?;
+                req_buf.read_from(&mut self.stream).no_block()?;
                 let state = if let Some(req) = Request::parse(&mut req_buf)? {
                     let resp = req.reply()?;
                     HandshakeState::SendingResponse(Cursor::new(resp))
@@ -104,7 +105,7 @@ impl<Stream: Read + Write> Handshake for ServerHandshake<Stream> {
                 }))
             }
             HandshakeState::SendingResponse(mut resp) => {
-                let size = self.stream.write(Buf::bytes(&resp))?;
+                let size = self.stream.write(Buf::bytes(&resp)).no_block()?.unwrap_or(0);
                 Buf::advance(&mut resp, size);
                 if resp.has_remaining() {
                     Ok(HandshakeResult::Incomplete(ServerHandshake {