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

//! SRT on loopback, this crate's caller against its listener: every payload delivered in
//! order and whole, with no loss and through a proxy that drops datagrams each way; the
//! account of every sequence number; encryption, a wrong passphrase, a missing one, the
//! stream ID reaching the admission and its refusal codes; and the latency window: a packet
//! is not delivered before it.

use joule_contrib::lossy::{Counts, Impair, Proxy};
use joule_contrib::srt::packet::{rej, Packet};
use std::collections::BTreeSet;
use joule_contrib::srt::{self, conn::CallError, Admitted, CallConfig, ConnRequest, ListenerConfig, Role};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
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 * 7 + k as u64 * 13 + j as u64) % 251) as u8;
        }
    }
    p
}

async fn listener(pass: Option<&str>, latency: u16, seen: Arc<Mutex<Vec<ConnRequest>>>) -> srt::Listener {
    let pass = pass.map(str::to_string);
    let adm = move |r: &ConnRequest| {
        seen.lock().unwrap().push(r.clone());
        match r.stream_id.as_deref() {
            Some("#!::m=request") => Err(rej::BAD_MODE),
            Some("#!::r=nobody,m=publish") => Err(rej::NOTFOUND),
            _ => Ok(Admitted { passphrase: pass.clone(), role: Role::Receive, latency_ms: Some(latency), tag: 42 }),
        }
    };
    srt::Listener::bind(ListenerConfig { latency_ms: latency, ..ListenerConfig::default() }, Arc::new(adm)).await.unwrap()
}

/// Send `n` payloads at `pps` packets a second; receive them all; return what arrived.
async fn carry(tx: &joule_contrib::link::Link, rx: &mut joule_contrib::link::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;
            send.send(payload(i)).await.unwrap();
        }
    });
    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
}


/// The sequence numbers whose first send the proxy dropped on its way to the receiver:
/// each must be found missing and recovered. Control packets and retransmissions it
/// dropped are not among them.
fn first_sends_lost(c: &Counts) -> BTreeSet<u32> {
    c.dropped_log.lock().unwrap().iter().filter_map(|b| match Packet::read(b) {
        Some(Packet::Data(d)) if !d.rexmit => Some(d.seq),
        _ => None,
    }).collect()
}

