diff --git a/Cargo.toml b/Cargo.toml index b339e0e..7e380f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,3 +35,4 @@ version = "0.2.3" [dev-dependencies] env_logger = "0.7.1" +net2 = "0.2.33" diff --git a/tests/connection_reset.rs b/tests/connection_reset.rs index a698a1a..d95ee81 100644 --- a/tests/connection_reset.rs +++ b/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>>; + +fn do_test(port: u16, client_task: CT, server_task: ST) +where + CT: FnOnce(Sock) + Send + 'static, + ST: FnOnce(WebSocket), +{ + 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(); }