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

//! A protected room through a relay (WAI moq-streaming-format §4.9): the video and the audio
//! track of `wai/live-conformance/protected-cmsf/`, and the catalogs that state both, one for
//! each key period, published with the gateway's own publisher, each group's receipts signed
//! over the protected bytes, through a relay whose door checks every object against its
//! receipt.
//!
//! - **The catalogs.** Each key period's catalog is a group of the owner's catalog track,
//!   signed as the room's other objects are (the door holds the catalog track to its receipts
//!   too), published before the period's first group, and the subscriber holds it before the
//!   publisher sends that group. A subscriber reads the three catalogs as the corpus holds
//!   them (each track's namespace added, as a room's catalog states it), and from what it
//!   holds alone finds each group's key period: the last catalog whose `from_group` is at
//!   most the group, whose header for the track names the KID its entry states. The two
//!   tracks never share a KID, and no two periods of a track do.
//! - **The objects.** A subscriber reads every object of both tracks byte for byte as the
//!   corpus holds it, verifies every receipt under the publisher's key, and holds no clear
//!   chunk: a protected object carries nothing of this section's on the wire, so the relay
//!   forwards it as any object, and the receipts are over what was published.
//!
//! That each object decrypts, under the key of its group's key period, to its clear chunk
//! is wai-pack's `tests/cmsf.rs` and `live-conformance/protected_cmsf_verify.py`, over the
//! same catalogs and the same bytes. This crate packages no CMAF: the tracks are the corpus's.

mod common;

use base64::Engine;
use common::*;
use joule_moq::carriage::ReceiptBundle;
use joule_moq::receipt::{ObjectReceipt, ReceiptVerifier};
use joule_moq::{GroupId, ObjectId, TrackName};
use joule_moq_ingest::catalog;
use joule_moq_ingest::package::{rfc3339, sign_group, Emit, Signer};
use joule_moq_ingest::relay::{Published, Publisher};
use serde_json::Value;
use sha2::Digest;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::timeout;

const CATALOG_ALIAS: u64 = 98;
const CATALOG_RECEIPTS_ALIAS: u64 = 99;

fn wai_dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../../../wai")
}

/// One track of the corpus: its name, its aliases, its media kind (2 video, 1 audio), and
/// its objects and clear chunks.
struct Track {
    name: String,
    kind: u8,
    media: u64,
    receipts: u64,
    objects: Vec<(u64, u64, Vec<u8>)>,
    clear: Vec<Vec<u8>>,
}

fn track(doc: &Value, case: &str, clear: &Value, name: &str, kind: u8, media: u64) -> Track {
    let case = doc["cases"].as_array().unwrap().iter().find(|c| c["name"] == case).expect("the case");
    assert_eq!(case["outcome"], "accept");
    let objects = case["objects"].as_array().unwrap().iter().map(|o| (o["group"].as_u64().unwrap(), o["object"].as_u64().unwrap(), hex::decode(o["hex"].as_str().unwrap()).unwrap())).collect();
    let clear = clear["objects"]
        .as_array()
        .unwrap()
        .iter()
        .map(|o| {
            let b = std::fs::read(wai_dir().join(o["file"].as_str().unwrap())).expect("a clear chunk");
            // A chunk that is part of its file.
            let b = match (o["at"].as_u64(), o["n"].as_u64()) {
                (Some(at), Some(n)) => b[at as usize..(at + n) as usize].to_vec(),
                _ => b,
            };
            assert_eq!(hex::encode(sha2::Sha256::digest(&b)), o["sha256"].as_str().unwrap(), "the clear chunk the corpus states");
            b
        })
        .collect();
    Track { name: name.into(), kind, media, receipts: media + 1, objects, clear }
}

