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

//! VP8, VP9 and AV1 through WHIP: each a video-only publisher, its frames depacketised by
//! the WebRTC stack, packaged in its capability's framing (an IVF frame for VP8 and VP9,
//! SPEC §5's AV1 framing for AV1), received at the relay with every receipt verified,
//! and the same bytes the publisher sent.

mod common;

use common::*;
use joule_moq_ingest::frame::{self, Codec};
use std::time::Duration;
use str0m::media::MediaKind;
use tokio::time::timeout;

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn vp8_frames_arrive_whole_in_ivf() {
    timeout(Duration::from_secs(120), video_only(Codec::Vp8, (640, 360))).await.expect("the test finishes in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn vp9_frames_arrive_whole_in_ivf() {
    timeout(Duration::from_secs(120), video_only(Codec::Vp9, (1280, 720))).await.expect("the test finishes in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn av1_temporal_units_arrive_whole() {
    timeout(Duration::from_secs(120), video_only(Codec::Av1, (1280, 720))).await.expect("the test finishes in time");
}

fn body(seed: u64, n: usize) -> Vec<u8> {
    (0..n).map(|j| ((seed.wrapping_mul(31).wrapping_add(j as u64 * 7)) % 255 + 1) as u8).collect()
}

/// Bits, most significant first.
struct W(Vec<bool>);
impl W {
    fn put(&mut self, v: u64, n: usize) {
        for i in (0..n).rev() {
            self.0.push((v >> i) & 1 == 1);
        }
    }
    fn bytes(mut self) -> Vec<u8> {
        while !self.0.len().is_multiple_of(8) {
            self.0.push(false);
        }
        self.0.chunks(8).map(|c| c.iter().fold(0u8, |a, &b| (a << 1) | b as u8)).collect()
    }
}

fn leb128(mut v: usize) -> Vec<u8> {
    let mut o = Vec::new();
    loop {
        let b = (v & 0x7f) as u8;
        v >>= 7;
        if v == 0 {
            o.push(b);
            return o;
        }
        o.push(b | 0x80);
    }
}

fn obu(ty: u8, payload: &[u8]) -> Vec<u8> {
    let mut o = vec![(ty << 3) | 0x02];
    o.extend(leb128(payload.len()));
    o.extend_from_slice(payload);
    o
}

/// Frame `i` of `codec`: a keyframe when `key`, with a header that states `size`.
fn video(codec: Codec, i: u64, key: bool, size: (u32, u32)) -> Vec<u8> {
    let n = if key { 2500 } else { 600 + (i % 4) as usize * 50 };
    match codec {
        Codec::Vp8 => {
            // RFC 6386 §9.1: the frame tag (bit 0 clear on a keyframe, show_frame set), and
            // on a keyframe the start code and the size.
            let mut f = vec![if key { 0x10 } else { 0x11 }, 0x02, 0x00];
            if key {
                f.extend_from_slice(&[0x9d, 0x01, 0x2a]);
                f.extend_from_slice(&(size.0 as u16).to_le_bytes());
                f.extend_from_slice(&(size.1 as u16).to_le_bytes());
            }
            f.extend(body(i, n));
            f
        }
        Codec::Vp9 => {
            // The uncompressed header: frame_marker, profile 0, show_existing_frame 0,
            // frame_type, show_frame 1, error_resilient_mode 0; a keyframe's sync code,
            // color_config (BT.601, studio range) and frame_size.
            let mut w = W(Vec::new());
            w.put(2, 2);
            w.put(0, 2);
            w.put(0, 1);
            w.put(if key { 0 } else { 1 }, 1);
            w.put(1, 1);
            w.put(0, 1);
            if key {
                w.put(0x49_83_42, 24);
                w.put(1, 3);
                w.put(0, 1);
                w.put(size.0 as u64 - 1, 16);
                w.put(size.1 as u64 - 1, 16);
            }
            let mut f = w.bytes();
            f.extend(body(i, n));
            f
        }
        Codec::Av1 => {
            let mut tu = Vec::new();
            if key {
                // A sequence header: profile 0, one operating point at level 8 tier 0,
                // and the largest frame `size`.
                let mut w = W(Vec::new());
                w.put(0, 3);
                w.put(0, 1);
                w.put(0, 1);
                w.put(0, 1);
                w.put(0, 1);
                w.put(0, 5);
                w.put(0, 12);
                w.put(8, 5);
                w.put(0, 1);
                w.put(10, 4);
                w.put(10, 4);
                w.put(size.0 as u64 - 1, 11);
                w.put(size.1 as u64 - 1, 11);
                tu.extend(obu(1, &w.bytes()));
            }
            tu.extend(obu(6, &body(i, n)));
            tu
        }
        _ => unreachable!(),
    }
}

async fn video_only(codec: Codec, size: (u32, u32)) {
    let w = World::start(|_| {}).await;
    let auth = w.auth(&w.publisher, 0);
    let p = publish_with(&w, &auth, &[codec], false, true).await;
    let watcher = watch_kinds(&w, &[2]).await;
    let start = std::time::Instant::now();
    let mut sent = Vec::new();
    for i in 0..45u64 {
        tokio::time::sleep_until(tokio::time::Instant::from_std(start + Duration::from_micros(i * 33_333))).await;
        let f = video(codec, i, i % 15 == 0, size);
        p.write(MediaKind::Video, codec, f.clone(), 1_000 + i * 3000);
        sent.push(f);
    }
    assert!(received(&w, &p, sent.len()).await, "every frame reached the gateway");
    assert_eq!(delete(&w, &p.location, &auth).await.status, 200);
    assert!(until(STEP, || watcher.ended.lock().unwrap().len() == 1).await, "PUBLISH_DONE");
    settle(&watcher, Duration::from_millis(500)).await;
    let t = &watcher.tracks[0];
    let (frames, groups) = verified(&watcher, t, &w.gateway_key, codec);
    same(&frames, &sent, "every frame, byte for byte");
    assert_eq!(groups, 3, "one group per keyframe");
    // Each keyframe's payload states its size: the IVF header for VP8 and VP9, the AV1
    // framing's height and width.
    let env = envelopes(&watcher, t);
    for (i, e) in env.iter().enumerate().filter(|(i, _)| i % 15 == 0) {
        assert!(e.key, "object {i} is a keyframe");
        let (pw, ph) = match codec {
            Codec::Vp8 | Codec::Vp9 => (u16::from_le_bytes([e.payload[12], e.payload[13]]) as u32, u16::from_le_bytes([e.payload[14], e.payload[15]]) as u32),
            _ => (u32::from_le_bytes(e.payload[8..12].try_into().unwrap()), u32::from_le_bytes(e.payload[4..8].try_into().unwrap())),
        };
        assert_eq!((pw, ph), size, "the size {codec:?}'s framing states");
    }
    assert_eq!(env[0].manifest["model_requirement"]["capability"], codec.capability());
    assert!(frame::unframe(codec, &env[1].payload).is_some());
    p.close();
    w.stop().await;
}