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

//! The reference tools into the gateway: a transport stream of H.264 or HEVC and AAC made
//! by a locally run encoder, sent by the reference SRT caller with the session's stream ID
//! and passphrase, and by the reference RIST sender in the Main Profile under PSK, each
//! through a path that drops datagrams. A subscriber at the relay verifies every receipt;
//! every access unit and ADTS frame of the file arrives; the encoder's own probe counts as
//! many video units; the link record verifies and shows the losses recovered.
//!
//! Ignored unless asked for; with `--ignored`, a tool not named is a failure, never a skip:
//!
//! ```text
//! WAI_SRT_LIVE_TRANSMIT=… WAI_RISTSENDER=… WAI_FFMPEG=… WAI_FFPROBE=… \
//! cargo test -p joule-moq-ingest --features rist --test contrib_interop -- --ignored
//! ```

mod common;

use common::*;
use joule_contrib::lossy::{Impair, Proxy};
use joule_contrib::ts::{self, Demux, Pacer};
use joule_moq_ingest::frame::Codec;
use joule_moq_ingest::linkrec;
use joule_moq_ingest::room::media;
use std::process::{Command, Stdio};
use std::time::Duration;

fn tool(var: &str) -> String {
    std::env::var(var).unwrap_or_else(|_| panic!("{var} names the tool this test runs; it is not set"))
}

struct Child(std::process::Child);

impl Drop for Child {
    fn drop(&mut self) {
        let _ = self.0.kill();
        let _ = self.0.wait();
    }
}

fn spawn(bin: &str, args: &[&str]) -> Child {
    Child(Command::new(bin).args(args).stdin(Stdio::null()).stdout(Stdio::null()).stderr(Stdio::null()).spawn().unwrap_or_else(|e| panic!("{bin}: {e}")))
}

