omnivox-rs/tests/protocol_v2.rs

//! Attack regressions: every v1 vulnerability class from the red-team
//! review must now fail closed, provably, with zero callbacks fired.

mod support;

use omnivox::{Config, ConnectionContext, ConnectionStatus, Decision, OmniVox, handshake::perform};
use std::{
    sync::{
        Arc,
        atomic::{AtomicU8, AtomicU32, Ordering},
    },
    time::Duration,
};
use support::{Msg, TEST_PSK, client_session, frame};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream,
};

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 counting_server_with(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 counting_server(port: u16, rx: Arc<AtomicU32>) -> Arc<OmniVox<Msg>> {
    counting_server_with(port, rx, test_config()).await
}

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")
        .unwrap();
    assert_eq!(n, 0);
}

#[tokio::test]
async fn replayed_frame_is_dropped_and_the_connection_dies() {
    let port = 18301;
    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 good = frame(&s.send_key, 1, b"ok-transfer-1000");
    w.write_all(&good).await.unwrap();
    tokio::time::sleep(Duration::from_millis(200)).await;
    assert_eq!(rx.load(Ordering::SeqCst), 1, "first delivery");

    // Reinject the verbatim captured frame.
    w.write_all(&good).await.unwrap();
    expect_eof(&mut r).await; // server must drop the connection outright
    tokio::time::sleep(Duration::from_millis(200)).await;
    assert_eq!(
        rx.load(Ordering::SeqCst),
        1,
        "replay must never reach the app"
    );
    server.stop().await;
}

#[tokio::test]
async fn frames_from_another_session_are_undecryptable() {
    // Frame portability: keys are session-bound by the handshake nonces,
    // so a frame valid on connection A is junk on connection B.
    let port = 18302;
    let rx = Arc::new(AtomicU32::new(0));
    let server = counting_server(port, Arc::clone(&rx)).await;

    let (_r1, _w1, s1) = client_session(port, &TEST_PSK).await;
    let (_r2, mut w2, s2) = client_session(port, &TEST_PSK).await;

    let portable = frame(&s1.send_key, 1, b"ok-portable");
    w2.write_all(&portable).await.unwrap();
    // Auth failure with a matching counter would continue (junk policy), but
    // the counter check passes while decryption fails → logged, skipped;
    // the connection survives and the payload never fires.
    tokio::time::sleep(Duration::from_millis(200)).await;
    assert_eq!(
        rx.load(Ordering::SeqCst),
        0,
        "cross-session frame must not decrypt"
    );
    // connection B is still alive with its own keys:
    w2.write_all(&frame(&s2.send_key, 1, b"ok-native"))
        .await
        .unwrap();
    tokio::time::sleep(Duration::from_millis(200)).await;
    assert_eq!(
        rx.load(Ordering::SeqCst),
        1,
        "native frame on B must arrive"
    );
    server.stop().await;
}

#[tokio::test]
async fn wrong_version_hello_is_closed() {
    let port = 18303;
    let rx = Arc::new(AtomicU32::new(0));
    let server = counting_server(port, rx).await;

    let raw = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
    let (mut r, mut w) = raw.into_split();
    let mut hello = [0u8; 69];
    hello[..4].copy_from_slice(b"XXXX");
    hello[4] = 99;
    w.write_all(&hello).await.unwrap();
    expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn v1_traffic_is_rejected_at_the_handshake() {
    // A v1 OmniVox client speaks u32-length-prefixed AeadPackets. The v2
    // server must refuse before any allocation path.
    let port = 18304;
    let rx = Arc::new(AtomicU32::new(0));
    let server = counting_server_with(
        port,
        rx,
        Config {
            handshake_timeout_ms: 700,
            ..test_config()
        },
    )
    .await;

    let raw = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
    let (mut r, mut w) = raw.into_split();
    w.write_all(&36u32.to_le_bytes()).await.unwrap(); // v1 header
    w.write_all(&[0u8; 36]).await.unwrap();
    expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn panicking_callback_does_not_kill_the_connection() {
    let port = 18305;
    let rx = Arc::new(AtomicU32::new(0));
    let server = Arc::new(OmniVox::<Msg>::new());
    server.set_encrypt_key(TEST_PSK);
    let ok_rx = Arc::clone(&rx);
    let data = Arc::new(move |msg: &Msg, _ctx: &ConnectionContext| {
        if msg.0 == b"boom" {
            let _v: [u8; 8] = msg.0[..8].try_into().unwrap(); // explodes (len 4)
        }
        if msg.0.starts_with(b"ok") {
            ok_rx.fetch_add(1, Ordering::SeqCst);
        }
    });
    server
        .start_server(
            port,
            test_config(),
            data,
            Arc::new(|_c| {}),
            Arc::new(|_c| Decision::Allow),
        )
        .await
        .unwrap();

    let (mut r, mut w, s) = client_session(port, &TEST_PSK).await;
    w.write_all(&frame(&s.send_key, 1, b"boom")).await.unwrap();
    w.write_all(&frame(&s.send_key, 2, b"ok-after-panic"))
        .await
        .unwrap();
    tokio::time::sleep(Duration::from_millis(300)).await;
    assert_eq!(
        rx.load(Ordering::SeqCst),
        1,
        "survives a panicking callback"
    );
    let mut buf = [0u8; 1];
    assert!(
        tokio::time::timeout(Duration::from_millis(300), r.read(&mut buf))
            .await
            .is_err(),
        "connection must still be open"
    );
    server.stop().await;
}

#[tokio::test]
async fn malicious_server_silence_times_out_the_client() {
    // v2 retires B5: the client side now has the same idle/stall discipline.
    let port = 18306;
    let listener = tokio::net::TcpListener::bind(("127.0.0.1", port))
        .await
        .unwrap();

    let accepted = tokio::spawn(async move {
        let (stream, _) = listener.accept().await.unwrap();
        let (mut r, mut w) = stream.into_split();
        // Complete the handshake as the server, then go silent forever.
        perform(&mut r, &mut w, &TEST_PSK, true).await.unwrap();
        tokio::time::sleep(Duration::from_secs(60)).await;
    });

    let client = OmniVox::<Msg>::new();
    client.set_encrypt_key(TEST_PSK);
    let dropped = Arc::new(AtomicU8::new(0));
    let d = Arc::clone(&dropped);
    let cfg = Config {
        idle_timeout_ms: 400,
        ..test_config()
    };
    client
        .connect_to_server(
            "127.0.0.1",
            port,
            cfg,
            Some(Arc::new(|_m: &Msg, _c: &ConnectionContext| {})),
            Arc::new(move |ctx: &ConnectionContext| {
                if ctx.connected == ConnectionStatus::NotConnected {
                    d.store(1, Ordering::SeqCst);
                }
            }),
        )
        .await
        .unwrap();

    tokio::time::sleep(Duration::from_millis(1500)).await;
    assert_eq!(
        dropped.load(Ordering::SeqCst),
        1,
        "client must time out a silent server and report NotConnected"
    );
    accepted.abort();
    client.stop().await;
}