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

//! A supplied floor ([packaging] §7.3): a directory of AV1 streams, one per video
//! component, named `<component>.obu` in the low-overhead OBU format, whose listing
//! (SPEC §3.1) the policy pins by SHA-256 (`floor.supplied_listing`). The floor is checked
//! against the plan before anything is written:
//!
//! - the directory's listing hashes to `floor.supplied_listing` (`floor-listing-mismatch`);
//! - the component's stream is in it (`floor-stream-missing`) and parses
//!   (`floor-av1-malformed`): temporal units, a sequence header in every unit that starts
//!   at a key frame, every sequence header the same;
//! - it has one temporal unit per source frame (`floor-frame-count`);
//! - a unit starts at a shown key frame at every piece start of the plan, and at no other
//!   frame unless `floor.extra_keyframes` is true (`floor-keyframe-mismatch`).
//!
//! [packaging]: ../../../extensions/packaging.md

use crate::av1::{self, SequenceHeader};
use std::collections::BTreeMap;

/// Why a supplied floor is refused.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FloorReason {
    pub code: &'static str,
    pub detail: String,
}

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

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

fn refuse<T>(code: &'static str, detail: impl Into<String>) -> Result<T, FloorReason> {
    Err(FloorReason { code, detail: detail.into() })
}

/// One piece of a component's floor: its temporal units, each as the stream holds it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FloorPiece {
    pub units: Vec<Vec<u8>>,
    /// Which units start at a shown key frame.
    pub key: Vec<bool>,
}

/// A component's checked floor: its sequence header and the OBU that carries it, and its
/// pieces.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Floor {
    pub seq: SequenceHeader,
    pub seq_obu: Vec<u8>,
    pub pieces: Vec<FloorPiece>,
}

/// The SHA-256 of a set's listing, as 64 lowercase hexadecimal digits.
pub fn listing_sha256(files: &BTreeMap<String, Vec<u8>>) -> Result<String, FloorReason> {
    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| FloorReason { code: "floor-listing-mismatch", detail: e.to_string() })?;
    Ok(crate::sha256_hex(&listing))
}

/// Check component `component`'s supplied floor against its plan: `starts` (the piece
/// starts, in frames) of a source of `frames` frames.
pub fn floor_supplied(files: &BTreeMap<String, Vec<u8>>, listing: Option<&str>, component: &str, starts: &[u64], frames: u64, extra_keyframes: bool) -> Result<Floor, FloorReason> {
    let got = listing_sha256(files)?;
    if listing != Some(got.as_str()) {
        return refuse("floor-listing-mismatch", format!("the supplied floor's listing is {got}, the policy pins {}", listing.unwrap_or("none")));
    }
    let name = format!("{component}.obu");
    let stream = files.get(&name).ok_or_else(|| FloorReason { code: "floor-stream-missing", detail: format!("no {name} in the supplied floor") })?;
    let malformed = |e: av1::Av1Error| FloorReason { code: "floor-av1-malformed", detail: e.0 };
    let tus = av1::temporal_units(stream).map_err(malformed)?;
    let first = tus[0].sequence_header().ok_or_else(|| FloorReason { code: "floor-av1-malformed", detail: "the first temporal unit carries no sequence header".into() })?;
    let seq_obu = first.bytes.to_vec();
    let seq = SequenceHeader::parse(first.payload).map_err(malformed)?;
    if tus.len() as u64 != frames {
        return refuse("floor-frame-count", format!("{} temporal units for {frames} source frames", tus.len()));
    }
    let mut key = Vec::with_capacity(tus.len());
    for (i, tu) in tus.iter().enumerate() {
        if let Some(s) = tu.sequence_header()
            && s.bytes != seq_obu.as_slice()
        {
            return refuse("floor-av1-malformed", format!("temporal unit {i} carries another sequence header"));
        }
        let k = av1::starts_with_shown_key_frame(tu, &seq).map_err(malformed)?;
        if k && tu.sequence_header().is_none() {
            return refuse("floor-av1-malformed", format!("temporal unit {i} starts a key frame without a sequence header"));
        }
        key.push(k);
    }
    for (i, &k) in key.iter().enumerate() {
        let at_start = starts.binary_search(&(i as u64)).is_ok();
        if at_start && !k {
            return refuse("floor-keyframe-mismatch", format!("no key frame at piece start {i}"));
        }
        if !at_start && k && !extra_keyframes {
            return refuse("floor-keyframe-mismatch", format!("a key frame at {i}, inside a piece"));
        }
    }
    let mut pieces = Vec::with_capacity(starts.len());
    for (p, &s) in starts.iter().enumerate() {
        let e = starts.get(p + 1).copied().unwrap_or(frames);
        let (s, e) = (s as usize, e as usize);
        pieces.push(FloorPiece { units: tus[s..e].iter().map(av1::TemporalUnit::bytes).collect(), key: key[s..e].to_vec() });
    }
    Ok(Floor { seq, seq_obu, pieces })
}

/// The payload of a `wai.video.av1` object as the reference sink decodes it
/// (`wai-rs` `codecs::video`): `n_frames`, `h`, `w`, `fps_num`, `fps_den` as u32
/// little-endian, then each temporal unit as a u32 little-endian length and its bytes.
pub fn av1_payload(units: &[Vec<u8>], h: u32, w: u32, fps_num: u32, fps_den: u32) -> Vec<u8> {
    let mut o = Vec::with_capacity(20 + units.iter().map(|u| 4 + u.len()).sum::<usize>());
    for v in [units.len() as u32, h, w, fps_num, fps_den] {
        o.extend_from_slice(&v.to_le_bytes());
    }
    for u in units {
        o.extend_from_slice(&(u.len() as u32).to_le_bytes());
        o.extend_from_slice(u);
    }
    o
}

/// A reference-framing payload's header after its unit count: `(h, w, fps_num, fps_den)`.
pub type Av1Header = (u32, u32, u32, u32);

/// The temporal units of a `wai.video.av1` payload in the reference framing, and its
/// header; `None` when it is not that framing.
pub fn av1_units(payload: &[u8]) -> Option<(Vec<Vec<u8>>, Av1Header)> {
    let u32_at = |i: usize| payload.get(i..i + 4).map(|b| u32::from_le_bytes(b.try_into().expect("four bytes")));
    let n = u32_at(0)? as usize;
    let hdr = (u32_at(4)?, u32_at(8)?, u32_at(12)?, u32_at(16)?);
    let mut at = 20usize;
    let mut units = Vec::new();
    while at < payload.len() {
        let len = u32_at(at)? as usize;
        let end = (at + 4).checked_add(len).filter(|&e| e <= payload.len())?;
        units.push(payload[at + 4..end].to_vec());
        at = end;
    }
    (units.len() == n).then_some((units, hdr))
}