protocol: handle "Connection reset by peer" in closing state

Signed-off-by: Alexey Galakhov <agalakhov@snapview.de>

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);
+        self.frame
-        // An expected "Connection reset by peer" is not fatal
+            .write_frame(stream, frame)
-        match res {
+            .check_connection_reset(&self.state)
-            Err(Error::Io(err)) => Err({
-                if !self.state.can_read() && err.kind() == IoErrorKind::ConnectionReset {
-                    Error::ConnectionClosed
-                } else {
-                    Error::Io(err)
-                }
-            }),
-            x => x,
-        }
     }
 }
 
@@ -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<T> CheckConnectionReset for Result<T> {
+    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};