open-standards/joule-code/joule-code-rs/crates/joule-contrib/tests/rist_link.rs

//! RIST on loopback, this crate's sender against its receiver: the Simple Profile on P and
//! P+1 with range and bitmask requests through proxies that drop datagrams; the Main
//! Profile's tunnel in the clear and under PSK (128 and 256), in the VSF and the legacy
//! formats, with loss; a CNAME that is not the expected one; a wrong passphrase.

#![cfg(feature = "rist")]

use joule_contrib::link::Link;
use joule_contrib::lossy::{Impair, Proxy};
use joule_contrib::rist::recv::RecvConfig;
use joule_contrib::rist::{self, NackKind, Profile, RistListener, RistListenerConfig, SendConfig};
use std::net::SocketAddr;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::time::timeout;

const STEP: Duration = Duration::from_secs(20);

fn payload(i: u64) -> Vec<u8> {
    let mut p = vec![0u8; 1316];
    for (k, c) in p.chunks_mut(188).enumerate() {
        c[0] = 0x47;
        for (j, b) in c.iter_mut().enumerate().skip(1) {
            *b = ((i * 11 + k as u64 * 3 + j as u64) % 251) as u8;
        }
    }
    p
}

async fn carry(tx: &Link, rx: &mut Link, n: u64, pps: u64) -> Vec<Vec<u8>> {
    let send = tx.send.clone().unwrap();
    let t = tokio::spawn(async move {
        let start = tokio::time::Instant::now();
        // Ten more than are counted: a receiver cannot know of a packet lost after the
        // last it receives.
        for i in 0..n + 10 {
            tokio::time::sleep_until(start + Duration::from_micros(i * 1_000_000 / pps)).await;
            if send.send(payload(i)).await.is_err() {
                return;
            }
        }
    });
    let mut got = Vec::new();
    while got.len() < n as usize + 10 {
        match timeout(Duration::from_secs(5), rx.data.recv()).await {
            Ok(Some(p)) => got.push(p),
            _ => break,
        }
    }
    t.await.unwrap();
    got.truncate(n as usize);
    got
}

fn listener_cfg(profile: Profile, secret: Option<&str>, cname: Option<&str>, nack: NackKind) -> RistListenerConfig {
    RistListenerConfig {
        profile,
        bind_ip: "127.0.0.1".parse().unwrap(),
        port: 0,
        recv: RecvConfig { buffer_us: 600_000, nack, ..RecvConfig::default() },
        secret: secret.map(str::to_string),
        expect_cname: cname.map(str::to_string),
        tag: 0,
    }
}

async fn simple(nack: NackKind, seed: u64) {
    let mut l = RistListener::bind(listener_cfg(Profile::Simple, None, Some("token-1"), nack)).await.unwrap();
    let p = l.addr;
    // Two proxies, one per port. The sender's RTCP must reach P + 1, so the RTCP proxy
    // listens on an address whose port is the RTP proxy's plus one.
    let imp = Impair { loss_ppm: 40_000, drop_index: &[30, 31, 32, 33], spare_first: 8, ..Impair::default() };
    let (rtp_px, rtcp_px) = proxy_pair(p, imp, Impair { loss_ppm: 40_000, ..Impair::default() }, seed).await;
    let tx = rist::send(rtp_px.addr, SendConfig { cname: "token-1".into(), buffer_ms: 1000, ..SendConfig::default() }).await.unwrap();
    let mut rx = timeout(STEP, l.accept()).await.unwrap().unwrap();
    assert_eq!(rx.info.cname.as_deref(), Some("token-1"));
    let got = carry(&tx, &mut rx, 800, 400).await;
    assert!(rtp_px.up.dropped.load(Ordering::Relaxed) > 10);
    tokio::time::sleep(Duration::from_millis(900)).await;
    let part = rx.account.lock().unwrap().cut(1, false);
    part.check().unwrap();
    assert_eq!(got.len(), 800, "{part:?}");
    assert!(got.iter().enumerate().all(|(i, q)| *q == payload(i as u64)));
    assert!(part.delivered - part.first_try >= 4 && part.requests > 0 && part.retransmitted > 0, "{part:?}");
    assert!(rx.rtt() > 0, "the RTT echo measured the round trip");
    let s = tx.sent.lock().unwrap().clone();
    assert!(s.retransmitted > 0);
    rtp_px.stop();
    rtcp_px.stop();
}