/// The longest run of consecutive sequence numbers in `s`, 31-bit and wrapping.
fn longest_run(s: &BTreeSet<u32>) -> u32 {
    let (mut best, mut run, mut prev) = (0u32, 0u32, None::<u32>);
    for &q in s {
        run = if prev == Some(q.wrapping_sub(1)) { run + 1 } else { 1 };
        best = best.max(run);
        prev = Some(q);
    }
    // A run across the wrap: the top of the space followed by 0.
    if s.contains(&0x7fff_ffff) && s.contains(&0) {
        let top = s.iter().rev().zip(0..).take_while(|(q, i)| **q == 0x7fff_ffff - i).count() as u32;
        let low = s.iter().zip(0..).take_while(|(q, i)| **q == *i).count() as u32;
        best = best.max(top + low);
    }
    best
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn every_payload_arrives_in_order_and_is_accounted_for() {
    let seen = Arc::new(Mutex::new(Vec::new()));
    let mut l = listener(None, 120, seen.clone()).await;
    let cfg = CallConfig { stream_id: Some("#!::r=room/owner,m=publish,s=abc".into()), latency_ms: 80, ..CallConfig::default() };
    let tx = timeout(STEP, srt::call(l.addr, cfg)).await.unwrap().unwrap();
    let mut rx = timeout(STEP, l.accept()).await.unwrap().unwrap();
    assert_eq!(rx.info.stream_id.as_deref(), Some("#!::r=room/owner,m=publish,s=abc"));
    assert_eq!(rx.info.tag, 42);
    assert_eq!(rx.info.latency_us, 120_000, "the larger of the two ends' latencies");
    assert_eq!(seen.lock().unwrap()[0].latency_ms, 80);
    let got = carry(&tx, &mut rx, 400, 500).await;
    assert_eq!(got.len(), 400);
    assert!(got.iter().enumerate().all(|(i, p)| *p == payload(i as u64)), "every payload whole and in order");
    tokio::time::sleep(Duration::from_millis(300)).await;
    let part = rx.account.lock().unwrap().cut(1, false);
    part.check().unwrap();
    assert_eq!((part.seq_count, part.first_try, part.dropped), (410, 410, 0), "the 400 and the ten after them");
    assert!(part.rtt.samples > 0, "ACK/ACKACK pairs measured the RTT");
    tx.close();
    let mut rx = rx;
    assert_eq!(timeout(STEP, rx.wait_closed()).await.unwrap(), "peer_shutdown");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn losses_each_way_are_recovered_within_the_window() {
    let seen = Arc::new(Mutex::new(Vec::new()));
    let mut l = listener(Some("correct horse battery"), 300, seen).await;
    // 3% of datagrams dropped toward the listener, 3% back (NAKs and ACKs among them), and
    // datagrams 50 to 54 in a row.
    let up = Impair { loss_ppm: 30_000, drop_index: &[50, 51, 52, 53, 54], ..Impair::default() };
    let down = Impair { loss_ppm: 30_000, ..Impair::default() };
    let px = Proxy::start(l.addr, up, down, 0x5eed).await.unwrap();
    let cfg = CallConfig { passphrase: Some("correct horse battery".into()), keylen: 32, latency_ms: 300, ..CallConfig::default() };
    let tx = timeout(STEP, srt::call(px.addr, cfg)).await.unwrap().unwrap();
    let mut rx = timeout(STEP, l.accept()).await.unwrap().unwrap();
    assert_eq!(rx.info.cipher_bits, 256);
    let got = carry(&tx, &mut rx, 1000, 500).await;
    assert!(px.up.dropped.load(std::sync::atomic::Ordering::Relaxed) > 20, "the proxy dropped datagrams");
    tokio::time::sleep(Duration::from_millis(500)).await;
    let part = rx.account.lock().unwrap().cut(1, false);
    part.check().unwrap();
    assert!(part.delivered >= got.len() as u64);
    // Everything was recovered: the window is many round trips on loopback.
    assert_eq!(part.dropped, 0, "{part:?}");
    assert_eq!(got.len(), 1000);
    assert!(got.iter().enumerate().all(|(i, p)| *p == payload(i as u64)));
    // The burst of five dropped datagrams holds as many first sends of data as the timing of
    // the control packets between them leaves: count them, rather than assume all five.
    let lost = first_sends_lost(&px.up);
    assert!(lost.len() >= 2, "the path lost first sends of data: {lost:?}");
    assert_eq!(part.delivered - part.first_try, lost.len() as u64, "each first send lost was found and recovered: {part:?}");
    assert!(part.requests > 0 && part.retransmitted > 0);
    let burst = longest_run(&lost);
    assert!(part.runs.iter().any(|r| r.count >= burst), "the longest burst of losses ({burst}) is one run: {:?}", part.runs);
    assert!(part.recovery.max < 300_000 && part.margin_min > 0);
    px.stop();
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_packet_is_held_for_the_latency_window() {
    let seen = Arc::new(Mutex::new(Vec::new()));
    let mut l = listener(None, 400, seen).await;
    let tx = srt::call(l.addr, CallConfig { latency_ms: 400, ..CallConfig::default() }).await.unwrap();
    let mut rx = l.accept().await.unwrap();
    let t = Instant::now();
    tx.send.clone().unwrap().send(payload(0)).await.unwrap();
    let p = timeout(STEP, rx.data.recv()).await.unwrap().unwrap();
    let held = t.elapsed();
    assert_eq!(p, payload(0));
    assert!(held >= Duration::from_millis(380) && held < Duration::from_millis(900), "delivered at its TSBPD time: {held:?}");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn refusals_carry_their_codes() {
    let seen = Arc::new(Mutex::new(Vec::new()));
    let l = listener(Some("the right passphrase"), 120, seen.clone()).await;
    let call = |pass: Option<&str>, sid: Option<&str>| {
        let cfg = CallConfig { passphrase: pass.map(str::to_string), stream_id: sid.map(str::to_string), connect_timeout_ms: 3_000, ..CallConfig::default() };
        srt::call(l.addr, cfg)
    };
    assert_eq!(call(Some("a wrong passphrase"), None).await.err(), Some(CallError::Rejected(rej::BADSECRET)));
    assert_eq!(call(None, None).await.err(), Some(CallError::Rejected(rej::UNSECURE)));
    assert_eq!(call(Some("the right passphrase"), Some("#!::m=request")).await.err(), Some(CallError::Rejected(rej::BAD_MODE)));
    assert_eq!(call(Some("the right passphrase"), Some("#!::r=nobody,m=publish")).await.err(), Some(CallError::Rejected(rej::NOTFOUND)));
    assert!(call(Some("the right passphrase"), None).await.is_ok());
    assert_eq!(l.refused.load(std::sync::atomic::Ordering::Relaxed), 4);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_listener_that_sends_and_a_caller_that_receives() {
    let adm = |_: &ConnRequest| Ok(Admitted { passphrase: Some("both ways round".into()), role: Role::Send, latency_ms: None, tag: 0 });
    let mut l = srt::Listener::bind(ListenerConfig::default(), Arc::new(adm)).await.unwrap();
    let px = Proxy::start(l.addr, Impair { loss_ppm: 20_000, ..Impair::default() }, Impair { loss_ppm: 20_000, ..Impair::default() }, 99).await.unwrap();
    let cfg = CallConfig { passphrase: Some("both ways round".into()), role: Role::Receive, latency_ms: 250, ..CallConfig::default() };
    let mut rx = timeout(STEP, srt::call(px.addr, cfg)).await.unwrap().unwrap();
    let tx = timeout(STEP, l.accept()).await.unwrap().unwrap();
    let got = carry(&tx, &mut rx, 600, 400).await;
    assert_eq!(got.len(), 600);
    assert!(got.iter().enumerate().all(|(i, p)| *p == payload(i as u64)));
    let s = tx.sent.lock().unwrap().clone();
    assert!(s.retransmitted > 0 && s.requests > 0, "{s:?}");
}