You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
tungstenite-rs/tests/connection_reset.rs

163 lines
5.4 KiB

//! Verifies that the server returns a `ConnectionClosed` error when the connection
//! is closed from the server's point of view and drop the underlying tcp socket.
#![cfg(any(
feature = "native-tls",
feature = "rustls-tls-native-roots",
feature = "rustls-tls-webpki-roots"
))]
use std::{
net::{TcpListener, TcpStream},
process::exit,
thread::{sleep, spawn},
time::Duration,
};
use net2::TcpStreamExt;
use tungstenite::{accept, connect, stream::Stream, Error, Message, WebSocket};
use url::Url;
#[cfg(feature = "native-tls")]
type Sock = WebSocket<Stream<TcpStream, native_tls_crate::TlsStream<TcpStream>>>;
#[cfg(all(
any(feature = "rustls-tls-native-roots", feature = "rustls-tls-webpki-roots"),
not(feature = "native-tls")
))]
type Sock = WebSocket<Stream<TcpStream, rustls::StreamOwned<rustls::ClientSession, 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));
println!("Unit test executed too long, perhaps stuck on WOULDBLOCK...");
exit(1);
});
let server =
TcpListener::bind(("127.0.0.1", port)).expect("Can't listen, is port already in use?");
let client_thread = spawn(move || {
let (client, _) = connect(Url::parse(&format!("ws://localhost:{}/socket", port)).unwrap())
.expect("Can't connect to port");
client_task(client);
});
let client_handler = server.incoming().next().unwrap();
let client_handler = accept(client_handler.unwrap()).unwrap();
server_task(client_handler);
client_thread.join().unwrap();
}
#[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),
}
},
);
}
#[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);
},
);
}
#[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),
}
},
);
}