open-standards/joule-code/joule-code-rs/crates/joule-moq-ingest/tests/contrib_rist.rs

//! RIST into a room's tracks, end to end on loopback: a publisher POSTs `/rist` for the
//! Simple Profile (a port pair of its own and a CNAME) and for the Main Profile (a tunnel
//! port, a CNAME and a PSK passphrase), and sends a transport stream of H.264 and AAC; a
//! subscriber at the relay verifies every receipt and gets every unit as sent; the link
//! record names the profile, the PSK key length and the receive window. A flow with another
//! CNAME opens no session.

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

mod common;

use common::*;
use joule_contrib::lossy::{Impair, Proxy};
use joule_contrib::rist::{self, Profile, SendConfig};
use joule_contrib::ts;
use joule_moq_ingest::frame::Codec;
use joule_moq_ingest::linkrec;
use joule_moq_ingest::room::media;
use std::time::Duration;
use tokio::time::timeout;

async fn open(w: &World, auth: &Auth, profile: &str) -> serde_json::Value {
    let body = serde_json::json!({"profile": profile}).to_string();
    let r = auth.send("POST", &format!("{}/rist", w.base()), &[("content-type", "application/json")], body.as_bytes()).await;
    assert_eq!(r.status, 201, "POST /rist: {}", r.text());
    serde_json::from_slice(&r.body).unwrap()
}

async fn one(profile: Profile) {
    let w = World::start(|c| {
        c.contrib.rist.enabled = true;
        c.contrib.rist.aes = 256;
        c.contrib.rist.buffer_ms = 800;
        c.contrib.record_interval_ms = 1_000;
    })
    .await;
    let auth = w.auth(&w.publisher, 0);
    let s = open(&w, &auth, if profile == Profile::Main { "main" } else { "simple" }).await;
    let port = s["port"].as_u64().unwrap() as u16;
    let target: std::net::SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
    assert_eq!(s["passphrase"].is_string(), profile == Profile::Main);
    // A wrong CNAME first: no session.
    let stranger = rist::send(target, SendConfig { profile, cname: "not-the-session".into(), secret: s["passphrase"].as_str().map(str::to_string), aes_256: true, ..SendConfig::default() }).await.unwrap();
    let tx = stranger.send.clone().unwrap();
    for _ in 0..30 {
        let mut p = vec![0x47u8; 1316];
        p[1] = 0x1f;
        p[2] = 0xff;
        let _ = tx.send(p).await;
    }
    tokio::time::sleep(Duration::from_millis(500)).await;
    assert_eq!(w.gateway.gateway.contrib.active.load(std::sync::atomic::Ordering::SeqCst), 0, "no session for another CNAME");
    stranger.close();
    tokio::time::sleep(Duration::from_millis(300)).await;

    let (target, px) = if profile == Profile::Main {
        let imp = Impair { loss_ppm: 25_000, spare_first: 16, ..Impair::default() };
        let px = Proxy::start(target, imp, imp, 5).await.unwrap();
        (px.addr, Some(px))
    } else {
        (target, None)
    };
    let cfg = SendConfig { profile, cname: s["cname"].as_str().unwrap().into(), secret: s["passphrase"].as_str().map(str::to_string), aes_256: true, buffer_ms: 1500, ..SendConfig::default() };
    let link = rist::send(target, cfg).await.unwrap();
    // The flow is the session's once its first RTCP names the CNAME; then the tracks are
    // published, and a subscriber joins before the first frame.
    assert!(contrib_live(&w, s["session"].as_str().unwrap()).await, "the session publishes its tracks");
    let watcher = watch(&w).await;
    let sent = send_ts(&link, 3, 30, ts::ST_H264).await;
    tokio::time::sleep(Duration::from_millis(1_500)).await;
    link.close();
    // RIST says nothing when a sender stops: the session ends when the flow goes idle.
    assert!(until(Duration::from_secs(15), || watcher.ended.lock().unwrap().len() == 2).await, "PUBLISH_DONE on both media tracks");
    settle(&watcher, Duration::from_millis(500)).await;
    let video = watcher.tracks.iter().find(|t| t.kind == media::VIDEO).unwrap();
    let audio = watcher.tracks.iter().find(|t| t.kind == media::AUDIO).unwrap();
    let (vf, _) = verified(&watcher, video, &w.gateway_key, Codec::Avc);
    let (af, _) = verified(&watcher, audio, &w.gateway_key, Codec::Aac);
    same(&vf, &sent.video, "every access unit");
    same(&af, &sent.audio, "every ADTS frame");
    let ps = linkrec::verify_series(&w.gateway.gateway.link_records()).expect("the parts verify as a series");
    let last = ps.last().unwrap();
    assert!(last.last && last.ended == linkrec::ended::IDLE);
    let (want_proto, want_cipher) = if profile == Profile::Main { (3, 256) } else { (2, 0) };
    assert!(ps.iter().all(|p| p.protocol == want_proto && p.cipher == want_cipher && p.window == 800_000));
    if px.is_some() {
        let recovered: u64 = ps.iter().map(|p| p.account.delivered - p.account.first_try).sum();
        assert!(recovered > 0, "losses on the tunnel were recovered");
        assert_eq!(ps.iter().map(|p| p.account.dropped).sum::<u64>(), 0);
    }
    w.stop().await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_simple_profile_flow_reaches_the_relay() {
    timeout(Duration::from_secs(120), one(Profile::Simple)).await.expect("in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_main_profile_tunnel_under_psk_reaches_the_relay() {
    timeout(Duration::from_secs(120), one(Profile::Main)).await.expect("in time");
}