diff --git a/src/handshake/client.rs b/src/handshake/client.rs index 3fdb2ef..86b2c60 100644 --- a/src/handshake/client.rs +++ b/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::{ - WebSocket, Role, -}; +use util::NonBlockingResult; /// Client request. pub struct Request { @@ -87,7 +88,7 @@ impl Handshake for ClientHandshake { 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 Handshake for ClientHandshake { 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, diff --git a/src/handshake/server.rs b/src/handshake/server.rs index dcab055..f94f3b3 100644 --- a/src/handshake/server.rs +++ b/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 Handshake for ServerHandshake { 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 Handshake for ServerHandshake { })) } 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 {