/// Proxies to P and P + 1 on two ports one apart, the first even, as a Simple Profile
/// sender addresses them.
async fn proxy_pair(p: SocketAddr, up: Impair, down: Impair, seed: u64) -> (Proxy, Proxy) {
    for _ in 0..100 {
        let probe = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
        let q = probe.local_addr().unwrap().port();
        drop(probe);
        let q = if q.is_multiple_of(2) { q } else { q + 1 };
        if q >= 65_534 {
            continue;
        }
        let Ok(a) = Proxy::start_on(SocketAddr::new(p.ip(), q), p, up, down, seed).await else { continue };
        match Proxy::start_on(SocketAddr::new(p.ip(), q + 1), SocketAddr::new(p.ip(), p.port() + 1), down, down, seed + 1).await {
            Ok(b) => return (a, b),
            Err(_) => a.stop(),
        }
    }
    panic!("no adjacent proxy ports");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_simple_profile_recovers_with_range_requests() {
    simple(NackKind::Range, 3).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_simple_profile_recovers_with_bitmask_requests() {
    simple(NackKind::Bitmask, 4).await;
}

async fn main_profile(secret: Option<&str>, aes_256: bool, vsf: bool, loss: u32) {
    let mut l = RistListener::bind(listener_cfg(Profile::Main, secret, Some("tunnel-token"), NackKind::Range)).await.unwrap();
    let imp = Impair { loss_ppm: loss, spare_first: 16, ..Impair::default() };
    let px = Proxy::start(l.addr, imp, imp, 17).await.unwrap();
    let cfg = SendConfig { profile: Profile::Main, cname: "tunnel-token".into(), secret: secret.map(str::to_string), aes_256, vsf, buffer_ms: 1000, ..SendConfig::default() };
    let tx = rist::send(px.addr, cfg).await.unwrap();
    let mut rx = timeout(STEP, l.accept()).await.unwrap().unwrap();
    assert_eq!(rx.info.cipher_bits, if secret.is_some() { if aes_256 { 256 } else { 128 } } else { 0 });
    let got = carry(&tx, &mut rx, 600, 400).await;
    tokio::time::sleep(Duration::from_millis(900)).await;
    let part = rx.account.lock().unwrap().cut(1, false);
    part.check().unwrap();
    assert_eq!(got.len(), 600, "{part:?}");
    assert!(got.iter().enumerate().all(|(i, q)| *q == payload(i as u64)));
    if loss > 0 {
        assert!(part.delivered > part.first_try, "losses were recovered: {part:?}");
    }
    px.stop();
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_main_profile_in_the_clear() {
    main_profile(None, false, true, 30_000).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_main_profile_under_psk_128_and_256() {
    main_profile(Some("Reliable Internet Stream Transport"), false, true, 30_000).await;
    main_profile(Some("another tunnel passphrase"), true, true, 0).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_main_profile_in_the_legacy_format() {
    main_profile(Some("legacy formats"), false, false, 20_000).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_wrong_cname_and_the_wrong_passphrase_are_refused() {
    let mut l = RistListener::bind(listener_cfg(Profile::Main, Some("the tunnel's passphrase"), Some("tok"), NackKind::Range)).await.unwrap();
    let bad = rist::send(l.addr, SendConfig { profile: Profile::Main, cname: "tok".into(), secret: Some("not the passphrase".into()), ..SendConfig::default() }).await.unwrap();
    let send = bad.send.clone().unwrap();
    for i in 0..50 {
        send.send(payload(i)).await.unwrap();
    }
    assert!(timeout(Duration::from_millis(1500), l.accept()).await.is_err(), "no flow under another passphrase");
    assert!(l.refused.undecryptable.load(Ordering::Relaxed) > 0);
    bad.close();

    let mut s = RistListener::bind(listener_cfg(Profile::Simple, None, Some("tok"), NackKind::Range)).await.unwrap();
    let other = rist::send(s.addr, SendConfig { cname: "another".into(), ..SendConfig::default() }).await.unwrap();
    let send = other.send.clone().unwrap();
    for i in 0..20 {
        send.send(payload(i)).await.unwrap();
    }
    assert!(timeout(Duration::from_millis(1500), s.accept()).await.is_err(), "no flow under another CNAME");
    assert!(s.refused.cname.load(Ordering::Relaxed) > 0);
}