/// Four seconds of test pattern and tone, encoded by the local encoder into a transport
/// stream: `codec` "libx264" or "libx265".
fn encode(codec: &str) -> Vec<u8> {
    let dir = std::env::temp_dir().join(format!("wai-contrib-interop-{}-{codec}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    let out = dir.join("in.ts");
    let st = Command::new(tool("WAI_FFMPEG"))
        .args(["-hide_banner", "-loglevel", "error", "-y", "-f", "lavfi", "-i", "testsrc2=size=640x360:rate=30", "-f", "lavfi", "-i", "sine=frequency=440:sample_rate=48000", "-t", "4", "-c:v", codec, "-preset", "veryfast", "-g", "30", "-bf", "0", "-pix_fmt", "yuv420p"])
        .args(if codec == "libx265" { vec!["-x265-params", "log-level=error"] } else { vec![] })
        .args(["-c:a", "aac", "-b:a", "96k", "-muxrate", "3000000", "-f", "mpegts"])
        .arg(&out)
        .status()
        .unwrap();
    assert!(st.success(), "the encoder runs");
    std::fs::read(out).unwrap()
}

/// The encoder's probe's count of the video stream's packets.
fn probe_video_units(ts: &[u8]) -> u64 {
    let f = std::env::temp_dir().join(format!("wai-contrib-probe-{}.ts", std::process::id()));
    std::fs::write(&f, ts).unwrap();
    let o = Command::new(tool("WAI_FFPROBE")).args(["-v", "error", "-select_streams", "v:0", "-count_packets", "-show_entries", "stream=nb_read_packets", "-of", "csv=p=0"]).arg(&f).output().unwrap();
    let out = String::from_utf8_lossy(&o.stdout).into_owned();
    out.split(|c: char| !c.is_ascii_digit()).find(|t| !t.is_empty()).and_then(|t| t.parse().ok()).unwrap_or_else(|| panic!("a packet count: {out:?} {}", String::from_utf8_lossy(&o.stderr)))
}

/// The access units and ADTS frames of a transport stream, as the gateway should carry
/// them.
fn units(ts_bytes: &[u8]) -> (Vec<Vec<u8>>, Vec<Vec<u8>>) {
    let mut d = Demux::new();
    let mut p = d.push(ts_bytes);
    p.extend(d.flush());
    let video = p.iter().filter(|x| matches!(x.stream_type, ts::ST_H264 | ts::ST_HEVC)).map(|x| x.data.clone()).collect();
    let audio = p.iter().filter(|x| x.stream_type == ts::ST_AAC_ADTS).flat_map(|x| ts::adts_frames(&x.data).into_iter().map(|f| f.to_vec()).collect::<Vec<_>>()).collect();
    (video, audio)
}

async fn feed(ts_bytes: &[u8], to: std::net::SocketAddr) {
    let s = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let chunks = Pacer::new(ts_bytes, 7).chunks().to_vec();
    let start = tokio::time::Instant::now();
    for (at, c) in chunks {
        tokio::time::sleep_until(start + Duration::from_micros(at)).await;
        s.send_to(&c, to).await.unwrap();
    }
}

fn free_port() -> u16 {
    let s = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
    s.local_addr().unwrap().port()
}

async fn check(w: &World, watcher: &Watcher, input: &[u8], codec: Codec) {
    assert!(until(Duration::from_secs(30), || watcher.ended.lock().unwrap().len() == 2).await, "PUBLISH_DONE on both media tracks");
    settle(watcher, Duration::from_millis(500)).await;
    let (want_v, want_a) = units(input);
    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, vg) = verified(watcher, video, &w.gateway_key, codec);
    let (af, _) = verified(watcher, audio, &w.gateway_key, Codec::Aac);
    // The link's own account first: a packet the link dropped explains a missing frame, and
    // one it did not drop places the loss outside the link.
    let ps = linkrec::verify_series(&w.gateway.gateway.link_records()).expect("the link record's parts verify");
    assert!(ps.last().unwrap().last);
    let sum = |f: fn(&linkrec::LinkRecord) -> u64| ps.iter().map(f).sum::<u64>();
    let (delivered, first_try, dropped, outstanding) = (sum(|p| p.account.delivered), sum(|p| p.account.first_try), sum(|p| p.account.dropped), sum(|p| p.account.outstanding));
    let bytes = sum(|p| p.account.bytes);
    eprintln!("the link's account: delivered {delivered} ({first_try} at the first try), dropped {dropped}, outstanding {outstanding}, {bytes} bytes of the input's {}", input.len());
    assert_eq!(dropped, 0, "the link dropped nothing");
    assert!(delivered > first_try, "losses were recovered");
    same(&vf, &want_v, "every access unit of the file");
    same(&af, &want_a, "every ADTS frame of the file");
    assert_eq!(vf.len() as u64, probe_video_units(input), "as many video units as the encoder's probe counts");
    assert!(vg >= 4, "a group per keyframe, at least one a second: {vg}");
    assert_eq!(w.relay.relay().stats().refused_at_door, 0);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
#[ignore = "runs the reference SRT tool and a local encoder: set WAI_SRT_LIVE_TRANSMIT, WAI_FFMPEG, WAI_FFPROBE"]
async fn the_reference_srt_caller_into_the_gateway() {
    for (codec, wai) in [("libx264", Codec::Avc), ("libx265", Codec::Hevc)] {
        let input = encode(codec);
        let w = World::start(|c| {
            c.contrib.srt.listen = Some("127.0.0.1:0".parse().unwrap());
            c.contrib.record_interval_ms = 1_000;
            c.contrib.srt.latency_ms = 300;
        })
        .await;
        let auth = w.auth(&w.publisher, 0);
        let r = auth.send("POST", &format!("{}/srt", w.base()), &[], b"").await;
        assert_eq!(r.status, 201, "{}", r.text());
        let s: serde_json::Value = serde_json::from_slice(&r.body).unwrap();
        let imp = Impair { loss_ppm: 20_000, spare_first: 10, ..Impair::default() };
        let px = Proxy::start(w.gateway.srt_addr.unwrap(), imp, imp, 41).await.unwrap();
        let sid = s["stream_id"].as_str().unwrap().replace('#', "%23");
        let url = format!("srt://127.0.0.1:{}?latency=300&streamid={sid}&passphrase={}&pbkeylen=16", px.addr.port(), s["passphrase"].as_str().unwrap());
        let udp_in = free_port();
        let child = spawn(&tool("WAI_SRT_LIVE_TRANSMIT"), &["-q", "-a", "no", &format!("udp://127.0.0.1:{udp_in}"), &url]);
        assert!(contrib_live(&w, s["session"].as_str().unwrap()).await, "the reference caller is admitted");
        let watcher = watch(&w).await;
        feed(&input, format!("127.0.0.1:{udp_in}").parse().unwrap()).await;
        tokio::time::sleep(Duration::from_millis(1_500)).await;
        // The reference caller ends the connection when it is stopped.
        drop(child);
        check(&w, &watcher, &input, wai).await;
        px.stop();
        w.stop().await;
    }
}

#[cfg(feature = "rist")]
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
#[ignore = "runs the reference RIST tool and a local encoder: set WAI_RISTSENDER, WAI_FFMPEG, WAI_FFPROBE"]
async fn the_reference_rist_sender_into_the_gateway() {
    let input = encode("libx264");
    let w = World::start(|c| {
        c.contrib.rist.enabled = true;
        c.contrib.rist.aes = 128;
        c.contrib.record_interval_ms = 1_000;
    })
    .await;
    let auth = w.auth(&w.publisher, 0);
    let r = auth.send("POST", &format!("{}/rist", w.base()), &[("content-type", "application/json")], br#"{"profile":"main"}"#).await;
    assert_eq!(r.status, 201, "{}", r.text());
    let s: serde_json::Value = serde_json::from_slice(&r.body).unwrap();
    let imp = Impair { loss_ppm: 20_000, spare_first: 16, ..Impair::default() };
    let target: std::net::SocketAddr = format!("127.0.0.1:{}", s["port"]).parse().unwrap();
    let px = Proxy::start(target, imp, imp, 42).await.unwrap();
    let url = format!("rist://127.0.0.1:{}?cname={}&buffer=1000&secret={}&aes-type=128", px.addr.port(), s["cname"].as_str().unwrap(), s["passphrase"].as_str().unwrap());
    let udp_in = free_port();
    let child = spawn(&tool("WAI_RISTSENDER"), &["-i", &format!("udp://@127.0.0.1:{udp_in}"), "-o", &url, "-p", "1", "-v", "2"]);
    assert!(contrib_live(&w, s["session"].as_str().unwrap()).await, "the reference sender's flow is the session's");
    let watcher = watch(&w).await;
    feed(&input, format!("127.0.0.1:{udp_in}").parse().unwrap()).await;
    tokio::time::sleep(Duration::from_millis(2_500)).await;
    drop(child);
    check(&w, &watcher, &input, Codec::Avc).await;
    let ps = linkrec::verify_series(&w.gateway.gateway.link_records()).unwrap();
    assert!(ps.iter().all(|p| p.protocol == 3 && p.cipher == 128));
    px.stop();
    w.stop().await;
}