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

//! A WAI-aware WHEP player states in its offer which WAI capabilities its sink dispatches, in
//! its order (`extensions/live-ingest.md` §9.5), and the gateway forwards it the rendition
//! the first of them names. An Enhanced RTMP publisher sends two video layers under the
//! room's `camera-720` profile, layer 0 in AV1 and layer 1 in H.264, and Opus; a viewer whose
//! grant reaches layer 1 POSTs offers to the owner's WHEP resource (layer 0, the floor):
//!
//! - H.264 first: the answer's video section is the configuration that names H.264, with its
//!   `a=acfg`, and the player gets layer 1's frames, byte for byte from where it joined;
//! - AV1 first: layer 0 in AV1, by its configuration;
//! - H.264 first, from a viewer whose grant stops at layer 0: AV1, the next capability named;
//! - VP8 alone, which the owner does not publish: the floor, no `a=acfg`;
//! - no capability attributes: answered as before, layer 0, no `a=acfg`, and a record with no
//!   `caps` member;
//! - a required option the gateway does not support (`caps_unsupported`): answered as without
//!   attributes, with the `a=csup` RFC 5939 §3.6.2 asks for, and the reason in the record;
//! - attributes that do not hold together: `422` with the reason.
//!
//! Each session's signed record keeps what it forwarded and why.

mod common;

use common::rtmp::*;
use common::*;
use joule_moq_ingest::frame::Codec;
use joule_moq_ingest::rtmp::amf0::Value;
use joule_moq_ingest::{sdp, sdp_caps};
use std::time::Duration;
use str0m::media::{Direction, MediaKind as RtcKind}; // allow-name: the str0m crate's own type name, imported once under a local alias
use tokio::net::TcpStream;
use tokio::time::timeout;

const GOP: u64 = 15;

/// The configuration number and payload type `offer`'s `a=pcfg` gives capability `n`'s
/// configuration `c`.
fn pcfg_pt(offer: &str, c: u32) -> String {
    let p = format!("a=pcfg:{c} m={c} pt={c}:");
    offer.split("\r\n").find_map(|l| l.strip_prefix(&p)).unwrap_or_else(|| panic!("the offer has {p}")).to_string()
}

/// The lines of media section `i` of `sdp` that start with `prefix`.
fn in_section<'a>(sdp: &'a str, i: usize, prefix: &str) -> Vec<&'a str> {
    sdp.split("\r\nm=").nth(i + 1).unwrap_or("").split("\r\n").filter(|l| l.starts_with(prefix)).collect()
}

/// The session level of `sdp`: the lines before its first `m=` line.
fn session_level(sdp: &str) -> &str {
    sdp.split("\r\nm=").next().unwrap()
}

/// An audio and video offer with `caps` written into it (none for `&[]`), and `edit`
/// applied.
async fn offering(caps: &[&str], edit: impl Fn(String) -> String) -> (Offering, String) {
    let o = offer(&[Codec::Avc, Codec::Av1, Codec::Vp8, Codec::Opus], Direction::RecvOnly, true, true).await;
    let text = if caps.is_empty() {
        o.offer.clone()
    } else {
        let (t, left) = sdp_caps::write(&o.offer, caps).expect("the offer takes capability attributes");
        assert!(left.is_empty(), "every capability is offered: {left:?}");
        t
    };
    let text = edit(text);
    (o, text)
}

/// POST `text` to the owner's WHEP resource until the owner is live.
async fn post(w: &World, va: &Auth, text: &str) -> Resp {
    let deadline = std::time::Instant::now() + STEP;
    loop {
        let r = va.send("POST", &w.whep_url(), &[("content-type", "application/sdp")], text.as_bytes()).await;
        if r.status != 409 {
            return r;
        }
        assert!(std::time::Instant::now() < deadline, "the owner's catalog arrives");
        tokio::time::sleep(Duration::from_millis(200)).await;
    }
}

