Merge pull request #66 from snapview/fix-close

Fix error `ConnectionClosed` never returned to server when server initiated close

Cargo.toml

@@ -7,9 +7,9 @@ authors = ["Alexey Galakhov"]
 license = "MIT/Apache-2.0"
 readme = "README.md"
 homepage = "https://github.com/snapview/tungstenite-rs"
-documentation = "https://docs.rs/tungstenite/0.9.0"
+documentation = "https://docs.rs/tungstenite/0.9.1"
 repository = "https://github.com/snapview/tungstenite-rs"
-version = "0.9.0"
+version = "0.9.1"
 
 [features]
 default = ["tls"]

src/protocol/mod.rs

@@ -18,7 +18,7 @@ use self::frame::coding::{OpCode, Data as OpData, Control as OpCtl, CloseCode};
 use util::NonBlockingResult;
 
 /// Indicates a Client or Server role of the websocket
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Role {
     /// This socket is a server
     Server,
@@ -295,8 +295,13 @@ impl WebSocketContext {
         // If we get to this point, the send queue is empty and the underlying socket is still
         // willing to take more data.
 
+        let closing_state = match self.state {
+            WebSocketState::ClosedByPeer | WebSocketState::CloseAcknowledged => true,
+            _ => false,
+        };
+
         // If we're closing and there is nothing to send anymore, we should close the connection.
-        if let (Role::Server, WebSocketState::ClosedByPeer) = (&self.role, &self.state) {
+        if self.role == Role::Server && closing_state {
             // The underlying TCP connection, in most normal cases, SHOULD be closed
             // first by the server, so that it holds the TIME_WAIT state and not the
             // client (as this would prevent it from re-opening the connection for 2
@@ -567,8 +572,7 @@ impl WebSocketState {
     /// Check if the state is active, return error if not.
     fn check_active(&self) -> Result<()> {
         match self {
-            WebSocketState::Terminated
-                => Err(Error::AlreadyClosed),
+            WebSocketState::Terminated => Err(Error::AlreadyClosed),
             _ => Ok(()),
         }
     }

tests/connection_reset.rs

@@ -0,0 +1,62 @@
+//! Verifies that the server returns a `ConnectionClosed` error when the connection
+//! is closedd from the server's point of view and drop the underlying tcp socket.
+
+extern crate env_logger;
+extern crate tungstenite;
+extern crate url;
+
+use std::net::TcpListener;
+use std::process::exit;
+use std::thread::{spawn, sleep};
+use std::time::Duration;
+
+use tungstenite::{accept, connect, Error, Message};
+use url::Url;
+
+#[test]
+fn test_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:3012").unwrap();
+
+    let client_thread = spawn(move || {
+        let (mut client, _) = connect(Url::parse("ws://localhost:3012/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"),
+        }
+    });
+
+    let client_handler = server.incoming().next().unwrap();
+    let mut client_handler = accept(client_handler.unwrap()).unwrap();
+
+    let message = client_handler.read_message().unwrap();
+    assert_eq!(message.into_data(), b"Hello WebSocket");
+
+    client_handler.close(None).unwrap(); // send close to client
+
+    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"),
+    }
+
+    drop(client_handler);
+
+    client_thread.join().unwrap();
+}