open-standards/wai/wai-pack/src/pack.rs

//! The packer ([packaging] §4–§9): a policy and its sources in, a package out, every
//! file a pure function of the source bytes, the effective policy, the supplied floor, the
//! signing key and the packager's identity, except the energy figure and the measurement
//! claim of a metered pack.
//!
//! This revision packs one video component with a `wai.video.int_motion` base (its step
//! chosen per piece by [`crate::encode::choose_q`]) and a supplied AV1 floor. What it
//! does not pack yet it refuses by name rather than leaving out: another medium
//! (`pack-media-unsupported`), refinement rungs (`pack-refinement-unsupported`), a rule's
//! ladder subset (`pack-rule-ladder-unsupported`), a floor to encode
//! (`pack-floor-encode-unsupported`), a rung of a capability it does not encode
//! (`pack-capability-unsupported:<cap>`). Its own failures are refused by name too
//! (`pack-encode`, `pack-piece-length`, `pack-catalog`, `pack-meter`,
//! `pack-measurement-claim`, `pack-listing`; packaging §11.4), and a package is written
//! only into a new or empty directory ([`out_dir_ready`], `pack-out-not-empty`).
//!
//! The package:
//!
//! ```text
//! policy.toml  policy.json  plan-<c>.txt  index-<c>.wsix  package.receipt.json
//! objects/<content hash>.wai …
//! floor/<c>/init.mp4  floor/<c>/<sha256>.m4s …            (CMAF; never in an index)
//! hls/classic/…  hls/<class>/…  dash/…  moq/catalog.json  stages/<class>.json
//! claims/<receipt hash>.measurement.json                  (a metered pack)
//! LISTING                                                 (SPEC §3.1, of every other file)
//! ```
//!
//! [packaging]: ../../../extensions/packaging.md

use crate::catalog::{self, CompView, FloorView};
use crate::classify::{self, PieceKind};
use crate::encode::{choose_q, MotionParams};
use crate::floor::{self, Floor};
use crate::objects;
use crate::plan::{class_plan, plan_piece};
use crate::policy::{Media, Overrides, Policy, Role as RungRole, Rung};
use crate::scenecut::{SegParams, Segmenter};
use ed25519_dalek::SigningKey;
use std::collections::BTreeMap;
use wai::package::{PackageReceiptV2, Parent};
use wai::quantum_ops::GrantRef;
use wai::stage_index::{self as wsix, ClassRec, EntryRec, Header, MediaClass, PieceRec, Role, StageIndex};

/// The packager's identity, `builtin:<crate>@<version>`.
pub fn packager_id() -> String {
    format!("builtin:{}@{}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"))
}

/// Why nothing was packed: a reason of §12 and what failed. Nothing is written.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PackError {
    pub code: String,
    pub detail: String,
}

impl PackError {
    fn new(code: impl Into<String>, detail: impl Into<String>) -> PackError {
        PackError { code: code.into(), detail: detail.into() }
    }
}

impl std::fmt::Display for PackError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}: {}", self.code, self.detail)
    }
}

impl std::error::Error for PackError {}

/// What a pack is made with beside the policy and its sources.
pub struct PackOpts {
    pub signer: SigningKey,
    pub signer_id: String,
    /// The supplied floor's files, by listing path.
    pub floor: BTreeMap<String, Vec<u8>>,
    /// The override a command line applies (§3.2).
    pub overrides: Overrides,
    /// The meter that brackets the pack, or none (the unmetered form).
    pub meter: Option<Box<dyn wai::energy_meter::Meter>>,
    /// What a meter's report names as its power reader.
    pub power_reader: String,
    /// The `issued_at` of a measurement claim (RFC 3339, UTC).
    pub issued_at: String,
    pub grant: GrantRef,
}

impl PackOpts {
    /// Unmetered options with the grant `wai.video.package` and no ceilings.
    pub fn new(signer: SigningKey, signer_id: impl Into<String>, floor: BTreeMap<String, Vec<u8>>) -> PackOpts {
        PackOpts {
            signer,
            signer_id: signer_id.into(),
            floor,
            overrides: Overrides::default(),
            meter: None,
            power_reader: String::new(),
            issued_at: String::new(),
            grant: GrantRef::unbounded("wai.video.package"),
        }
    }
}

/// A package: its files by path, and its receipt.
#[derive(Debug, Clone)]
pub struct PackOutput {
    pub files: BTreeMap<String, Vec<u8>>,
    pub receipt: PackageReceiptV2,
}

