omnivox-rs/tests/loopback.rs

//! Loopback behavior tests mirroring the contract's runtime clauses
//! (B1 unreasonable header, B2 tampered packet survives, framing).
//! Raw keys via `set_encrypt_key` keep these fast (no scrypt).

use omnivox::{
    AeadPacket, Cereal, ConnectionContext, Decision, MAX_MESSAGE_SIZE, OmniVox, encrypt,
};
use std::{
    sync::{
        Arc,
        atomic::{AtomicU32, Ordering},
    },
    time::Duration,
};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream,
};

#[derive(Debug, Default)]
struct Msg(Vec<u8>);

impl Cereal for Msg {
    fn encode(&self, out: &mut Vec<u8>) {
        omnivox::write_byte_vec(out, &self.0);
    }
    fn decode(bytes: &[u8]) -> Result<Self, omnivox::CodecError> {
        Ok(Self(omnivox::Reader::new(bytes).byte_vec()?))
    }
}

fn frame(key: &[u8; 32], msg: &[u8]) -> Vec<u8> {
    let mut plain = Vec::new();
    Msg(msg.to_vec()).encode(&mut plain);
    let packet = encrypt(key, &plain);
    let mut bytes = Vec::new();
    packet.encode(&mut bytes);
    let mut out = Vec::new();
    out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
    out.extend_from_slice(&bytes);
    out
}

async fn start_server(port: u16, rx: Arc<AtomicU32>) -> Arc<OmniVox<Msg>> {
    let server = Arc::new(OmniVox::<Msg>::new());
    server.set_encrypt_key([42u8; 32]);
    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,
            1,
            2000,
            data,
            Arc::new(|_ctx| {}),
            Arc::new(|_ctx| Decision::Allow),
        )
        .await
        .unwrap();
    server
}

fn free_port(hint: u16) -> u16 {
    // bind-and-drop probe; tests run sequentially per port
    hint
}

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

    let client = OmniVox::<Msg>::new();
    client.set_encrypt_key([42u8; 32]);
    let got = Arc::new(AtomicU32::new(0));
    let g = Arc::clone(&got);
    client
        .connect_to_server(
            "127.0.0.1",
            port,
            Some(Arc::new(move |_d, _c| {
                g.fetch_add(1, Ordering::SeqCst);
            })),
            Arc::new(|_c| {}),
        )
        .await
        .unwrap();

    // Valid through the API.
    assert!(client.send_to_server(&Msg(b"ok1".to_vec())).await);
    // Tampered raw frame: flip a ciphertext byte.
    let mut bad = frame(&[42u8; 32], b"ok2");
    let last = bad.len() - 1;
    bad[last] ^= 0x01;
    let mut raw = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
    raw.write_all(&bad).await.unwrap();
    // Valid again on the tampered connection's server slot (new conn).
    assert!(client.send_to_server(&Msg(b"ok3".to_vec())).await);

    tokio::time::sleep(Duration::from_millis(300)).await;
    assert_eq!(
        rx.load(Ordering::SeqCst),
        2,
        "tampered frame must not decode"
    );
    // Contract B2: the tampering connection is NOT dropped — server keeps
    // its header wait armed (rx-timeout is 2 s, so a 500 ms read must time
    // out without EOF).
    let r = tokio::time::timeout(Duration::from_millis(500), raw.read(&mut [0u8; 1])).await;
    assert!(r.is_err(), "tampering connection must stay open (B2)");
    assert_eq!(
        got.load(Ordering::SeqCst),
        0,
        "server does not echo in this test"
    );
    server.stop().await;
    client.stop().await;
}

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

    let mut raw = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
    raw.write_all(&0u32.to_le_bytes()).await.unwrap();
    // C++ B1: server logs, disconnects, then the 0-byte read "succeeds" and
    // the empty packet fails to decode (logged). The socket must close.
    let mut buf = [0u8; 1];
    let n = tokio::time::timeout(Duration::from_secs(2), raw.read(&mut buf))
        .await
        .expect("server must close, not hang")
        .unwrap();
    assert_eq!(n, 0, "EOF after unreasonable header");
    server.stop().await;
}

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

    let mut raw = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
    raw.write_all(&(MAX_MESSAGE_SIZE + 1).to_le_bytes())
        .await
        .unwrap();
    let mut buf = [0u8; 1];
    let n = tokio::time::timeout(Duration::from_secs(2), raw.read(&mut buf))
        .await
        .expect("server must close, not hang")
        .unwrap();
    assert_eq!(n, 0);
    server.stop().await;
}

#[tokio::test]
async fn header_stall_times_out_server_side() {
    let port = free_port(18104);
    let rx = Arc::new(AtomicU32::new(0));
    // start_server timeout_ms = 300 for this one
    let server = Arc::new(OmniVox::<Msg>::new());
    server.set_encrypt_key([42u8; 32]);
    server
        .start_server(
            port,
            1,
            300,
            Arc::new(move |_: &Msg, _: &ConnectionContext| {
                rx.fetch_add(1, Ordering::SeqCst);
            }),
            Arc::new(|_c| {}),
            Arc::new(|_c| Decision::Allow),
        )
        .await
        .unwrap();

    let mut raw = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
    raw.write_all(&[0x10, 0x00]).await.unwrap(); // partial header, then silence
    let mut buf = [0u8; 1];
    let n = tokio::time::timeout(Duration::from_secs(3), raw.read(&mut buf))
        .await
        .expect("stall timeout must fire")
        .unwrap();
    assert_eq!(n, 0);
    server.stop().await;
}

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

#[test]
fn aead_packet_roundtrip_matches_wire_shape() {
    let pkt = AeadPacket {
        nonce: [9; 12],
        tag: [8; 16],
        payload: vec![7; 3],
    };
    let mut b = Vec::new();
    pkt.encode(&mut b);
    assert_eq!(b.len(), 12 + 16 + 8 + 3);
    assert_eq!(AeadPacket::decode(&b).unwrap(), pkt);
}