open-standards/wai/wai-rs/tests/int_payload.rs

// wai-conformance: int-kernels
//! The payloads of the WAI-native integer codecs (`extensions/integer-payloads.md`)
//! against the committed corpus `wai/int-payload-conformance/`:
//!
//! - every refusal fixture is refused with the code its index states, at the layer it
//!   states (parsing alone, or the decode);
//! - every strict prefix of every committed payload is refused, and every changed byte
//!   in a sweep over each payload is refused or decoded — never a panic, in either
//!   profile (CI runs this file with overflow checks on and off);
//! - a tool's JSON output carried as an envelope's payload under each pinned string is
//!   refused (`bad_magic`) by the framing check `wai wrap` runs and by `decode_int`.
#![cfg(all(feature = "video_int", feature = "audio_int", feature = "splat_int"))]

use std::collections::BTreeMap;
use std::path::PathBuf;

use serde_json::Value;

use wai::container::{Conditioning, Manifest, ModelRequirement, Rendition, Wai, WaiMulti};
use wai::int_payload::{decode_int, wire, IntParams, Refusal};
use wai::pin::HeldBytes;

fn wai_dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}
fn read(rel: &str) -> Vec<u8> {
    let p = wai_dir().join(rel);
    std::fs::read(&p).unwrap_or_else(|e| panic!("read {}: {e}", p.display()))
}
fn index() -> Value {
    serde_json::from_slice(&read("int-payload-conformance/vectors.json")).expect("the index parses")
}

/// The set an index entry names, read: member name → bytes.
fn set_of(entry: &Value) -> BTreeMap<String, Vec<u8>> {
    entry["set"].as_object().into_iter().flatten().map(|(k, v)| (k.clone(), read(v.as_str().unwrap()))).collect()
}

fn decode(cap: &str, payload: &[u8], set: &BTreeMap<String, Vec<u8>>) -> Result<wai::int_payload::IntDecoded, wai::int_payload::IntPayloadError> {
    let view: BTreeMap<String, &[u8]> = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
    let params = wai::int_payload::resolve(cap, None, &[HeldBytes::Set { listing: b"", files: &view }], None)
        .unwrap_or_else(|_| IntParams::none());
    decode_int(cap, payload, &params)
}

fn parse_err(cap: &str, p: &[u8]) -> Option<wire::WireError> {
    match wai::codecs::pinned_int_payload(cap)? {
        wire::MAGIC_WIV1 => wire::parse_wiv1(p, Some((cap == wai::codecs::CAP_VIDEO_INT_HYPER) as u8)).err(),
        wire::MAGIC_WIH1 => wire::parse_wih1(p).err(),
        wire::MAGIC_WIA1 => wire::parse_wia1(p).err(),
        _ => wire::parse_wis1(p).err(),
    }
}

/// Every fixture is refused with its code, at its layer; and every code a payload's
/// bytes can cause has a fixture.
#[test]
fn every_refusal_fixture_is_refused_with_its_code() {
    let idx = index();
    let mut seen = std::collections::BTreeSet::new();
    for f in idx["refusals"].as_array().unwrap() {
        let (name, cap, code) = (f["name"].as_str().unwrap(), f["capability"].as_str().unwrap(), f["code"].as_str().unwrap());
        let payload = read(f["payload"].as_str().unwrap());
        let parsed = parse_err(cap, &payload);
        match f["layer"].as_str().unwrap() {
            "parse" => assert_eq!(parsed.map(|e| e.code()), Some(code), "{name}: refused by parsing alone"),
            _ => assert_eq!(parsed, None, "{name}: parses, and is refused by the decode"),
        }
        let r = decode(cap, &payload, &set_of(f));
        assert_eq!(r.as_ref().map(|_| ()).map_err(|e| e.code()), Err(code), "{name}");
        seen.insert(code);
    }
    for code in ["bad_magic", "format_capability_mismatch", "keyframe_capability_mismatch", "version", "reserved_nonzero",
                 "field_range", "too_large", "cdf_malformed", "stream_length", "nested_length", "trailing",
                 "stream_exhausted", "stream_noncanonical", "geometry_mismatch", "model_mismatch"] {
        assert!(Refusal::from_code(code).is_some(), "{code} is a code");
        assert!(seen.contains(code), "no fixture is refused with {code}");
    }
}