impl PackOutput {
    /// Write every file below `dir`.
    pub fn write(&self, dir: &std::path::Path) -> std::io::Result<()> {
        for (p, b) in &self.files {
            let path = dir.join(p);
            if let Some(parent) = path.parent() {
                std::fs::create_dir_all(parent)?;
            }
            std::fs::write(path, b)?;
        }
        Ok(())
    }

    /// The SHA-256 of `LISTING`: the pack digest.
    pub fn pack_digest(&self) -> String {
        crate::sha256_hex(&self.files["LISTING"])
    }
}

/// Whether `dir` may take a package: a directory that does not exist yet, or one that
/// holds nothing (`pack-out-not-empty` otherwise, §11.4). A package written among another
/// one's files would leave the files it does not write behind (a measurement claim, an
/// object no entry names), and would not verify.
pub fn out_dir_ready(dir: &std::path::Path) -> Result<(), PackError> {
    let refuse = |why: String| Err(PackError::new("pack-out-not-empty", format!("{}: {why}", dir.display())));
    match std::fs::read_dir(dir) {
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
        Err(e) => refuse(e.to_string()),
        Ok(mut entries) => match entries.next() {
            None => Ok(()),
            Some(_) => refuse("not empty; a package is written only into a new or empty directory".into()),
        },
    }
}

/// Whether `id` can name a file: one listing-path segment (SPEC §3.1).
pub fn id_is_path_segment(id: &str) -> bool {
    let b = id.as_bytes();
    let first = |c: u8| c.is_ascii_alphanumeric() || c == b'_' || c == b'-';
    !b.is_empty() && first(b[0]) && b.iter().all(|&c| first(c) || c == b'.') && b[b.len() - 1] != b'.'
}

/// The SHA-256 clip digest of frames (SPEC §7, *Digests*): the SHA-256 of their RGB8
/// SHA-256 digests, concatenated in order.
pub fn clip_digest<F: AsRef<[u8]>>(frames: &[F]) -> [u8; 32] {
    use sha2::Digest as _;
    let mut h = sha2::Sha256::new();
    for f in frames {
        h.update(wai::output_tree::out_digest(f.as_ref()));
    }
    h.finalize().into()
}

/// An object's manifest and payload lengths. Where WAI checks its payload's format, they
/// are its counted bytes (staged-measurement §3, `measure_framing::count_object`), which
/// also checks the payload parses as its capability's pinned format; a floor's payload,
/// whose format WAI does not check, is read from the envelope's length fields.
pub fn object_lengths(env: &[u8]) -> Result<(u32, u64), String> {
    let cap = wai::container::Wai::from_bytes(env).map_err(|e| e.to_string())?.manifest.model_requirement.capability;
    if wai::measure_framing::has_format_check(&cap) {
        let c = wai::measure_framing::count_object(env).map_err(|e| e.to_string())?;
        return Ok((c.manifest as u32, c.payload));
    }
    let man = u32::from_le_bytes(env[4..8].try_into().expect("an envelope"));
    Ok((man, env.len() as u64 - 12 - man as u64))
}

fn rungs_of<'a>(p: &'a Policy, comp: &str, role: RungRole) -> Vec<&'a Rung> {
    p.ladder.iter().filter(|r| r.component == comp && r.role == role).collect()
}

/// An entry describing `env`.
#[allow(clippy::too_many_arguments)]
fn entry(env: &[u8], piece: u64, capability: u16, role: Role, prefix_tier: u8, prefix_digest: [u8; 32], sse: u64, flags: u8) -> Result<EntryRec, PackError> {
    let (man, payload_len) = object_lengths(env).map_err(|e| PackError::new("pack-encode", format!("an object the packer wrote is not countable: {e}")))?;
    Ok(EntryRec {
        content_hash: objects::object_hash(env),
        byte_len: env.len() as u64,
        piece,
        refines: wsix::NO_REFINES,
        capability,
        role,
        layer: 0,
        prefix_tier,
        flags,
        deadline_ms: 0,
        prefix_digest,
        sse,
        manifest_len: man,
        payload_len,
    })
}

/// A component's CMAF files: the initialization segment's path and bytes, and each
/// piece's fragment's.
pub type CmafFiles = (String, Vec<u8>, Vec<(String, Vec<u8>)>);

