Merge pull request #74 from najamelan/bugfix/no_send_after_close

Do not allow sending messages after sending a close frame

src/protocol/mod.rs

@@ -230,9 +230,14 @@ impl WebSocketContext {
     where
         Stream: Read + Write,
     {
-        // Do not write to already closed connections.
+        // When terminated, return AlreadyClosed.
         self.state.check_active()?;
 
+        // Do not write after sending a close frame.
+        if !self.state.is_active() {
+            return Err(Error::Protocol("Sending after closing is not allowed".into()));
+        }
+
         if let Some(max_send_queue) = self.config.max_send_queue {
             if self.send_queue.len() >= max_send_queue {
                 // Try to make some room for the new message.
@@ -332,6 +337,9 @@ impl WebSocketContext {
         Stream: Read + Write,
     {
         if let Some(mut frame) = self.frame.read_frame(stream, self.config.max_frame_size)? {
+            if !self.state.can_read() {
+                return Err(Error::Protocol("Remote sent frame after having sent a Close Frame".into()));
+            }
             // MUST be 0 unless an extension is negotiated that defines meanings
             // for non-zero values.  If a nonzero value is received and none of
             // the negotiated extensions defines the meaning of such a nonzero
@@ -383,9 +391,6 @@ impl WebSocketContext {
                         OpCtl::Reserved(i) => Err(Error::Protocol(
                             format!("Unknown control frame type {}", i).into(),
                         )),
-                        OpCtl::Ping | OpCtl::Pong if !self.state.can_read() => {
-                            Ok(None)
-                        }
                         OpCtl::Ping => {
                             let data = frame.into_data();
                             // No ping processing after we sent a close frame.
@@ -398,11 +403,6 @@ impl WebSocketContext {
                     }
                 }
 
-                OpCode::Data(_) if !self.state.can_read() => {
-                    // No data processing while closing.
-                    Ok(None)
-                }
-
                 OpCode::Data(data) => {
                     let fin = frame.header().is_final;
                     match data {

tests/no_send_after_close.rs

@@ -0,0 +1,55 @@
+//! 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_no_send_after_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();
+
+        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
+
+    let err = client_handler
+        .write_message(Message::Text("Hello WebSocket".into()));
+
+    assert!( err.is_err() );
+
+    match err.unwrap_err() {
+        Error::Protocol(s) => { assert_eq!( "Sending after closing is not allowed", s )}
+        e => panic!("unexpected error: {:?}", e),
+    }
+
+    drop(client_handler);
+
+    client_thread.join().unwrap();
+}