/// The KID a header's `tenc` names: the sixteen bytes after its version, flags and four more.
fn header_kid(header: &[u8]) -> String {
    let at = header.windows(4).position(|w| w == b"tenc").expect("a tenc") + 4 + 8;
    let h = hex::encode(&header[at..at + 16]);
    format!("{}-{}-{}-{}-{}", &h[..8], &h[8..12], &h[12..16], &h[16..20], &h[20..])
}

/// What a catalog states of a track: its key period (index, first group) and the KID its
/// entry names, which is the one its header names.
fn stated(catalog: &Value, name: &str) -> (u64, u64, String) {
    let t = catalog["tracks"].as_array().unwrap().iter().find(|t| t["name"] == name).expect("the track");
    let refs = t["contentProtectionRefIDs"].as_array().expect("the track is protected");
    assert_eq!(refs.len(), 1);
    let entry = catalog["contentProtections"].as_array().unwrap().iter().find(|e| e["refID"] == refs[0]).expect("the entry it names");
    let kid = entry["defaultKID"][0].as_str().expect("a KID").to_string();
    let init = catalog["initDataList"].as_array().unwrap().iter().find(|e| e["id"] == t["initRef"]).expect("the header it names");
    let header = base64::engine::general_purpose::STANDARD.decode(init["data"].as_str().unwrap()).expect("Base64");
    assert_eq!(header_kid(&header), kid, "{name}: the header names the KID the entry states");
    (t["wai_key_period"]["index"].as_u64().expect("an index"), t["wai_key_period"]["from_group"].as_u64().expect("a first group"), kid)
}

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

