diff --git a/src/client.rs b/src/client.rs index 9301939..c24785e 100644 --- a/src/client.rs +++ b/src/client.rs @@ -178,6 +178,39 @@ where client_with_config(request, stream, None) } +/// Do the client handshake over the given stream given a web socket configuration. Passing `None` +/// as configurationis equal to calling `client_allow_blocking()` function. +/// +/// This works like `client_with_config()`, but does not return an error if the stream would block on +/// the handshake. +pub fn client_with_config_allow_blocking( + request: Req, + stream: Stream, + config: Option, +) -> StdResult<(WebSocket, Response), HandshakeError>> +where + Stream: Read + Write, + Req: IntoClientRequest, +{ + ClientHandshake::start(stream, request.into_client_request()?, config)? + .handshake_allow_blocking() +} + +/// Do the client handshake over the given stream. +/// +/// This works like `client()`, but does not return an error if the stream would block on +/// the handshake. +pub fn client_allow_blocking( + request: Req, + stream: Stream, +) -> StdResult<(WebSocket, Response), HandshakeError>> +where + Stream: Read + Write, + Req: IntoClientRequest, +{ + client_with_config_allow_blocking(request, stream, None) +} + /// Trait for converting various types into HTTP requests used for a client connection. /// /// This trait is implemented by default for string slices, strings, `url::Url`, `http::Uri` and diff --git a/src/handshake/mod.rs b/src/handshake/mod.rs index e063d4a..83f9177 100644 --- a/src/handshake/mod.rs +++ b/src/handshake/mod.rs @@ -35,13 +35,32 @@ impl MidHandshake { &mut self.machine } - /// Restarts the handshake process. - pub fn handshake(mut self) -> Result> { + /// Restarts the handshake process, but returns an error if the stream would block + pub fn handshake(self) -> Result> { + self.handshake_internal(false) + } + + /// Restarts the handshake process, but allows the stream to block + pub fn handshake_allow_blocking(self) -> Result> { + self.handshake_internal(true) + } + + fn handshake_internal( + mut self, + allow_blocking: bool, + ) -> Result> { let mut mach = self.machine; loop { mach = match mach.single_round()? { RoundResult::WouldBlock(m) => { - return Err(HandshakeError::Interrupted(MidHandshake { machine: m, ..self })) + if allow_blocking { + m + } else { + return Err(HandshakeError::Interrupted(MidHandshake { + machine: m, + ..self + })); + } } RoundResult::Incomplete(m) => m, RoundResult::StageFinished(s) => match self.role.stage_finished(s)? {