/// The floor's CMAF files of a component: the initialization segment and each piece's
/// fragment, by path. `floor` is the checked floor, `starts` the piece starts.
pub fn cmaf_files(comp: &str, floor: &Floor, starts: &[u64], tb_den: u32, timescale: u32) -> Result<CmafFiles, String> {
    let init = crate::cmaf::write_init_av1(&floor.seq, &floor.seq_obu, floor.seq.max_frame_width, floor.seq.max_frame_height, timescale);
    let mut frags = Vec::with_capacity(floor.pieces.len());
    for (p, piece) in floor.pieces.iter().enumerate() {
        let mut samples = Vec::with_capacity(piece.units.len());
        for (u, k) in piece.units.iter().zip(&piece.key) {
            let tus = crate::av1::temporal_units(u).map_err(|e| e.to_string())?;
            let [tu] = tus.as_slice() else { return Err("a floor unit is not one temporal unit".into()) };
            samples.push(crate::cmaf::Sample { data: tu.sample(), duration: tb_den, sync: *k });
        }
        let tfdt = starts[p].checked_mul(tb_den as u64).ok_or("a fragment's decode time past 2^64 - 1")?;
        let f = crate::cmaf::write_fragment(p as u32 + 1, tfdt, &samples);
        frags.push((format!("floor/{comp}/{}.m4s", crate::sha256_hex(&f)), f));
    }
    Ok((format!("floor/{comp}/init.mp4"), init, frags))
}

/// The keyframe plan's file: each piece start, one decimal per line.
pub fn plan_file(starts: &[u64]) -> String {
    starts.iter().map(|s| format!("{s}\n")).collect()
}

