omnivox-rs/tests/loopback.rs

//! Loopback behavior tests for the v2 transport: framing discipline,
//! timeouts, junk handling, and the send-path contracts.

mod support;

use omnivox::{Config, ConnectionContext, Decision, OmniVox};
use std::{
    sync::{
        Arc,
        atomic::{AtomicU32, Ordering},
    },
    time::Duration,
};
use support::{Msg, TEST_PSK, client_session, frame};
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_connections: 16,
        allow_insecure_defaults: true,
    }
}

async fn start_server(port: u16, rx: Arc<AtomicU32>, config: Config) -> Arc<OmniVox<Msg>> {
    let server = Arc::new(OmniVox::<Msg>::new());
    server.set_encrypt_key(TEST_PSK);
    let data = Arc::new(move |msg: &Msg, _ctx: &ConnectionContext| {
        if msg.0.starts_with(b"ok") {
            rx.fetch_add(1, Ordering::SeqCst);
        }
    });
    server
        .start_server(
            port,
            config,
            data,
            Arc::new(|_ctx| {}),
            Arc::new(|_ctx| Decision::Allow),
        )
        .await
        .unwrap();
    server
}

async fn read_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_body_with_valid_counter_is_skipped_connection_survives() {
    let port = 18201;
    let rx = Arc::new(AtomicU32::new(0));
    let server = start_server(port, Arc::clone(&rx), test_config()).await;

    let (r, mut w, s) = client_session(port, &TEST_PSK).await;

    // Counter 1, tampered ciphertext: auth fails, frame dropped, alive.
    let mut bad = frame(&s.send_key, 1, b"ok-junk");
    let last = bad.len() - 1;
    bad[last] ^= 0x01;
    w.write_all(&bad).await.unwrap();

    // Counter 1 valid (the tampered frame consumed nothing).
    w.write_all(&frame(&s.send_key, 1, b"ok1")).await.unwrap();
    w.write_all(&frame(&s.send_key, 2, b"ok2")).await.unwrap();

    tokio::time::sleep(Duration::from_millis(300)).await;
    assert_eq!(rx.load(Ordering::SeqCst), 2);
    // still alive: another frame gets through
    w.write_all(&frame(&s.send_key, 3, b"ok3")).await.unwrap();
    tokio::time::sleep(Duration::from_millis(300)).await;
    assert_eq!(rx.load(Ordering::SeqCst), 3);
    server.stop().await;
    drop(r);
    drop(w);
}

#[tokio::test]
async fn zero_size_header_disconnects_immediately() {
    // v2 retires the v1 no-return quirk: bad header closes, no follow-up.
    let port = 18202;
    let rx = Arc::new(AtomicU32::new(0));
    let server = start_server(port, Arc::clone(&rx), test_config()).await;

    let (mut r, mut w, _s) = client_session(port, &TEST_PSK).await;
    w.write_all(&0u32.to_le_bytes()).await.unwrap();
    read_expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn oversized_header_disconnects() {
    let port = 18203;
    let rx = Arc::new(AtomicU32::new(0));
    let server = start_server(port, Arc::clone(&rx), test_config()).await;

    let (mut r, mut w, _s) = client_session(port, &TEST_PSK).await;
    w.write_all(&8193u32.to_le_bytes()).await.unwrap();
    read_expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn header_stall_times_out() {
    let port = 18204;
    let rx = Arc::new(AtomicU32::new(0));
    let server = start_server(port, Arc::clone(&rx), test_config()).await;

    let (mut r, mut w, _s) = client_session(port, &TEST_PSK).await;
    w.write_all(&[0x10, 0x00]).await.unwrap(); // partial header, then silence
    read_expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn send_to_unknown_client_is_false() {
    let server = OmniVox::<Msg>::new();
    server.set_encrypt_key(TEST_PSK);
    assert!(!server.send_to_client(&Msg(b"x".to_vec()), 999).await);
    assert!(!server.send_to_server(&Msg(b"x".to_vec())).await);
}

#[tokio::test]
async fn valid_handshake_then_api_roundtrip() {
    let port = 18205;
    let rx = Arc::new(AtomicU32::new(0));
    let server = Arc::new(OmniVox::<Msg>::new());
    server.set_encrypt_key(TEST_PSK);
    let data = Arc::new({
        let rx = Arc::clone(&rx);
        move |msg: &Msg, ctx: &ConnectionContext| {
            if msg.0.starts_with(b"ping") {
                rx.fetch_add(1, Ordering::SeqCst);
                let server = Arc::clone(SERVER_FOR_ECHO.get().unwrap());
                let id = ctx.id;
                let reply = Msg(b"pong".to_vec());
                tokio::spawn(async move {
                    server.send_to_client(&reply, id).await;
                });
            }
        }
    });
    let _ = SERVER_FOR_ECHO.set(server.clone());
    server
        .start_server(
            port,
            test_config(),
            data,
            Arc::new(|_c| {}),
            Arc::new(|_c| Decision::Allow),
        )
        .await
        .unwrap();

    let client = OmniVox::<Msg>::new();
    client.set_encrypt_key(TEST_PSK);
    let got = Arc::new(AtomicU32::new(0));
    let g = Arc::clone(&got);
    client
        .connect_to_server(
            "127.0.0.1",
            port,
            test_config(),
            Some(Arc::new(move |msg: &Msg, _c| {
                if msg.0 == b"pong" {
                    g.fetch_add(1, Ordering::SeqCst);
                }
            })),
            Arc::new(|_c| {}),
        )
        .await
        .unwrap();

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

static SERVER_FOR_ECHO: std::sync::OnceLock<Arc<OmniVox<Msg>>> = std::sync::OnceLock::new();