Merge pull request #72 from najamelan/bugfix/receive_after_close

Keep processing incoming data even after we have initiated a close handshake.

src/protocol/mod.rs

@@ -396,7 +396,7 @@ impl WebSocketContext {
                     }
                 }
 
-                OpCode::Data(_) if !self.state.is_active() => {
+                OpCode::Data(_) if !self.state.can_read() => {
                     // No data processing while closing.
                     Ok(None)
                 }
@@ -533,7 +533,7 @@ impl WebSocketContext {
 }
 
 /// The current connection state.
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
 enum WebSocketState {
     /// The connection is active.
     Active,
@@ -550,12 +550,23 @@ enum WebSocketState {
 impl WebSocketState {
     /// Tell if we're allowed to process normal messages.
     fn is_active(&self) -> bool {
-        match *self {
+        match self {
             WebSocketState::Active => true,
             _ => false,
         }
     }
 
+    /// Tell if we should process incoming data. Note that if we send a close frame
+    /// but the remote hasn't confirmed, they might have sent data before they receive our
+    /// close frame, so we should still pass those to client code, hence ClosedByUs is valid.
+    fn can_read(&self) -> bool {
+        match self {
+            WebSocketState::Active     |
+            WebSocketState::ClosedByUs => true,
+            _ => false,
+        }
+    }
+
     /// Check if the state is active, return error if not.
     fn check_active(&self) -> Result<()> {
         match self {

tests/connection_reset.rs

@@ -34,7 +34,7 @@ fn test_close() {
         let err = client.read_message().unwrap_err(); // now we should get ConnectionClosed
         match err {
             Error::ConnectionClosed => {}
-            _ => panic!("unexpected error"),
+            _ => panic!("unexpected error: {:?}", err),
         }
     });
 
@@ -51,7 +51,7 @@ fn test_close() {
     let err = client_handler.read_message().unwrap_err(); // now we should get ConnectionClosed
     match err {
         Error::ConnectionClosed => {}
-        _ => panic!("unexpected error"),
+        _ => panic!("unexpected error: {:?}", err),
     }
 
     drop(client_handler);

tests/receive_after_init_close.rs

@@ -0,0 +1,61 @@
+//! Verifies that we can read data messages even if we have initiated a close handshake,
+//! but before we got confirmation.
+
+use std::net::TcpListener;
+use std::process::exit;
+use std::thread::{sleep, spawn};
+use std::time::Duration;
+
+use tungstenite::{accept, connect, Error, Message};
+use url::Url;
+
+#[test]
+fn test_receive_after_init_close() {
+    env_logger::init();
+
+    spawn(|| {
+        sleep(Duration::from_secs(5));
+        println!("Unit test executed too long, perhaps stuck on WOULDBLOCK...");
+        exit(1);
+    });
+
+    let server = TcpListener::bind("127.0.0.1:3013").unwrap();
+
+    let client_thread = spawn(move || {
+        let (mut client, _) = connect(Url::parse("ws://localhost:3013/socket").unwrap()).unwrap();
+
+        client
+            .write_message(Message::Text("Hello WebSocket".into()))
+            .unwrap();
+
+        let message = client.read_message().unwrap(); // receive close from server
+        assert!(message.is_close());
+
+        let err = client.read_message().unwrap_err(); // now we should get ConnectionClosed
+        match err {
+            Error::ConnectionClosed => {}
+            _ => panic!("unexpected error: {:?}", err),
+        }
+    });
+
+    let client_handler = server.incoming().next().unwrap();
+    let mut client_handler = accept(client_handler.unwrap()).unwrap();
+
+    client_handler.close(None).unwrap(); // send close to client
+
+    // This read should succeed even though we already initiated a close
+    let message = client_handler.read_message().unwrap();
+    assert_eq!(message.into_data(), b"Hello WebSocket");
+
+    assert!(client_handler.read_message().unwrap().is_close()); // receive acknowledgement
+
+    let err = client_handler.read_message().unwrap_err(); // now we should get ConnectionClosed
+    match err {
+        Error::ConnectionClosed => {}
+        _ => panic!("unexpected error: {:?}", err),
+    }
+
+    drop(client_handler);
+
+    client_thread.join().unwrap();
+}