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

use super::*;
use dever_runtime::http::websocket::{self as ws, Message, WebSocket};

mod lifecycle;
mod protocol;
mod shared;
#[path = "../../../performance/websocket/shared/mod.rs"]
mod shared_reference;

fn limits() -> ws::Limits {
    ws::Limits {
        idle_ms: 1000,
        write_ms: 1000,
        ..ws::Limits::default()
    }
}

async fn echo(
    _: Request,
    reply: http::HttpReply,
    scope: Arc<Scope<'static>>,
) -> Result<(), String> {
    let socket = ws::accept(&reply, &limits(), &scope)
        .await
        .map_err(|error| error.to_string())?;
    while let Ok(Some(message)) = socket.receive().await {
        if matches!(message, Message::Text(_) | Message::Binary(_)) {
            socket
                .send(message)
                .await
                .map_err(|error| error.to_string())?;
        }
    }
    Ok(())
}

async fn server(
    scope: &Scope<'_>,
) -> (
    net::Listener,
    dever_runtime::task::Task<Result<(), http::ServeError>>,
) {
    let listener = net::listen("127.0.0.1", 0, scope).await.unwrap();
    let incoming = listener.clone();
    let task = scope
        .spawn(move |child| async move {
            http::serve_live(
                &incoming,
                &Limits::default(),
                &http::LiveLimits::default(),
                &child,
                echo,
            )
            .await
        })
        .unwrap();
    (listener, task)
}

#[test]
fn websocket_live_upgrade_echo_control_frames_and_close() {
    runtime()
        .run(async |scope| {
            let (listener, server) = server(scope).await;
            let socket = ws::connect(
                "127.0.0.1",
                listener.port().unwrap(),
                "/echo?key=value",
                &limits(),
                scope,
            )
            .await
            .unwrap();
            for message in [
                Message::Text("你好".into()),
                Message::Binary(vec![0, 255, 17].into()),
            ] {
                socket.send(message.clone()).await.unwrap();
                assert_eq!(socket.receive().await.unwrap(), Some(message));
            }
            socket
                .send(Message::Ping(b"ping".to_vec().into()))
                .await
                .unwrap();
            assert_eq!(
                socket.receive().await.unwrap(),
                Some(Message::Pong(b"ping".to_vec().into()))
            );
            assert!(
                socket
                    .send(Message::Ping(vec![0; 126].into()))
                    .await
                    .is_err()
            );
            assert!(socket.close(1005, "").await.is_err());
            assert!(socket.close(1000, &"界".repeat(42)).await.is_err());
            socket
                .send(Message::Text("still open".into()))
                .await
                .unwrap();
            assert_eq!(
                socket.receive().await.unwrap(),
                Some(Message::Text("still open".into()))
            );
            socket.close(1000, "结束").await.unwrap();
            assert!(socket.receive().await.unwrap().is_none());
            listener.close().unwrap();
            server.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn websocket_aliases_send_whole_messages_while_receiving() {
    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 writers = futures_util::future::join_all((0u8..4).map(|writer| {
                let socket = &socket;
                async move {
                    for sequence in 0u8..16 {
                        let mut bytes = vec![writer; 65536];
                        bytes[1] = sequence;
                        socket.send(Message::Binary(bytes.into())).await.unwrap();
                    }
                }
            }));
            let receiving = async {
                let mut next = [0u8; 4];
                for _ in 0..64 {
                    let Some(Message::Binary(bytes)) = socket.receive().await.unwrap() else {
                        panic!("expected a complete binary message");
                    };
                    assert_eq!(bytes.len(), 65536);
                    let writer = bytes[0];
                    assert_eq!(bytes[1], next[usize::from(writer)]);
                    assert!(bytes[2..].iter().all(|byte| *byte == writer));
                    next[usize::from(writer)] += 1;
                }
                assert_eq!(next, [16; 4]);
            };
            tokio::join!(writers, receiving);
            socket.close(1000, "").await.unwrap();
            listener.close().unwrap();
            server.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn websocket_wss_uses_supplied_ca_and_dns_with_no_insecure_fallback() {
    runtime()
        .run(async |scope| {
            let listener = net::listen("127.0.0.1", 0, scope).await.unwrap();
            let port = listener.port().unwrap();
            let incoming = listener.clone();
            let server = scope
                .spawn(move |child| async move {
                    let tls = dever_runtime::tls::server(tls::CERTIFICATE, tls::KEY).unwrap();
                    http::serve_live_tls(
                        &incoming,
                        &Limits::default(),
                        &http::LiveLimits::default(),
                        &child,
                        &tls,
                        echo,
                    )
                    .await
                })
                .unwrap();
            let url = format!("wss://localhost:{port}/echo");
            assert!(
                ws::open(&url, &dever_runtime::tls::system(), &limits(), scope)
                    .await
                    .is_err()
            );
            let trusted = dever_runtime::tls::client(tls::ROOT).unwrap();
            let socket = ws::open(&url, &trusted, &limits(), scope).await.unwrap();
            socket
                .send(Message::Text("verified TLS".into()))
                .await
                .unwrap();
            assert_eq!(
                socket.receive().await.unwrap(),
                Some(Message::Text("verified TLS".into()))
            );
            socket.close(1000, "").await.unwrap();
            listener.close().unwrap();
            server.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn websocket_invalid_inputs_fail_before_connecting() {
    runtime()
        .run(async |scope| {
            for limit in [
                ws::Limits {
                    message_bytes: 0,
                    ..limits()
                },
                ws::Limits {
                    idle_ms: 0,
                    ..limits()
                },
                ws::Limits {
                    write_ms: -1,
                    ..limits()
                },
            ] {
                assert!(
                    ws::connect("127.0.0.1", 0, "/", &limit, scope)
                        .await
                        .is_err()
                );
            }
            for url in [
                "http://localhost/",
                "ws://user:password@localhost/",
                "ws://localhost/#fragment",
                "ws://localhost:99999/",
            ] {
                assert!(
                    ws::open(url, &dever_runtime::tls::system(), &limits(), scope)
                        .await
                        .is_err()
                );
            }
            for target in ["http://localhost/", "*", "/#fragment"] {
                assert!(
                    ws::connect("127.0.0.1", 0, target, &limits(), scope)
                        .await
                        .is_err()
                );
            }
        })
        .unwrap();
}