async fn run() {
    let doc: Value = serde_json::from_slice(&std::fs::read(wai_dir().join("live-conformance/protected-cmsf/vectors.json")).expect("the corpus")).expect("JSON");
    let tracks = [
        track(&doc, "both-tracks-video", &doc["clear"], doc["track"].as_str().unwrap(), 2, 100),
        track(&doc, "both-tracks-audio", &doc["audio"]["clear"], doc["audio"]["track"].as_str().unwrap(), 1, 102),
    ];
    assert_eq!((tracks[0].objects.len(), tracks[1].objects.len()), (5, 6));
    // The catalogs of the two cases are one set: a catalog for each key period, of both tracks.
    let case = |name: &str| doc["cases"].as_array().unwrap().iter().find(|c| c["name"] == name).unwrap()["catalogs"].clone();
    let corpus = case("both-tracks-video");
    assert_eq!(corpus, case("both-tracks-audio"));
    let corpus = corpus.as_array().unwrap().clone();
    assert_eq!(corpus.len(), 3, "three key periods");

    let w = World::start(|_| {}).await;
    let ns = |kind: u8| catalog::media_namespace(&w.room, &w.owner(), kind);
    let cat_ns = catalog::catalog_namespace(&w.room, &w.owner());
    // A room's catalog states each track's namespace: the corpus's, with it.
    let published: Vec<Vec<u8>> = corpus
        .iter()
        .map(|c| {
            let mut c = c.clone();
            for entry in c["tracks"].as_array_mut().unwrap() {
                let of = tracks.iter().find(|t| entry["name"] == t.name.as_str() || entry["name"] == format!("{}_receipts", t.name).as_str()).expect("a track of the room");
                entry["namespace"] = catalog::render(&ns(of.kind)).into();
            }
            serde_json::to_vec(&c).unwrap()
        })
        .collect();
    for t in &tracks {
        let entry = corpus[0]["tracks"].as_array().unwrap().iter().find(|e| e["name"] == t.name.as_str()).expect("the track");
        assert_eq!((entry["packaging"].as_str(), entry["wai_receipts_track"].as_str().map(TrackName::from_str)), (Some("cmaf"), Some(catalog::receipts_name(&t.name))));
    }

    let (client, endpoint) = relay_client(&w).await;
    let mut publishing = vec![
        Published { ns: cat_ns.clone(), name: TrackName::from_str(catalog::CATALOG_NAME), alias: CATALOG_ALIAS, priority: 1 },
        Published { ns: cat_ns.clone(), name: catalog::receipts_name(catalog::CATALOG_NAME), alias: CATALOG_RECEIPTS_ALIAS, priority: 1 },
    ];
    for t in &tracks {
        publishing.push(Published { ns: ns(t.kind), name: TrackName::from_str(&t.name), alias: t.media, priority: 1 });
        publishing.push(Published { ns: ns(t.kind), name: catalog::receipts_name(&t.name), alias: t.receipts, priority: 1 });
    }
    let publisher = Publisher::start(client, endpoint, publishing, 1 << 24).await.expect("PUBLISH the catalog and both tracks");
    // The two media tracks, then the catalog track, each with its receipts track.
    let watcher = watch_tracks(&w, tracks.iter().map(|t| (t.kind, ns(t.kind), t.name.clone())).chain([(0, cat_ns, catalog::CATALOG_NAME.to_string())]).collect()).await;
    let (cat_alias, cat_receipts) = (watcher.tracks[2].media_alias, watcher.tracks[2].receipts_alias);

    let signer = Signer { key: Arc::new(w.gateway_key.clone()), signer_id: "protected-cmsf-test".into(), origin_label: "protected-cmsf-test".into() };
    let mut parents = [None, None];
    let mut catalog_parent = None;
    let t0 = now_us();
    let periods: Vec<(u64, u64)> = corpus.iter().map(|c| (stated(c, &tracks[0].name).1, stated(c, &tracks[1].name).1)).collect();
    let mut sent = 0;
    for (k, bytes) in published.iter().enumerate() {
        // The period's catalog, a group of the catalog track; the subscriber holds it before
        // the publisher sends the period's first group.
        let (bundle, link) = sign_group(&signer, CATALOG_ALIAS, k as u64, &[ObjectReceipt::hash_payload(bytes)], &rfc3339(t0), &rfc3339(t0 + 1), catalog_parent);
        catalog_parent = Some(link);
        assert!(publisher.send(Emit::Single { alias: CATALOG_RECEIPTS_ALIAS, group: k as u64, priority: 1, bytes: bundle }));
        assert!(publisher.send(Emit::Single { alias: CATALOG_ALIAS, group: k as u64, priority: 1, bytes: bytes.clone() }));
        assert!(until(STEP, || watcher.objects.lock().unwrap().contains_key(&(cat_alias, k as u64, 0))).await, "the catalog of key period {k} arrived");
        for (i, t) in tracks.iter().enumerate() {
            let (from, next) = (if i == 0 { periods[k].0 } else { periods[k].1 }, periods.get(k + 1).map(|p| if i == 0 { p.0 } else { p.1 }));
            assert!(!watcher.objects.lock().unwrap().keys().any(|(alias, g, _)| *alias == watcher.tracks[i].media_alias && *g >= from), "{}: no object of key period {k} before its catalog", t.name);
            for (group, object, bytes) in t.objects.iter().filter(|o| o.0 >= from && next.is_none_or(|n| o.0 < n)) {
                assert_eq!(*object, 0, "one chunk a group in this corpus");
                let (bundle, link) = sign_group(&signer, t.media, *group, &[ObjectReceipt::hash_payload(bytes)], &rfc3339(t0), &rfc3339(t0 + 1), parents[i]);
                parents[i] = Some(link);
                assert!(publisher.send(Emit::Object { alias: t.media, group: *group, object: 0, first: true, priority: 1, bytes: bytes.clone() }));
                assert!(publisher.send(Emit::EndGroup { alias: t.media, group: *group }));
                assert!(publisher.send(Emit::Single { alias: t.receipts, group: *group, priority: 1, bytes: bundle }));
                sent += 2;
            }
        }
        // The period's objects are with the subscriber before the next period's catalog is sent.
        assert!(until(STEP, || watcher.objects.lock().unwrap().len() == sent + 2 * (k + 1)).await, "every object and every receipts bundle of key period {k} arrived");
    }
    assert_eq!(sent, 2 * (tracks[0].objects.len() + tracks[1].objects.len()));
    settle(&watcher, Duration::from_millis(300)).await;
    let got = watcher.objects.lock().unwrap().clone();
    assert_eq!(got.len(), sent + 2 * published.len(), "and nothing else");
    let verifier = ReceiptVerifier::with_anchors(vec![w.gateway_key.verifying_key()]);

    // The catalogs, as the subscriber holds them: the published bytes, which are the
    // corpus's with each track's namespace.
    let mut held = Vec::new();
    for (k, bytes) in published.iter().enumerate() {
        let c = got.get(&(cat_alias, k as u64, 0)).expect("the catalog");
        assert_eq!(c, bytes, "the catalog of key period {k}, byte for byte");
        let bundle = ReceiptBundle::from_payload(got.get(&(cat_receipts, k as u64, 0)).expect("the catalog's receipts")).expect("a receipts bundle");
        let verified = verifier.verify_group(&bundle.group).expect("the catalog's group receipt verifies");
        verifier.verify_payload(c, &bundle.objects[0], &verified).expect("the receipt is over the catalog");
        let mut read: Value = serde_json::from_slice(c).expect("JSON");
        held.push(read.clone());
        for entry in read["tracks"].as_array_mut().unwrap() {
            assert!(entry.as_object_mut().unwrap().remove("namespace").is_some());
        }
        assert_eq!(read, corpus[k], "the corpus's catalog of key period {k}");
    }
    let mut kids = std::collections::BTreeSet::new();
    for (i, t) in tracks.iter().enumerate() {
        let of: Vec<(u64, u64, String)> = held.iter().map(|c| stated(c, &t.name)).collect();
        assert!(of.windows(2).all(|w| w[0].0 < w[1].0 && w[0].1 < w[1].1), "{}: each catalog's key period follows the one before", t.name);
        for (_, _, kid) in &of {
            assert!(kids.insert(kid.clone()), "{}: a KID of its own for each key period, and none of the other track's", t.name);
        }
        // From what the subscriber holds alone, each group's key period: the last catalog
        // whose first group is at most it.
        let wt = &watcher.tracks[i];
        let mut by_period = [0usize; 3];
        for (j, (group, _, bytes)) in t.objects.iter().enumerate() {
            let period = of.iter().rev().find(|p| p.1 <= *group).unwrap_or_else(|| panic!("{}: group {group} is of a key period a catalog states", t.name));
            by_period[period.0 as usize] += 1;
            let payload = got.get(&(wt.media_alias, *group, 0)).unwrap_or_else(|| panic!("{}: group {group} arrived", t.name));
            assert_eq!(payload, bytes, "{}: group {group}: the protected object, byte for byte", t.name);
            let bundle = ReceiptBundle::from_payload(got.get(&(wt.receipts_alias, *group, 0)).expect("its receipts")).expect("a receipts bundle");
            let verified = verifier.verify_group(&bundle.group).expect("the group receipt verifies");
            let r = &bundle.objects[0];
            r.check_moqt(bundle.group.track_alias, GroupId(*group), ObjectId(0)).expect("the receipt names its object");
            verifier.verify_payload(payload, r, &verified).expect("the receipt is over the protected bytes");
            // What was published is not the clear chunk, and no receipt is over that.
            assert!(verifier.verify_payload(&t.clear[j], r, &verified).is_err(), "{}: group {group}: no receipt is over the clear chunk", t.name);
            assert!(payload.len() > t.clear[j].len() && *payload != t.clear[j]);
        }
        assert!(by_period.iter().all(|n| *n >= 1) && by_period.iter().sum::<usize>() == t.objects.len(), "{}: objects of every key period: {by_period:?}", t.name);
    }
    assert_eq!(kids.len(), 6, "three key periods of two tracks");
    publisher.end(Duration::from_millis(100)).await;
    w.stop().await;
}