/// A player: its offer's capabilities, whether a `creq` the gateway does not support is
/// added, who asks, the configuration its video section is answered by (none: no `a=acfg`),
/// and the codec it gets.
struct Case<'a> {
    name: &'a str,
    caps: &'a [&'a str],
    unsupported: bool,
    who: &'a Auth,
    config: Option<u32>,
    codec: Codec,
}

fn case<'a>(name: &'a str, caps: &'a [&'a str], unsupported: bool, who: &'a Auth, config: Option<u32>, codec: Codec) -> Case<'a> {
    Case { name, caps, unsupported, who, config, codec }
}

/// What the publisher sent of layer 1 (H.264, Annex B), in order.
struct Sender {
    c: Client<TcpStream>,
    hi: Vec<Vec<u8>>,
    next: u64,
}

impl Sender {
    /// `n` more frames of both layers and as many Opus packets, at 30 a second.
    async fn send(&mut self, n: u64) {
        let start = tokio::time::Instant::now();
        for k in 0..n {
            tokio::time::sleep_until(start + Duration::from_micros(k * 33_333)).await;
            let i = self.next;
            self.next += 1;
            let key = i.is_multiple_of(GOP);
            let ts = (i * 1000 / 30) as u32;
            let nals = avc_nals(i, key);
            let mut avc_body = vec![0, 0, 0];
            avc_body.extend(lengths(&nals));
            self.c.video(ts, ex_video_tracks(key, 1, &[(b"avc1", 0, avc_body), (b"av01", 1, av1_tu(i, key))])).await;
            let mut au = if key { vec![AVC_SPS.to_vec(), AVC_PPS.to_vec()] } else { Vec::new() };
            au.extend(nals);
            self.hi.push(annexb(&au));
            self.c.audio(ts, ex_audio(1, b"Opus", &opus_packet(i))).await;
        }
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_player_gets_the_rendition_its_sink_dispatches_first() {
    timeout(Duration::from_secs(150), run()).await.expect("the test finishes in time");
}

async fn run() {
    // A viewer may take up to layer 1; the publisher sends layers 0 and 1 by camera-720.
    let w = World::start_room(
        |d| d.roles[1].caps.max_layer = 1,
        |_, _, _| {},
        |c| {
            c.rtmp.listen = Some("127.0.0.1:0".parse().unwrap());
            c.timeline.max_skew_ms = 600_000;
        },
    )
    .await;
    let mut g = w.grant(&w.publisher, 0);
    g.profiles = vec![1];
    let auth = Auth::new(g.sign(&w.authority), &w.publisher);
    let key = stream_key(&w, &auth).await;
    let mut c = Client::start(TcpStream::connect(w.gateway.rtmp_addr.expect("an RTMP listener")).await.unwrap()).await;
    c.connect(0x0f).await;
    let s = c.publish(&key).await;
    assert_eq!(s.code, "NetStream.Publish.Start", "{}", s.description);
    let info = |w: f64, h: f64| Value::obj(&[("width", Value::Number(w)), ("height", Value::Number(h)), ("framerate", Value::Number(30.0))]);
    c.metadata(vec![
        ("width".into(), Value::Number(1280.0)),
        ("height".into(), Value::Number(720.0)),
        ("framerate".into(), Value::Number(30.0)),
        ("videoTrackIdInfoMap".into(), Value::obj(&[("1", info(640.0, 360.0))])),
    ])
    .await;
    c.video(0, ex_video_tracks(true, 0, &[(b"avc1", 0, avc_record()), (b"av01", 1, av1_record(640, 360))])).await;
    c.audio(0, ex_audio(0, b"Opus", &opus_head())).await;
    let mut tx = Sender { c, hi: Vec::new(), next: 0 };
    tx.send(30).await;

    let viewer = ed25519_dalek::SigningKey::from_bytes(&[0x61; 32]);
    let va = w.auth(&viewer, 1);
    // A second viewer, whose grant stops at layer 0.
    let viewer0 = ed25519_dalek::SigningKey::from_bytes(&[0x62; 32]);
    let mut g0 = w.grant(&viewer0, 1);
    g0.caps.max_layer = 0;
    let va0 = Auth::new(g0.sign(&w.authority), &viewer0);

    // Attributes that do not hold together: 422, with the reason, before anything is held.
    for (from, to, reason) in [("a=pcfg:2 ", "a=pcfg: 2 ", "caps_malformed"), ("a=csup:med-v0\r\n", "", "caps_option_missing")] {
        let (_, text) = offering(&["wai.video.avc", "wai.video.av1", "wai.audio.opus.packet"], |t| t.replacen(from, to, 1)).await;
        let r = post(&w, &va, &text).await;
        assert_eq!(r.status, 422, "{}", r.text());
        let body: serde_json::Value = serde_json::from_slice(&r.body).unwrap();
        assert_eq!(body["reason"], reason);
    }

    // The players. The writer numbers Opus's configuration 1 (the audio section comes first)
    // and the video capabilities from 2, in the sink's order.
    let order_avc = ["wai.video.avc", "wai.video.av1", "wai.audio.opus.packet"];
    let order_av1 = ["wai.video.av1", "wai.video.avc", "wai.audio.opus.packet"];
    let creq = |t: String| t.replacen("a=csup:med-v0\r\n", "a=csup:med-v0\r\na=creq:med-v0,x-later-v1\r\n", 1);
    let cases = [
        case("avc", &order_avc, false, &va, Some(2), Codec::Avc),
        case("av1", &order_av1, false, &va, Some(2), Codec::Av1),
        case("bounded", &order_avc, false, &va0, Some(3), Codec::Av1),
        case("vp8", &["wai.video.vp8"], false, &va, None, Codec::Av1),
        case("none", &[], false, &va, None, Codec::Av1),
        case("unsupported", &order_avc, true, &va, None, Codec::Av1),
    ];
    let mut players = Vec::new();
    for Case { name, caps, unsupported, who, config, codec } in cases {
        let (o, text) = offering(caps, |t| if unsupported { creq(t) } else { t }).await;
        let r = post(&w, who, &text).await;
        assert_eq!(r.status, 201, "{name}: {}", r.text());
        let answer = r.text();
        if let Some(c) = config {
            // The video section answers configuration `c` alone, with its a=acfg; the audio
            // section configuration 1.
            let pt = pcfg_pt(&text, c);
            assert_eq!(sdp::formats(&answer, 1).unwrap(), [pt.as_str()], "{name}: the configuration's payload type alone");
            assert_eq!(in_section(&answer, 1, "a=acfg:"), [format!("a=acfg:{c} m={c} pt={c}:{pt}")], "{name}");
            let enc = if codec == Codec::Avc { "H264/90000" } else { "AV1/90000" };
            assert_eq!(in_section(&answer, 1, &format!("a=rtpmap:{pt} ")), [format!("a=rtpmap:{pt} {enc}")], "{name}");
            let apt = pcfg_pt(&text, 1);
            assert_eq!(in_section(&answer, 0, "a=acfg:"), [format!("a=acfg:1 m=1 pt=1:{apt}")], "{name}");
            assert!(!answer.contains("a=csup:"), "{name}");
        } else if unsupported {
            assert!(!answer.contains("a=acfg:"), "{name}: not negotiated");
            assert_eq!(session_level(&answer).split("\r\n").filter(|l| l.starts_with("a=csup:")).collect::<Vec<_>>(), ["a=csup:med-v0"], "{name}: a=csup at the creq's level");
        } else {
            assert!(!answer.contains("a=acfg:") && !answer.contains("a=csup:") && !answer.contains("a=pcfg:"), "{name}: answered as §9.1 has it");
        }
        let id = r.header("location").unwrap().rsplit('/').next().unwrap().to_string();
        players.push((name, id, codec, o.start(&r)));
    }
    for (_, _, _, p) in &players {
        p.connected().await;
    }

    // Three more seconds, then a last keyframe closes the group before it.
    tx.send(90).await;
    tokio::time::sleep(Duration::from_millis(300)).await;
    tx.send(1).await;
    let closed = tx.hi.len() - 1;
    for (name, _, codec, p) in &players {
        assert!(until(STEP, || p.media(RtcKind::Video).len() >= 2 * GOP as usize && !p.media(RtcKind::Audio).is_empty()).await, "{name}: frames arrive");
        assert_eq!(p.codecs(RtcKind::Video), [Some(*codec)], "{name}: the rendition forwarded");
        assert_eq!(p.codecs(RtcKind::Audio), [Some(Codec::Opus)], "{name}");
    }
    // The H.264 player has layer 1's frames, byte for byte, from a keyframe on.
    let (_, _, _, pa) = &players[0];
    assert!(until(STEP, || pa.media(RtcKind::Video).last() == tx.hi.get(closed - 1)).await, "every frame up to the last closed group");
    let got = pa.media(RtcKind::Video);
    let j = run_of(&tx.hi, &got);
    assert_eq!(j % GOP as usize, 0, "the player joins at a keyframe");

    // The publisher leaves; each WHEP session ends with the tracks and signs its record.
    tx.c.delete_stream().await;
    assert!(until(STEP, || w.gateway.gateway.records().iter().filter(|r| r.record.protocol == "whep").count() == players.len()).await, "every WHEP session's record");
    let records = w.gateway.gateway.records();
    for (name, id, _, p) in &players {
        let rec = records.iter().find(|r| r.record.session == *id).unwrap_or_else(|| panic!("{name}: a record"));
        rec.verify().unwrap();
        let video_name: Vec<&str> = rec.record.tracks.iter().filter(|t| t.namespace.ends_with("-video")).map(|t| t.name.as_str()).collect();
        let caps = rec.record.detail.get("caps");
        match *name {
            "avc" => {
                assert_eq!(video_name, ["1", "1"], "{name}: layer 1, in and out");
                let pt = pcfg_pt_of(p);
                assert_eq!(caps.unwrap()["video"], serde_json::json!({ "acfg": format!("a=acfg:2 m=2 pt=2:{pt}"), "advertised": ["wai.video.avc", "wai.video.av1"], "capability": "wai.video.avc", "layer": 1 }));
                assert_eq!(caps.unwrap()["audio"]["capability"], "wai.audio.opus.packet");
                assert!(rec.record.subscriber_us[1] > 0 && rec.record.subscriber_us[0] == 0, "{name}: its time under layer 1");
            }
            "av1" | "bounded" => {
                assert_eq!(video_name, ["0", "0"], "{name}");
                assert_eq!(caps.unwrap()["video"]["capability"], "wai.video.av1", "{name}");
                assert_eq!(caps.unwrap()["video"]["layer"], 0, "{name}");
            }
            "vp8" => {
                assert_eq!(video_name, ["0", "0"], "{name}");
                assert_eq!(caps.unwrap()["video"], serde_json::json!({ "advertised": ["wai.video.vp8"], "floor": true, "layer": 0 }));
                assert_eq!(caps.unwrap()["audio"], serde_json::json!({ "advertised": [], "floor": true, "layer": 0 }));
            }
            "none" => {
                assert_eq!(video_name, ["0", "0"], "{name}");
                assert!(caps.is_none(), "{name}: the record is as it was");
                assert!(rec.record.subscriber_us[0] > 0);
            }
            _ => {
                assert_eq!(video_name, ["0", "0"], "{name}");
                assert_eq!(caps, Some(&serde_json::json!({ "not_negotiated": "caps_unsupported" })), "{name}");
            }
        }
        p.close();
    }
    tx.c.shutdown().await;
    w.stop().await;
}

/// The payload type the gateway answered the H.264 player's video section with.
fn pcfg_pt_of(p: &Peer) -> String {
    sdp::formats(&p.answer, 1).unwrap().remove(0)
}