omnivox-rs/tests/protocol_v2.rs

//! Attack regressions for protocol v2.1: every red-team finding class must
//! fail closed, provably, with zero callbacks fired on the attack input.

mod support;

use omnivox::{Config, ConnectionStatus, Decision};
use std::{
    sync::{
        Arc,
        atomic::{AtomicU8, AtomicU32, Ordering},
    },
    time::Duration,
};
use support::{Msg, TEST_PSK, client_session, msg_frame, raw_frame, stream_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_stream_frame: 8192,
        max_connections: 16,
        allow_insecure_defaults: true,
    }
}

async fn counting_server(port: u16, rx: Arc<AtomicU32>) -> omnivox::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);
            }
        })
        .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")
        .unwrap();
    assert_eq!(n, 0);
}

#[tokio::test]
async fn replayed_frame_kills_the_connection() {
    let port = 18501;
    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 = msg_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: nonce is derived from our own
    // counter (now 2), so authentication fails outright.
    w.write_all(&good).await.unwrap();
    expect_eof(&mut r).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_rejected() {
    let port = 18502;
    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 (mut r2, mut w2, _s2) = client_session(port, &TEST_PSK).await;

    // A perfectly valid frame from session A, injected into session B.
    let portable = msg_frame(&s1.send_key, 1, b"ok-portable");
    w2.write_all(&portable).await.unwrap();
    expect_eof(&mut r2).await; // v2.1: auth failure is fatal
    assert_eq!(
        rx.load(Ordering::SeqCst),
        0,
        "cross-session frame must not decrypt"
    );
    server.stop().await;
}

#[tokio::test]
async fn frame_type_tamper_is_rejected() {
    // Message -> stream downgrade attempt: type byte is AAD, so this is an
    // authentication failure, not a lane switch.
    let port = 18503;
    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;
    let mut frame = msg_frame(&s.send_key, 1, b"ok-msg");
    let type_idx = frame.len() - 17; // [len|type|ct|tag] -> type at offset 4
    frame[4] = 1;
    let _ = type_idx;
    w.write_all(&frame).await.unwrap();
    expect_eof(&mut r).await;
    server.stop().await;
}

#[tokio::test]
async fn wrong_version_hello_is_closed() {
    let port = 18504;
    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() {
    let port = 18505;
    let rx = Arc::new(AtomicU32::new(0));
    let server = omnivox::server::<Msg>()
        .key(TEST_PSK)
        .config(Config {
            handshake_timeout_ms: 700,
            ..test_config()
        })
        .on_data(move |_: &Msg, _| {
            rx.fetch_add(1, Ordering::SeqCst);
        })
        .serve(port)
        .await
        .unwrap();

    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_callbacks_never_kill_the_connection() {
    let port = 18506;
    let rx = Arc::new(AtomicU32::new(0));
    let ok_rx = Arc::clone(&rx);
    let server = omnivox::server::<Msg>()
        .key(TEST_PSK)
        .config(test_config())
        .on_data(move |msg: &Msg, _ctx| {
            if msg.0 == b"boom" {
                panic!("handler exploded on attacker input");
            }
            if msg.0.starts_with(b"ok") {
                ok_rx.fetch_add(1, Ordering::SeqCst);
            }
        })
        .on_stream(|_b, _ctx| {
            panic!("stream handler exploded");
        })
        .serve(port)
        .await
        .unwrap();

    let (mut r, mut w, s) = client_session(port, &TEST_PSK).await;
    w.write_all(&msg_frame(&s.send_key, 1, b"boom"))
        .await
        .unwrap();
    w.write_all(&stream_frame(&s.send_key, 2, b"stream-boom"))
        .await
        .unwrap();
    w.write_all(&msg_frame(&s.send_key, 3, b"ok-after-panics"))
        .await
        .unwrap();
    tokio::time::sleep(Duration::from_millis(300)).await;
    assert_eq!(
        rx.load(Ordering::SeqCst),
        1,
        "connection survives two panics"
    );
    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() {
    let port = 18507;
    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.
        omnivox::handshake::perform(&mut r, &mut w, &TEST_PSK, true)
            .await
            .unwrap();
        tokio::time::sleep(Duration::from_secs(60)).await;
    });

    let dropped = Arc::new(AtomicU8::new(0));
    let d = Arc::clone(&dropped);
    let client = omnivox::client::<Msg>()
        .key(TEST_PSK)
        .config(Config {
            idle_timeout_ms: 400,
            ..test_config()
        })
        .on_data(|_m: &Msg, _c| {})
        .on_state(move |ctx| {
            if ctx.status() == ConnectionStatus::NotConnected {
                d.store(1, Ordering::SeqCst);
            }
        })
        .connect("127.0.0.1", port)
        .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;
}

#[tokio::test]
async fn messages_and_streams_interleave_in_order() {
    // One sequence across both lanes: strict per-direction ordering of the
    // exact interleaving the sender produced.
    let port = 18508;
    let tape = Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
    let t = Arc::clone(&tape);
    let t2 = Arc::clone(&tape);
    let server = omnivox::server::<Msg>()
        .key(TEST_PSK)
        .config(test_config())
        .on_data(move |msg: &Msg, _ctx| {
            t.lock()
                .unwrap()
                .push(format!("msg:{}", String::from_utf8_lossy(&msg.0)));
        })
        .on_stream(move |bytes, _ctx| {
            t2.lock()
                .unwrap()
                .push(format!("stream:{}", String::from_utf8_lossy(bytes)));
        })
        .serve(port)
        .await
        .unwrap();

    let (_r, mut w, s) = client_session(port, &TEST_PSK).await;
    w.write_all(&msg_frame(&s.send_key, 1, b"alpha"))
        .await
        .unwrap();
    w.write_all(&stream_frame(&s.send_key, 2, b"S1"))
        .await
        .unwrap();
    w.write_all(&msg_frame(&s.send_key, 3, b"beta"))
        .await
        .unwrap();
    w.write_all(&stream_frame(&s.send_key, 4, b"S2"))
        .await
        .unwrap();
    tokio::time::sleep(Duration::from_millis(400)).await;
    assert_eq!(
        *tape.lock().unwrap(),
        vec![
            "msg:alpha".to_string(),
            "stream:S1".to_string(),
            "msg:beta".to_string(),
            "stream:S2".to_string(),
        ]
    );
    server.stop().await;
}

#[tokio::test]
async fn lying_decoder_length_is_skipped_connection_survives() {
    // A frame that authenticates but whose payload claims a huge bincode
    // length: decode fails, frame is skipped (authenticated junk policy),
    // connection continues.
    let port = 18509;
    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 lying = {
        // Msg(Vec<u8>) bincode: u64 LE length then bytes. Claim 2^40.
        let mut plain = Vec::new();
        plain.extend_from_slice(&(1u64 << 40).to_le_bytes());
        plain.extend_from_slice(b"ha");
        raw_frame(&s.send_key, 1, omnivox::FrameType::Message, &plain)
    };
    w.write_all(&lying).await.unwrap();
    w.write_all(&msg_frame(&s.send_key, 2, b"ok-real"))
        .await
        .unwrap();
    tokio::time::sleep(Duration::from_millis(300)).await;
    assert_eq!(
        rx.load(Ordering::SeqCst),
        1,
        "junk skipped, real frame arrives"
    );
    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 decision_deny_rejects_after_handshake() {
    // A fully authenticated client can still be refused; it must see EOF
    // and no data path may open.
    let port = 18510;
    let seen = Arc::new(AtomicU8::new(0));
    let s = Arc::clone(&seen);
    let server = omnivox::server::<Msg>()
        .key(TEST_PSK)
        .config(test_config())
        .on_data(move |_: &Msg, _| {
            s.store(9, Ordering::SeqCst); // must never run
        })
        .on_connect(move |_ctx| Decision::Deny)
        .serve(port)
        .await
        .unwrap();

    let (mut r, mut w, sess) = client_session(port, &TEST_PSK).await;
    // handshake completed (client_session); the Deny must have closed us.
    expect_eof(&mut r).await;
    let _ = w
        .write_all(&msg_frame(&sess.send_key, 1, b"ok-after-deny"))
        .await;
    tokio::time::sleep(Duration::from_millis(200)).await;
    assert_eq!(
        seen.load(Ordering::SeqCst),
        0,
        "denied client never reaches data path"
    );
    server.stop().await;
}