From e3c0ec30c8ed427d65abc8a7d76d003887f97016 Mon Sep 17 00:00:00 2001
From: Alexey Galakhov <agalakhov@snapview.de>
Date: Fri, 19 May 2017 17:54:02 +0200
Subject: [PATCH] Fix close handling while read.

Signed-off-by: Alexey Galakhov <agalakhov@snapview.de>
---
 src/protocol/mod.rs | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 6266dcb..b576fc3 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -293,7 +293,14 @@ impl<Stream: Read + Write> WebSocket<Stream> {
             } // match opcode
 
         } else {
-            Err(Error::Protocol("Connection reset without closing handshake".into()))
+            match replace(&mut self.state, WebSocketState::Terminated) {
+                WebSocketState::CloseAcknowledged(close) | WebSocketState::ClosedByPeer(close) => {
+                    Err(Error::ConnectionClosed(close))
+                }
+                _ => {
+                    Err(Error::Protocol("Connection reset without closing handshake".into()))
+                }
+            }
         }
     }
 
@@ -342,6 +349,7 @@ impl<Stream: Read + Write> WebSocket<Stream> {
                     }
                 }
             }
+            WebSocketState::Terminated => unreachable!(),
         }
     }
 
@@ -413,6 +421,8 @@ enum WebSocketState {
     ClosedByPeer(Option<CloseFrame<'static>>),
     /// The peer replied to our close handshake.
     CloseAcknowledged(Option<CloseFrame<'static>>),
+    /// The connection does not exist anymore.
+    Terminated,
 }
 
 impl WebSocketState {