/// Every strict prefix of every committed payload is refused; a sweep of single-byte
/// changes over each is refused or decoded, and nothing panics. The sweep changes
/// every byte of the fixed header and 96 bytes spread evenly over the rest (low bit
/// and high bit by turns), so each payload's tables, length fields and stream are
/// all reached while the decodes stay few enough for a build with overflow checks.
#[test]
fn every_prefix_and_byte_change_is_refused_or_decoded() {
    let idx = index();
    for v in idx["vectors"].as_array().unwrap() {
        let (name, cap) = (v["name"].as_str().unwrap(), v["capability"].as_str().unwrap());
        let payload = read(v["payload"].as_str().unwrap());
        let set = set_of(v);
        assert!(decode(cap, &payload, &set).is_ok(), "{name}: the committed payload decodes");
        for n in 0..payload.len() {
            assert!(parse_err(cap, &payload[..n]).is_some(), "{name}: the {n}-byte prefix parses");
        }
        let (mut refused, mut decoded) = (0, 0);
        const HEADER: usize = 48;
        let stride = (payload.len() - HEADER).div_ceil(96).max(1);
        for (n, at) in (0..payload.len()).filter(|&i| i < HEADER || (i - HEADER) % stride == 0).enumerate() {
            let mut p = payload.clone();
            p[at] ^= if n % 2 == 0 { 0x01 } else { 0x80 };
            match decode(cap, &p, &set) {
                Ok(_) => decoded += 1,
                Err(_) => refused += 1,
            }
        }
        assert!(refused > 0, "{name}: no change was refused");
        println!("{name}: {refused} changes refused, {decoded} decoded, no panic");
    }
}

/// A tool's JSON output wrapped as the payload under each pinned string: the framing
/// check (`Wai::check_int_payload_framing`, `WaiMulti::check_int_payload_framing`,
/// which `wai wrap` and `wai inspect` run) and `decode_int` both refuse it as
/// `bad_magic`; its binary conversion passes both.
#[test]
fn a_json_form_in_an_envelope_is_refused() {
    use wai::codecs as c;
    for (cap, json, bin) in [
        (c::CAP_VIDEO_INT_MOTION, "wai-rs/tests/vectors/int_video/int_video_bitstream.json", "wai-rs/tests/vectors/int_video/int_video_bitstream.wiv1"),
        (c::CAP_VIDEO_INT_HYPER, "wai-rs/tests/vectors/int_video/video_nk_bitstream.json", "wai-rs/tests/vectors/int_video/video_nk_bitstream.wiv1"),
        (c::CAP_NEURAL_INT_HYPER, "wai-rs/tests/vectors/int_hyper/tiny_bitstream.json", "wai-rs/tests/vectors/int_hyper/tiny_bitstream.wih1"),
        (c::CAP_AUDIO_INT_CODEC, "wai-rs/tests/vectors/int_audio/audio_bitstream.json", "wai-rs/tests/vectors/int_audio/audio_bitstream.wia1"),
        (c::CAP_SPLAT_INT_CODEC, "wai-rs/tests/vectors/int_splat/splat.json", "wai-rs/tests/vectors/int_splat/splat.wis1"),
    ] {
        let wrap = |payload: Vec<u8>| Wai::new(Manifest::new("binary", "improve", ModelRequirement::new(cap), Conditioning::new("bytes")), payload);
        let env = Wai::from_bytes(&wrap(read(json)).to_bytes().unwrap()).unwrap();
        assert_eq!(env.check_int_payload_framing().map_err(|e| e.code), Err("bad_magic"), "{cap}");
        assert_eq!(decode_int(cap, &env.payload, &IntParams::none()).map(|_| ()).map_err(|e| e.code()), Err("bad_magic"), "{cap}");
        assert_eq!(wrap(read(bin)).check_int_payload_framing(), Ok(()), "{cap}: the binary form");
        let multi = WaiMulti::new("binary", "improve", Value::Null, vec![Rendition::media(cap, "bytes", read(json))]);
        assert_eq!(multi.check_int_payload_framing().map_err(|(i, e)| (i, e.code)), Err((0, "bad_magic")), "{cap}");
    }
    // The fallback decodes the same payload (SPEC §4 step 5), so its framing is checked too.
    let env = Wai::new(
        Manifest::new("video", "improve", ModelRequirement::new(c::CAP_VIDEO_AV1).with_fallback(c::CAP_VIDEO_INT_MOTION), Conditioning::new("bytes")),
        read("wai-rs/tests/vectors/int_video/int_video_bitstream.json"));
    assert_eq!(env.check_int_payload_framing().map_err(|e| e.code), Err("bad_magic"));
}