omnivox-rs/tests/loopback.rs

//! Behavior tests for the v2.1 transport via the public builder API:
//! framing discipline, timeouts, the echo roundtrip, and send-path returns.

mod support;

use omnivox::{Config, Decision, Server};
use std::{
    sync::{
        Arc,
        atomic::{AtomicU32, Ordering},
    },
    time::Duration,
};
use support::{Msg, TEST_PSK, client_session};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

fn test_config() -> Config {
    Config {
        idle_timeout_ms: 1500,
        body_stall_ms: 800,
        send_timeout_ms: 1000,
        handshake_timeout_ms: 5000,
        max_message_size: 8192,
        max_stream_frame: 8192,
        max_connections: 16,
        allow_insecure_defaults: true,
    }
}

async fn counting_server(port: u16, rx: Arc<AtomicU32>) -> Server<Msg> {
    omnivox::server::<Msg>()
        .key(TEST_PSK)
        .config(test_config())
        .on_data(move |msg: &Msg, _ctx| {
            if msg.0.starts_with(b"ok") {
                rx.fetch_add(1, Ordering::SeqCst);
            }
        })
        .on_connect(|_ctx| Decision::Allow)
        .serve(port)
        .await
        .unwrap()
}

async fn expect_eof(r: &mut tokio::net::tcp::OwnedReadHalf) {
    let mut buf = [0u8; 1];
    let n = tokio::time::timeout(Duration::from_secs(4), r.read(&mut buf))
        .await
        .expect("connection must close, not hang")
        .unwrap();
    assert_eq!(n, 0, "expected EOF");
}

#[tokio::test]
async fn tampered_frame_disconnects_v21_policy() {
    // v2.1: authentication failure is fatal — no skip-and-continue.
    let port = 18401;
    let rx = Arc::new(AtomicU32::new(0));
    let server = counting_server(port, Arc::clone(&rx)).await;

    let (mut r, mut w, s) = client_session(port, &TEST_PSK).await;
    let mut good = support::msg_frame(&s.send_key, 1, b"ok1");
    let last = good.len() - 5; // flip a ciphertext byte
    good[last] ^= 0x01;
    w.write_all(&good).await.unwrap();
    expect_eof(&mut r).await;
    assert_eq!(rx.load(Ordering::SeqCst), 0);
    server.stop().await;
}

#[tokio::test]
async fn zero_size_header_disconnects() {
    let port = 18402;
    let rx = Arc::new(AtomicU32::new(0));
    let server = counting_server(port, rx).await;
    let (mut r, mut w, _s) = client_session(port, &TEST_PSK).await;
    w.write_all(&0u32.to_le_bytes()).await.unwrap();
    expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn oversized_header_disconnects() {
    let port = 18403;
    let rx = Arc::new(AtomicU32::new(0));
    let server = counting_server(port, rx).await;
    let (mut r, mut w, _s) = client_session(port, &TEST_PSK).await;
    w.write_all(&9000u32.to_le_bytes()).await.unwrap();
    expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn header_stall_times_out() {
    let port = 18404;
    let rx = Arc::new(AtomicU32::new(0));
    let server = counting_server(port, rx).await;
    let (mut r, mut w, _s) = client_session(port, &TEST_PSK).await;
    w.write_all(&[0x10, 0x00]).await.unwrap(); // partial header, then silence
    expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn send_to_unknown_client_is_false() {
    let rx = Arc::new(AtomicU32::new(0));
    let server = omnivox::server::<Msg>()
        .key(TEST_PSK)
        .config(test_config())
        .on_data(move |_: &Msg, _| {})
        .serve(18405)
        .await
        .unwrap();
    assert!(!server.send_to_client(999, &Msg(b"x".to_vec())).await);
    assert_eq!(server.broadcast(&Msg(b"x".to_vec())).await, 0);
    server.stop().await;
    let _ = rx;
}

#[tokio::test]
async fn echo_roundtrip_through_the_public_api() {
    let port = 18406;
    let server = omnivox::server::<Msg>()
        .key(TEST_PSK)
        .config(test_config())
        .on_data(|msg, ctx| {
            // echo via the synchronous callback lane
            let _ = ctx.reply(msg);
        })
        .serve(port)
        .await
        .unwrap();

    let got = Arc::new(AtomicU32::new(0));
    let g = Arc::clone(&got);
    let client = omnivox::client::<Msg>()
        .key(TEST_PSK)
        .config(test_config())
        .on_data(move |msg, _ctx| {
            if msg.0 == b"ping" {
                g.fetch_add(1, Ordering::SeqCst);
            }
        })
        .connect("127.0.0.1", port)
        .await
        .unwrap();

    assert!(client.send(&Msg(b"ping".to_vec())).await);
    tokio::time::sleep(Duration::from_millis(400)).await;
    assert_eq!(got.load(Ordering::SeqCst), 1, "encrypted echo returned");
    client.stop().await;
    server.stop().await;
}