core/test/rust/runtime/http/websocket/shared.rs

use super::shared_reference::{Limits, Owner, Socket};
use futures_util::{SinkExt, StreamExt};
use std::time::Duration;
use tokio::{
    net::{TcpListener, TcpStream},
    sync::oneshot,
};
use tokio_tungstenite::{
    WebSocketStream,
    tungstenite::{Message, protocol::WebSocketConfig},
};

fn limits() -> Limits {
    Limits {
        message: 16 * 1024 * 1024,
        idle: Duration::from_secs(2),
        write: Duration::from_secs(2),
    }
}

async fn pair(limits: Limits) -> (Socket, Owner, WebSocketStream<TcpStream>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();
    let config = WebSocketConfig::default()
        .max_message_size(Some(limits.message))
        .max_frame_size(Some(limits.message))
        .write_buffer_size(0);
    let client = async {
        let stream = TcpStream::connect(address).await.unwrap();
        stream.set_nodelay(true).unwrap();
        tokio_tungstenite::client_async_with_config(
            format!("ws://{address}/"),
            stream,
            Some(config),
        )
        .await
        .unwrap()
        .0
    };
    let server = async {
        let (stream, _) = listener.accept().await.unwrap();
        stream.set_nodelay(true).unwrap();
        tokio_tungstenite::accept_async_with_config(stream, Some(config))
            .await
            .unwrap()
    };
    let (client, server) = tokio::join!(client, server);
    let (socket, owner) = Socket::new(client, limits);
    (socket, owner, server)
}

#[tokio::test]
async fn websocket_rust_reference_owner_exit_wakes_reader_and_releases_tcp() {
    let (socket, owner, mut peer) = pair(limits()).await;
    let alias = socket.clone();
    let reading = alias.receive();
    tokio::pin!(reading);
    assert!(futures_util::poll!(&mut reading).is_pending());
    drop(owner);
    assert!(
        tokio::time::timeout(Duration::from_millis(100), reading)
            .await
            .unwrap()
            .is_err()
    );
    assert!(socket.send(Message::Text("closed".into())).await.is_err());
    assert!(
        tokio::time::timeout(Duration::from_millis(100), peer.next())
            .await
            .unwrap()
            .unwrap()
            .is_err()
    );
}

#[tokio::test]
async fn websocket_rust_reference_full_duplex_aliases_preserve_messages() {
    let (socket, _owner, mut peer) = pair(limits()).await;
    let server = tokio::spawn(async move {
        while let Some(message) = peer.next().await {
            let message = message.unwrap();
            if message.is_close() {
                peer.flush().await.unwrap();
                break;
            }
            peer.send(message).await.unwrap();
        }
    });
    let sender = |writer: u8| {
        let socket = socket.clone();
        async move {
            for sequence in 0..16 {
                let mut bytes = vec![writer; 65536];
                bytes[1] = sequence;
                socket.send(Message::Binary(bytes.into())).await.unwrap();
            }
        }
    };
    let receiver = async {
        let mut sequences = [0_u8; 2];
        for _ in 0..32 {
            let Message::Binary(bytes) = socket.receive().await.unwrap().unwrap() else {
                panic!("binary expected")
            };
            let writer = usize::from(bytes[0]);
            assert_eq!(bytes[1], sequences[writer]);
            assert_eq!(bytes.len(), 65536);
            assert!(bytes[2..].iter().all(|byte| *byte == writer as u8));
            sequences[writer] += 1;
        }
        assert_eq!(sequences, [16, 16]);
    };
    tokio::time::timeout(Duration::from_secs(5), async {
        tokio::join!(sender(0), sender(1), receiver);
    })
    .await
    .unwrap();
    socket.close(1000, "").await.unwrap();
    server.await.unwrap();
}

#[tokio::test]
async fn websocket_rust_reference_close_has_priority_over_readable_reply() {
    let (socket, _owner, mut peer) = pair(limits()).await;
    let (sent, ready) = oneshot::channel();
    let server = tokio::spawn(async move {
        assert!(peer.next().await.unwrap().unwrap().is_close());
        peer.flush().await.unwrap();
        sent.send(()).unwrap();
    });
    let reading = socket.receive();
    tokio::pin!(reading);
    assert!(futures_util::poll!(&mut reading).is_pending());
    let closing = socket.close(1000, "");
    tokio::pin!(closing);
    assert!(futures_util::poll!(&mut closing).is_pending());
    ready.await.unwrap();
    tokio::time::sleep(Duration::from_millis(25)).await;
    assert!(reading.await.is_err());
    closing.await.unwrap();
    server.await.unwrap();
}

#[tokio::test]
async fn websocket_rust_reference_cancelled_writer_only_closes_after_lock() {
    let (socket, _owner, _peer) = pair(limits()).await;
    let mut writing = Box::pin(socket.send(Message::Binary(vec![1; limits().message].into())));
    assert!(futures_util::poll!(&mut writing).is_pending());
    let mut queued = Box::pin(socket.send(Message::Text("queued".into())));
    assert!(futures_util::poll!(&mut queued).is_pending());
    drop(queued);
    assert!(futures_util::poll!(&mut writing).is_pending());
    drop(writing);
    assert!(socket.send(Message::Text("closed".into())).await.is_err());
}

#[tokio::test]
async fn websocket_rust_reference_limits_ping_and_idle_close() {
    let (socket, _owner, mut peer) = pair(Limits {
        message: 128,
        idle: Duration::from_millis(25),
        ..limits()
    })
    .await;
    assert!(
        socket
            .send(Message::Binary(vec![0; 129].into()))
            .await
            .is_err()
    );
    assert!(
        socket
            .send(Message::Ping(vec![0; 126].into()))
            .await
            .is_err()
    );
    peer.send(Message::Ping(vec![7].into())).await.unwrap();
    assert!(socket.receive().await.unwrap().unwrap().is_ping());
    assert_eq!(
        peer.next().await.unwrap().unwrap(),
        Message::Pong(vec![7].into())
    );
    assert!(socket.receive().await.unwrap_err().contains("timed out"));
    assert!(socket.send(Message::Text("closed".into())).await.is_err());
}