diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 3997803..7627b46 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -378,7 +378,11 @@ impl WebSocketContext { where Stream: Read + Write, { - if let Some(mut frame) = self.frame.read_frame(stream, self.config.max_frame_size)? { + if let Some(mut frame) = self + .frame + .read_frame(stream, self.config.max_frame_size) + .check_connection_reset(&self.state)? + { if !self.state.can_read() { return Err(Error::Protocol( "Remote sent frame after having sent a Close Frame".into(), @@ -560,18 +564,9 @@ impl WebSocketContext { } trace!("Sending frame: {:?}", frame); - let res = self.frame.write_frame(stream, frame); - // An expected "Connection reset by peer" is not fatal - match res { - Err(Error::Io(err)) => Err({ - if !self.state.can_read() && err.kind() == IoErrorKind::ConnectionReset { - Error::ConnectionClosed - } else { - Error::Io(err) - } - }), - x => x, - } + self.frame + .write_frame(stream, frame) + .check_connection_reset(&self.state) } } @@ -618,6 +613,26 @@ impl WebSocketState { } } +/// Translate "Connection reset by peer" into `ConnectionClosed` if appropriate. +trait CheckConnectionReset { + fn check_connection_reset(self, state: &WebSocketState) -> Self; +} + +impl CheckConnectionReset for Result { + fn check_connection_reset(self, state: &WebSocketState) -> Self { + match self { + Err(Error::Io(io_error)) => Err({ + if !state.can_read() && io_error.kind() == IoErrorKind::ConnectionReset { + Error::ConnectionClosed + } else { + Error::Io(io_error) + } + }), + x => x, + } + } +} + #[cfg(test)] mod tests { use super::{Message, Role, WebSocket, WebSocketConfig};