/// Pack `policy_bytes` over `sources` (each component's source bytes, by id).
pub fn pack(policy_bytes: &[u8], sources: &BTreeMap<String, Vec<u8>>, opts: &mut PackOpts) -> Result<PackOutput, PackError> {
    let policy = Policy::parse_bytes(policy_bytes).map_err(|e| PackError::new(e.code(), e.to_string()))?.effective(&opts.overrides);
    // The policy's own refusals, in the order a package's verifier applies them (§11.3).
    if let Some(r) = crate::verify::policy_refusal(&policy) {
        return Err(PackError::new(r.code, r.detail));
    }
    let [comp] = policy.components.as_slice() else {
        return Err(PackError::new("pack-media-unsupported", "this revision packs one component"));
    };
    if comp.media != Media::Video {
        return Err(PackError::new("pack-media-unsupported", format!("component {} is {}", comp.id, comp.media.label())));
    }
    if !rungs_of(&policy, &comp.id, RungRole::Refinement).is_empty() {
        return Err(PackError::new("pack-refinement-unsupported", "refinement rungs are packed by a later revision"));
    }
    if policy.rules.iter().any(|r| r.then.ladder.is_some()) {
        return Err(PackError::new("pack-rule-ladder-unsupported", "a rule's ladder subset is packed by a later revision"));
    }
    let bases = rungs_of(&policy, &comp.id, RungRole::Base);
    let [base] = bases.as_slice() else {
        let extra = bases.get(1).map_or("", |r| r.capability.as_str());
        return Err(PackError::new(format!("pack-capability-unsupported:{extra}"), "this revision packs one base rung"));
    };
    let mp = MotionParams::of(base).ok_or_else(|| PackError::new(format!("pack-capability-unsupported:{}", base.capability), "the base rung"))?;
    let floors = rungs_of(&policy, &comp.id, RungRole::Floor);
    let [floor_rung] = floors.as_slice() else {
        return Err(PackError::new("floor-missing", format!("component {} needs one floor rung (§7.1)", comp.id)));
    };
    if !matches!(floor_rung.capability.as_str(), "wai.video.av1" | "wai.video.av1.lossless") {
        return Err(PackError::new(format!("pack-capability-unsupported:{}", floor_rung.capability), "the floor rung"));
    }
    if floor_rung.params.get("source").and_then(|v| v.as_str()) != Some("supplied") {
        return Err(PackError::new("pack-floor-encode-unsupported", "this revision packs a supplied floor"));
    }
    let tb = comp.timebase;
    let src = sources.get(&comp.id).ok_or_else(|| PackError::new("pack-source-missing", format!("no source for component {}", comp.id)))?;
    let (hd, frames) = crate::source::y4m::read_all(src).map_err(|e| PackError::new(e.code(), e.to_string()))?;
    if hd.rate.is_some_and(|(n, d)| (n, d) != (tb.num, tb.den)) {
        return Err(PackError::new("pack-source-rate", "the source's frame rate is not the component's timebase"));
    }
    if frames.is_empty() {
        return Err(PackError::new("y4m-truncated", "a source of no frames"));
    }
    let (w, h) = (hd.w, hd.h);
    let mut seg = Segmenter::new(SegParams::from_policy(&policy.segmentation, tb), w, h);
    let mut pieces = Vec::new();
    for f in &frames {
        pieces.extend(seg.push(&f.rgb));
    }
    pieces.extend(seg.finish());
    let starts: Vec<u64> = pieces.iter().map(|p| p.start).collect();
    let floor_listing = policy.floor.supplied_listing.clone();
    let fl = floor::floor_supplied(&opts.floor, floor_listing.as_deref(), &comp.id, &starts, frames.len() as u64, policy.floor.extra_keyframes)
        .map_err(|r| PackError::new(r.code, r.detail))?;
    let listing = floor_listing.expect("checked by floor_supplied");
    if let Some(m) = opts.meter.as_mut() {
        m.start();
    }

    // The string table: the component, the classes, then capabilities by first use.
    let mut strings: Vec<String> = vec![comp.id.clone()];
    strings.extend(policy.classes.iter().map(|c| c.id.clone()));
    let mut intern = |s: &str| -> u16 {
        match strings.iter().position(|x| x == s) {
            Some(i) => i as u16,
            None => {
                strings.push(s.to_string());
                (strings.len() - 1) as u16
            }
        }
    };
    let base_cap = intern(&base.capability);
    let floor_cap = intern(&floor_rung.capability);
    let mut objects: BTreeMap<[u8; 32], Vec<u8>> = BTreeMap::new();
    let mut entries: Vec<EntryRec> = Vec::new();
    let mut piece_recs: Vec<PieceRec> = Vec::new();
    let mut layer_bytes: Vec<(Vec<u64>, Option<u64>)> = Vec::new();
    let (fw, fh) = (fl.seq.max_frame_width, fl.seq.max_frame_height);
    let mut add = |env: Vec<u8>, entries: &mut Vec<EntryRec>, e: EntryRec| {
        let mut e = e;
        if entries.iter().any(|x| x.content_hash == e.content_hash) {
            e.flags |= wsix::ENTRY_DUP;
        }
        objects.insert(e.content_hash, env);
        entries.push(e);
    };
    for (p, piece) in pieces.iter().enumerate() {
        let kind: PieceKind = classify::classify(&piece.stats, &policy.classify);
        let pp = plan_piece(&policy, tb, piece.start, piece.len, kind);
        let src_frames: Vec<Vec<u8>> = frames[piece.start as usize..(piece.start + piece.len) as usize].iter().map(|f| f.rgb.clone()).collect();
        let q = choose_q(&src_frames, h, w, &mp).map_err(|e| PackError::new("pack-encode", e.to_string()))?;
        let n = src_frames.len() as u64;
        let base_env = objects::base_envelope("video", &base.capability, q.encoded.wiv1, n, h, w);
        let base_hash = objects::object_hash(&base_env);
        let first = entries.len() as u32;
        let tier = wsix::capability_prefix_tier(&base.capability);
        let digest = if tier == wsix::TIER_DECODE { clip_digest(&q.encoded.recon) } else { [0u8; 32] };
        let be = entry(&base_env, p as u64, base_cap, Role::Base, tier, digest, q.encoded.sse, 0)?;
        let base_len = be.byte_len;
        add(base_env, &mut entries, be);
        let df = objects::supplied_derived_from(&base_hash, &base.capability, &listing);
        let payload = floor::av1_payload(&fl.pieces[p].units, fh, fw, tb.num, tb.den);
        let floor_env = objects::floor_envelope("video", &floor_rung.capability, payload, n, fh, fw, &df);
        let fe = entry(&floor_env, p as u64, floor_cap, Role::Floor, wsix::TIER_NONE, [0u8; 32], wsix::SSE_UNKNOWN, wsix::ENTRY_DERIVED_FROM)?;
        add(floor_env, &mut entries, fe);
        piece_recs.push(PieceRec {
            start: piece.start,
            duration: u32::try_from(piece.len).map_err(|_| PackError::new("pack-piece-length", "a piece past 2^32 - 1 units"))?,
            first_entry: first,
            n_entries: (entries.len() as u32 - first) as u16,
            kind: kind as u8,
            flags: piece.flags,
            rules: pp.rules,
            moments: pp.moments,
        });
        layer_bytes.push((vec![base_len], pp.max_layer));
    }
    let plan: Vec<Vec<u8>> = policy
        .classes
        .iter()
        .map(|c| piece_recs.iter().zip(&layer_bytes).map(|(pr, (lb, rm))| class_plan(c, lb, *rm, pr.duration as u64, tb)).collect())
        .collect();
    let classes: Vec<ClassRec> = policy
        .classes
        .iter()
        .enumerate()
        .map(|(i, c)| ClassRec {
            name: (1 + i) as u16,
            max_layer: c.max_layer.min(255) as u8,
            flags: if c.floor_only { wsix::CLASS_FLOOR_ONLY } else { 0 },
            max_kbps: c.max_kbps.map_or(0, |k| k.min(u32::MAX as u64) as u32),
        })
        .collect();
    let policy_digest = policy.digest();
    let index = StageIndex {
        header: Header {
            media_class: MediaClass::Video,
            tb_num: tb.num,
            tb_den: tb.den,
            policy_digest,
            prev_index_digest: [0u8; 32],
            first_piece: 0,
            first_entry: 0,
            plan_digest: wsix::plan_digest(&starts),
            component: 0,
        },
        pieces: piece_recs,
        entries,
        classes,
        plan,
        strings,
    };
    index.check().map_err(|r| PackError::new(r.code, format!("the packer wrote an index that does not check: {}", r.detail)))?;
    let (init_path, init, frags) = cmaf_files(&comp.id, &fl, &starts, tb.den, tb.num).map_err(|e| PackError::new("floor-av1-malformed", e))?;
    let view = CompView {
        index: &index,
        tb,
        w,
        h,
        objects: &objects,
        floor: Some(FloorView { codec: fl.seq.codec_string(), w: fw, h: fh, init: init_path.clone(), fragments: frags.iter().map(|(p, b)| (p.clone(), b.len() as u64)).collect() }),
    };
    let catalogs = catalog::render_all(&view, &policy.layer_lag_ms, policy.deadline_ms, &policy_digest).map_err(|e| PackError::new("pack-catalog", e.0))?;
    let index_bytes = index.to_bytes();
    // A meter that was asked for and gives no figure (a span too short, a product past
    // its range) refuses the pack: it never turns into an unmetered package.
    let energy = match opts.meter.as_mut() {
        None => wai::energy_meter::EnergyReading::unmetered(),
        Some(m) => m.stop().ok_or_else(|| PackError::new("pack-meter", "the meter gave no figure for the pack"))?.as_seal_args(),
    };
    let receipt = PackageReceiptV2::seal(
        &opts.signer,
        &opts.signer_id,
        &policy.asset_id,
        policy_digest,
        &[(comp.id.clone(), index_bytes.clone())],
        &packager_id(),
        energy,
        opts.grant.clone(),
        Parent::None,
    )
    .map_err(|r| PackError::new(r.code, r.detail))?;
    let mut files: BTreeMap<String, Vec<u8>> = BTreeMap::new();
    files.insert("policy.toml".into(), policy_bytes.to_vec());
    files.insert("policy.json".into(), policy.canonical_json().into_bytes());
    files.insert(format!("plan-{}.txt", comp.id), plan_file(&starts).into_bytes());
    files.insert(wai::package::index_file(&comp.id), index_bytes.clone());
    for (h, env) in objects {
        files.insert(objects::object_path(&h), env);
    }
    files.insert(init_path, init);
    for (p, b) in frags {
        files.insert(p, b);
    }
    files.extend(catalogs);
    files.insert("package.receipt.json".into(), receipt.to_json().into_bytes());
    if receipt.energy_provenance.is_some() {
        let claim = measurement_claim(&receipt, &index_bytes, opts).map_err(|e| PackError::new("pack-measurement-claim", e))?;
        files.insert(claim_path(&receipt), claim.into_bytes());
    }
    let view: BTreeMap<String, &[u8]> = files.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
    let listing = wai::pin::listing_from(&view).map_err(|e| PackError::new("pack-listing", e.to_string()))?;
    files.insert("LISTING".into(), listing);
    Ok(PackOutput { files, receipt })
}

/// Where a package's measurement claim is: `claims/<receipt hash>.measurement.json`.
pub fn claim_path(r: &PackageReceiptV2) -> String {
    format!("claims/{}.measurement.json", crate::hex(&r.receipt_hash()))
}

/// The `measurement` claim about a metered receipt's figure (jwp-receipts §2.4.8), signed
/// with the package key and checked against the receipt before it is returned.
fn measurement_claim(r: &PackageReceiptV2, index: &[u8], opts: &PackOpts) -> Result<String, String> {
    use wai::energy_figure::{meter_report, EnergyFigures};
    let indexes: Vec<&[u8]> = vec![index];
    let view = (r, indexes.as_slice()).figure_view().ok_or("the receipt does not verify")?;
    let claim = meter_report::emit(&view, None, None, wai::energy_report::BuildRecord::of_this_build(), &opts.power_reader, &opts.signer, &opts.signer_id, None, &opts.issued_at)?;
    Ok(claim.to_json())
}