test: check close with SO_LINGER

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

Cargo.toml

@@ -35,3 +35,4 @@ version = "0.2.3"
 
 [dev-dependencies]
 env_logger = "0.7.1"
+net2 = "0.2.33"

tests/connection_reset.rs

@@ -1,17 +1,24 @@
 //! 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.
 
-use std::net::TcpListener;
+use std::net::{TcpStream, TcpListener};
 use std::process::exit;
 use std::thread::{sleep, spawn};
 use std::time::Duration;
 
-use tungstenite::{accept, connect, Error, Message};
+use tungstenite::{accept, connect, Error, Message, WebSocket, stream::Stream};
+use native_tls::TlsStream;
 use url::Url;
+use net2::TcpStreamExt;
 
-#[test]
-fn test_close() {
-    env_logger::init();
+type Sock = WebSocket<Stream<TcpStream, TlsStream<TcpStream>>>;
+
+fn do_test<CT, ST>(port: u16, client_task: CT, server_task: ST)
+where
+    CT: FnOnce(Sock) + Send + 'static,
+    ST: FnOnce(WebSocket<TcpStream>),
+{
+    env_logger::try_init().ok();
 
     spawn(|| {
         sleep(Duration::from_secs(5));
@@ -19,42 +26,127 @@ fn test_close() {
         exit(1);
     });
 
-    let server = TcpListener::bind("127.0.0.1:3012").unwrap();
+    let server = TcpListener::bind(("127.0.0.1", port))
+        .expect("Can't listen, is port already in use?");
 
     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 (client, _) = connect(Url::parse(&format!("ws://localhost:{}/socket", port)).unwrap())
+            .expect("Can't connect to port");
 
-        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),
-        }
+        client_task(client);
     });
 
     let client_handler = server.incoming().next().unwrap();
-    let mut client_handler = accept(client_handler.unwrap()).unwrap();
+    let client_handler = accept(client_handler.unwrap()).unwrap();
 
-    let message = client_handler.read_message().unwrap();
-    assert_eq!(message.into_data(), b"Hello WebSocket");
+    server_task(client_handler);
 
-    client_handler.close(None).unwrap(); // send close to client
+    client_thread.join().unwrap();
+}
 
-    assert!(client_handler.read_message().unwrap().is_close()); // receive acknowledgement
+#[test]
+fn test_server_close() {
+    do_test(3012,
+        |mut cli_sock| {
+            cli_sock
+                .write_message(Message::Text("Hello WebSocket".into()))
+                .unwrap();
+
+            let message = cli_sock.read_message().unwrap(); // receive close from server
+            assert!(message.is_close());
+
+            let err = cli_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
+            match err {
+                Error::ConnectionClosed => {}
+                _ => panic!("unexpected error: {:?}", err),
+            }
+        },
+        |mut srv_sock| {
+            let message = srv_sock.read_message().unwrap();
+            assert_eq!(message.into_data(), b"Hello WebSocket");
+
+            srv_sock.close(None).unwrap(); // send close to client
+
+            let message = srv_sock.read_message().unwrap(); // receive acknowledgement
+            assert!(message.is_close());
+
+            let err = srv_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
+            match err {
+                Error::ConnectionClosed => {}
+                _ => panic!("unexpected error: {:?}", err),
+            }
+        });
+}
 
-    let err = client_handler.read_message().unwrap_err(); // now we should get ConnectionClosed
-    match err {
-        Error::ConnectionClosed => {}
-        _ => panic!("unexpected error: {:?}", err),
-    }
+#[test]
+fn test_evil_server_close() {
+    do_test(3013,
+        |mut cli_sock| {
+            cli_sock
+                .write_message(Message::Text("Hello WebSocket".into()))
+                .unwrap();
+
+            sleep(Duration::from_secs(1));
+
+            let message = cli_sock.read_message().unwrap(); // receive close from server
+            assert!(message.is_close());
+
+            let err = cli_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
+            match err {
+                Error::ConnectionClosed => {}
+                _ => panic!("unexpected error: {:?}", err),
+            }
+        },
+        |mut srv_sock| {
+            let message = srv_sock.read_message().unwrap();
+            assert_eq!(message.into_data(), b"Hello WebSocket");
+
+            srv_sock.close(None).unwrap(); // send close to client
+
+            let message = srv_sock.read_message().unwrap(); // receive acknowledgement
+            assert!(message.is_close());
+            // and now just drop the connection without waiting for `ConnectionClosed`
+            srv_sock.get_mut().set_linger(Some(Duration::from_secs(0))).unwrap();
+            drop(srv_sock);
+        });
+}
 
-    drop(client_handler);
+#[test]
+fn test_client_close() {
+    do_test(3014,
+        |mut cli_sock| {
+            cli_sock
+                .write_message(Message::Text("Hello WebSocket".into()))
+                .unwrap();
+
+            let message = cli_sock.read_message().unwrap(); // receive answer from server
+            assert_eq!(message.into_data(), b"From Server");
+
+            cli_sock.close(None).unwrap(); // send close to server
+
+            let message = cli_sock.read_message().unwrap(); // receive acknowledgement from server
+            assert!(message.is_close());
+
+            let err = cli_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
+            match err {
+                Error::ConnectionClosed => {}
+                _ => panic!("unexpected error: {:?}", err),
+            }
+        },
+        |mut srv_sock| {
+            let message = srv_sock.read_message().unwrap();
+            assert_eq!(message.into_data(), b"Hello WebSocket");
+
+            srv_sock.write_message(Message::Text("From Server".into())).unwrap();
+
+            let message = srv_sock.read_message().unwrap(); // receive close from client
+            assert!(message.is_close());
+
+            let err = srv_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
+            match err {
+                Error::ConnectionClosed => {}
+                _ => panic!("unexpected error: {:?}", err),
+            }
+        });
 
-    client_thread.join().unwrap();
 }