326 lines
13 KiB
Rust
326 lines
13 KiB
Rust
use super::*;
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use tokio::sync::oneshot;
|
|
|
|
#[test]
|
|
fn websocket_client_scope_ends_all_aliases_and_timeout_stream_fails_once() {
|
|
runtime()
|
|
.run(async |scope| {
|
|
let (listener, server) = server(scope).await;
|
|
let port = listener.port().unwrap();
|
|
let owner = scope
|
|
.spawn(move |child| async move {
|
|
ws::connect("127.0.0.1", port, "/", &limits(), &child)
|
|
.await
|
|
.unwrap()
|
|
})
|
|
.unwrap();
|
|
let escaped = owner.wait().await.unwrap();
|
|
assert!(escaped.send(Message::Text("closed".into())).await.is_err());
|
|
let socket = ws::connect(
|
|
"127.0.0.1",
|
|
port,
|
|
"/",
|
|
&ws::Limits {
|
|
idle_ms: 25,
|
|
..limits()
|
|
},
|
|
scope,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let stream = socket.messages(|message| message);
|
|
let alias = stream.clone();
|
|
let failure = stream.next().await.unwrap().unwrap().unwrap_err();
|
|
assert!(failure.contains("timed out"));
|
|
assert!(alias.next().await.unwrap().is_none());
|
|
assert!(socket.send(Message::Text("closed".into())).await.is_err());
|
|
listener.close().unwrap();
|
|
server.wait().await.unwrap().unwrap();
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn websocket_listener_drain_keeps_session_and_parent_stop_cleans_handler() {
|
|
runtime()
|
|
.run(async |scope| {
|
|
let (listener, mut server) = server(scope).await;
|
|
let socket = ws::connect("127.0.0.1", listener.port().unwrap(), "/", &limits(), scope)
|
|
.await
|
|
.unwrap();
|
|
listener.close().unwrap();
|
|
tokio::select! {
|
|
_ = &mut server => panic!("server must retain its upgraded session"),
|
|
() = tokio::time::sleep(Duration::from_millis(25)) => (),
|
|
}
|
|
socket
|
|
.send(Message::Text("after listener close".into()))
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(
|
|
socket.receive().await.unwrap(),
|
|
Some(Message::Text("after listener close".into()))
|
|
);
|
|
socket.close(1000, "").await.unwrap();
|
|
server.wait().await.unwrap().unwrap();
|
|
|
|
let listener = net::listen("127.0.0.1", 0, scope).await.unwrap();
|
|
let incoming = listener.clone();
|
|
let active = Arc::new(AtomicUsize::new(0));
|
|
let running = active.clone();
|
|
let (sender, receiver) = oneshot::channel::<WebSocket>();
|
|
let sender = Arc::new(std::sync::Mutex::new(Some(sender)));
|
|
let server = scope
|
|
.spawn(move |child| async move {
|
|
http::serve_live(
|
|
&incoming,
|
|
&Limits::default(),
|
|
&http::LiveLimits::default(),
|
|
&child,
|
|
move |_, reply, scope| {
|
|
let sender = sender.clone();
|
|
let running = running.clone();
|
|
async move {
|
|
let socket = ws::accept(&reply, &limits(), &scope).await.unwrap();
|
|
running.fetch_add(1, Ordering::SeqCst);
|
|
let _running = Running(running);
|
|
let _ = sender.lock().unwrap().take().unwrap().send(socket.clone());
|
|
let _ = socket.receive().await;
|
|
Ok(())
|
|
}
|
|
},
|
|
)
|
|
.await
|
|
})
|
|
.unwrap();
|
|
let socket = ws::connect("127.0.0.1", listener.port().unwrap(), "/", &limits(), scope)
|
|
.await
|
|
.unwrap();
|
|
let escaped = receiver.await.unwrap();
|
|
server.stop().await.unwrap();
|
|
assert_eq!(active.load(Ordering::SeqCst), 0);
|
|
assert!(escaped.send(Message::Text("gone".into())).await.is_err());
|
|
assert!(socket.receive().await.is_err());
|
|
listener.close().unwrap();
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
struct Running(Arc<AtomicUsize>);
|
|
impl Drop for Running {
|
|
fn drop(&mut self) {
|
|
self.0.fetch_sub(1, Ordering::SeqCst);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn websocket_owner_stop_wakes_receive_in_another_task() {
|
|
runtime()
|
|
.run(async |scope| {
|
|
let (listener, server) = server(scope).await;
|
|
let port = listener.port().unwrap();
|
|
let (sender, receiver) = oneshot::channel();
|
|
let owner = scope
|
|
.spawn(move |child| async move {
|
|
let socket = ws::connect("127.0.0.1", port, "/", &limits(), &child)
|
|
.await
|
|
.unwrap();
|
|
sender.send(socket).unwrap();
|
|
std::future::pending::<()>().await;
|
|
})
|
|
.unwrap();
|
|
let socket = receiver.await.unwrap();
|
|
let (started, ready) = oneshot::channel();
|
|
let reader = scope
|
|
.spawn_leaf(move || async move {
|
|
let receiving = socket.receive();
|
|
tokio::pin!(receiving);
|
|
assert!(futures_util::poll!(&mut receiving).is_pending());
|
|
started.send(()).unwrap();
|
|
tokio::time::timeout(Duration::from_millis(200), receiving)
|
|
.await
|
|
.unwrap()
|
|
.unwrap_err()
|
|
})
|
|
.unwrap();
|
|
ready.await.unwrap();
|
|
owner.stop().await.unwrap();
|
|
assert!(reader.wait().await.unwrap().contains("WebSocket is closed"));
|
|
listener.close().unwrap();
|
|
server.wait().await.unwrap().unwrap();
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn websocket_closing_message_stream_preserves_independent_socket_alias() {
|
|
runtime()
|
|
.run(async |scope| {
|
|
let (listener, server) = server(scope).await;
|
|
let socket = ws::connect("127.0.0.1", listener.port().unwrap(), "/", &limits(), scope)
|
|
.await
|
|
.unwrap();
|
|
let stream = socket.messages(|message| message);
|
|
stream.close();
|
|
assert!(stream.next().await.unwrap().is_none());
|
|
socket
|
|
.send(Message::Binary(vec![1, 2].into()))
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(
|
|
socket.receive().await.unwrap(),
|
|
Some(Message::Binary(vec![1, 2].into()))
|
|
);
|
|
socket.close(1000, "").await.unwrap();
|
|
listener.close().unwrap();
|
|
server.wait().await.unwrap().unwrap();
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn websocket_cancelled_writer_closes_only_after_it_owns_the_direction() {
|
|
runtime()
|
|
.run(async |scope| {
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let port = listener.local_addr().unwrap().port();
|
|
let (started, receiving) = oneshot::channel();
|
|
let (release, released) = oneshot::channel();
|
|
let peer = scope
|
|
.spawn_leaf(move || async move {
|
|
let (stream, _) = listener.accept().await.unwrap();
|
|
let wire = tokio_tungstenite::accept_async(stream).await.unwrap();
|
|
wire.get_ref().readable().await.unwrap();
|
|
let _ = started.send(());
|
|
let _ = released.await;
|
|
})
|
|
.unwrap();
|
|
let socket = ws::connect(
|
|
"127.0.0.1",
|
|
i64::from(port),
|
|
"/",
|
|
&ws::Limits {
|
|
message_bytes: 16 * 1024 * 1024,
|
|
write_ms: 1000,
|
|
..limits()
|
|
},
|
|
scope,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let sender = socket.clone();
|
|
let mut writing = scope
|
|
.spawn_leaf(move || async move {
|
|
sender
|
|
.send(Message::Binary(vec![1; 16 * 1024 * 1024].into()))
|
|
.await
|
|
})
|
|
.unwrap();
|
|
receiving.await.unwrap();
|
|
let waiting = socket.clone();
|
|
let waiter = scope
|
|
.spawn_leaf(
|
|
move || async move { waiting.send(Message::Text("queued".into())).await },
|
|
)
|
|
.unwrap();
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
waiter.stop().await.unwrap();
|
|
assert!(
|
|
tokio::time::timeout(Duration::from_millis(10), &mut writing)
|
|
.await
|
|
.is_err()
|
|
);
|
|
writing.stop().await.unwrap();
|
|
assert!(socket.send(Message::Text("closed".into())).await.is_err());
|
|
release.send(()).unwrap();
|
|
peer.wait().await.unwrap();
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn websocket_close_keeps_the_handshake_when_a_reader_is_waiting() {
|
|
runtime()
|
|
.run(async |scope| {
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let port = listener.local_addr().unwrap().port();
|
|
let (replied, reply) = oneshot::channel();
|
|
let peer = scope
|
|
.spawn_leaf(move || async move {
|
|
let (stream, _) = listener.accept().await.unwrap();
|
|
let mut wire = tokio_tungstenite::accept_async(stream).await.unwrap();
|
|
assert!(matches!(
|
|
wire.next().await,
|
|
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_)))
|
|
));
|
|
wire.flush().await.unwrap();
|
|
replied.send(()).unwrap();
|
|
})
|
|
.unwrap();
|
|
let socket = ws::connect("127.0.0.1", i64::from(port), "/", &limits(), scope)
|
|
.await
|
|
.unwrap();
|
|
let receiving = socket.receive();
|
|
tokio::pin!(receiving);
|
|
assert!(futures_util::poll!(&mut receiving).is_pending());
|
|
let closing = socket.close(1000, "");
|
|
tokio::pin!(closing);
|
|
assert!(futures_util::poll!(&mut closing).is_pending());
|
|
reply.await.unwrap();
|
|
tokio::time::sleep(Duration::from_millis(25)).await;
|
|
let received = receiving.await;
|
|
closing.await.unwrap();
|
|
assert!(received.unwrap_err().contains("WebSocket is closing"));
|
|
peer.wait().await.unwrap();
|
|
})
|
|
.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn websocket_close_deadline_releases_all_aliases() {
|
|
runtime()
|
|
.run(async |scope| {
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let port = listener.local_addr().unwrap().port();
|
|
let (release, released) = oneshot::channel();
|
|
let peer = scope
|
|
.spawn_leaf(move || async move {
|
|
let (stream, _) = listener.accept().await.unwrap();
|
|
let _wire = tokio_tungstenite::accept_async(stream).await.unwrap();
|
|
let _ = released.await;
|
|
})
|
|
.unwrap();
|
|
let socket = ws::connect(
|
|
"127.0.0.1",
|
|
i64::from(port),
|
|
"/",
|
|
&ws::Limits {
|
|
write_ms: 50,
|
|
..limits()
|
|
},
|
|
scope,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let alias = socket.clone();
|
|
let receiving = alias.receive();
|
|
tokio::pin!(receiving);
|
|
assert!(futures_util::poll!(&mut receiving).is_pending());
|
|
let (closed, received) = tokio::join!(socket.close(1000, ""), receiving);
|
|
assert!(closed.unwrap_err().contains("timed out"));
|
|
assert!(received.unwrap_err().contains("WebSocket is closing"));
|
|
assert!(
|
|
alias
|
|
.send(Message::Text("closed".into()))
|
|
.await
|
|
.unwrap_err()
|
|
.contains("WebSocket is closed")
|
|
);
|
|
release.send(()).unwrap();
|
|
peer.wait().await.unwrap();
|
|
})
|
|
.unwrap();
|
|
}
|