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

//! An integer payload whose header asks for more than the caps allow, or more than its
//! parameter set can fit, is refused before the decode allocates anything sized from
//! it (`extensions/integer-payloads.md` §7).
//!
//! Each case below is a committed payload with header fields changed. The parsers cap
//! every tensor a header sizes (`too_large`) before reading a section, and the decode
//! derives each network's shapes from its layer geometry before it allocates a latent;
//! this file checks both directly with a global allocator that records the largest
//! single allocation, and asserts that refusing each payload allocates at once no more
//! than parsing its parameter set needs (each parsed weight is held as an `i64`, so at
//! most four times the set's bytes) or a few kilobytes, whichever is more: nothing
//! sized from the header.
//!
//! Its own test binary with a single test, so that no other test's allocations land in
//! the window it measures.
#![cfg(all(feature = "video_int", feature = "audio_int", feature = "splat_int"))]

use std::alloc::{GlobalAlloc, Layout, System};
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};

use wai::int_payload::{decode_int, resolve, IntParams};
use wai::pin::HeldBytes;

/// The system allocator, recording the largest single request.
struct Largest;
static LARGEST: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Largest {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        LARGEST.fetch_max(l.size(), Relaxed);
        unsafe { System.alloc(l) }
    }
    unsafe fn alloc_zeroed(&self, l: Layout) -> *mut u8 {
        LARGEST.fetch_max(l.size(), Relaxed);
        unsafe { System.alloc_zeroed(l) }
    }
    unsafe fn realloc(&self, p: *mut u8, l: Layout, new: usize) -> *mut u8 {
        LARGEST.fetch_max(new, Relaxed);
        unsafe { System.realloc(p, l, new) }
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        unsafe { System.dealloc(p, l) }
    }
}

#[global_allocator]
static ALLOC: Largest = Largest;

/// The largest single allocation `f` makes, and its result.
fn largest_in<R>(f: impl FnOnce() -> R) -> (R, usize) {
    LARGEST.store(0, Relaxed);
    let r = f();
    (r, LARGEST.load(Relaxed))
}

fn read(rel: &str) -> Vec<u8> {
    let p = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors").join(rel);
    std::fs::read(&p).unwrap_or_else(|e| panic!("{}: {e}", p.display()))
}

/// `p` with the u32 fields at the given offsets replaced.
fn with(p: &[u8], fields: &[(usize, u32)]) -> Vec<u8> {
    let mut p = p.to_vec();
    for &(off, v) in fields {
        p[off..off + 4].copy_from_slice(&v.to_le_bytes());
    }
    p
}

const FEW_KB: usize = 16 << 10;

/// A rANS stream of `n` symbols, symbol `i` coded with `(start, freq)` from `sym(i)`,
/// from the state 2^31, so it ends canonically: what an encoder writes.
fn encode(n: u64, sym: impl Fn(u64) -> (u64, u64)) -> Vec<u8> {
    const L: u64 = 1 << 31;
    let mut state = L;
    let mut renorm: Vec<u32> = Vec::new();
    let mut i = n;
    while i > 0 {
        i -= 1;
        let (start, freq) = sym(i);
        let x_max = ((L >> 16) << 32) * freq;
        while state >= x_max {
            renorm.push(state as u32);
            state >>= 32;
        }
        state = ((state / freq) << 16) + (state % freq) + start;
    }
    renorm.reverse();
    let mut words = vec![state as u32, (state >> 32) as u32];
    words.extend(renorm);
    words.iter().flat_map(|w| w.to_le_bytes()).collect()
}

#[test]
fn a_header_past_the_caps_or_the_model_is_refused_before_anything_is_allocated() {
    use wai::codecs as c;
    let hyp: BTreeMap<String, Vec<u8>> = [("entropy.json", "int_hyper/tiny_entropy.json"), ("g_s.whm", "int_hyper/tiny_g_s.bin"), ("h_s.whm", "int_hyper/tiny_h_s.bin")]
        .into_iter().map(|(k, v)| (k.to_string(), read(v))).collect();
    let hyp_view: BTreeMap<String, &[u8]> = hyp.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
    let unpinned = |cap, held| resolve(cap, None, &[held], None).expect("an unpinned set resolves");
    let hyp = unpinned(c::CAP_NEURAL_INT_HYPER, HeldBytes::Set { listing: b"", files: &hyp_view });
    let (model, gs) = (read("int_audio/audio_model.bin"), read("int_splat/splat_gs.bin"));
    let none = IntParams::none();
    let (aud, spl) = (unpinned(c::CAP_AUDIO_INT_CODEC, HeldBytes::File(&model)), unpinned(c::CAP_SPLAT_INT_CODEC, HeldBytes::File(&gs)));
    let (iv, nk, ih) = (read("int_video/int_video_bitstream.wiv1"), read("int_video/video_nk_bitstream.wiv1"), read("int_hyper/tiny_bitstream.wih1"));
    let (au, sp) = (read("int_audio/audio_bitstream.wia1"), read("int_splat/splat.wis1"));

    // The counter sees a decode's allocations: the committed clip's largest is at least
    // one 48 × 48 frame of RGB8.
    let (r, largest) = largest_in(|| decode_int(c::CAP_VIDEO_INT_MOTION, &iv, &IntParams::none()));
    r.expect("the committed clip decodes");
    assert!(largest >= 48 * 48 * 3, "the counter saw a {largest}-byte allocation at most");

    let hyp_bytes: usize = hyp_view.values().map(|v| v.len()).sum();
    let cases: Vec<(&str, &str, Vec<u8>, &IntParams<'_>, &str, usize)> = vec![
        // A 2^20-frame clip of 48 × 48: each frame is within the cap, the clip is not.
        ("a 2^20-frame clip", c::CAP_VIDEO_INT_MOTION, with(&iv, &[(8, 1 << 20)]), &none, "too_large", 0),
        // An 8192 × 8192 frame.
        ("an 8192 × 8192 frame", c::CAP_VIDEO_INT_HYPER, with(&nk, &[(12, 8192), (16, 8192)]), &hyp, "too_large", hyp_bytes),
        // A hyper-latent of 2 × 6000 × 6000 = 72 million elements, past the parser's
        // 2^26 cap.
        ("a 72-million-element hyper-latent", c::CAP_NEURAL_INT_HYPER, with(&ih, &[(24, 6000), (28, 6000)]), &hyp, "too_large", hyp_bytes),
        // A hyper-latent of 2 × 2000 × 2000, within the cap, that h_s would expand to
        // 8 × 4000 × 4000, past it.
        ("a hyper-latent h_s expands past the cap", c::CAP_NEURAL_INT_HYPER, with(&ih, &[(24, 2000), (28, 2000)]), &hyp, "too_large", hyp_bytes),
        // An audio latent of 16 × 2^20, within the cap, that the synthesis would
        // upsample past it.
        ("an audio latent the synthesis expands past the cap", c::CAP_AUDIO_INT_CODEC, with(&au, &[(16, 1 << 20)]), &aud, "too_large", model.len()),
        // A splat latent of 32 × 1024 × 1024, within the cap, that g_s would expand to
        // 14 × 4096 × 4096, past it.
        ("a splat latent g_s expands past the cap", c::CAP_SPLAT_INT_CODEC, with(&sp, &[(12, 1024), (16, 1024)]), &spl, "too_large", gs.len()),
    ];
    let parsed = wai::int_payload::wire::parse_wih1(&with(&ih, &[(24, 6000), (28, 6000)])).map(|_| ()).map_err(|e| e.code());
    assert_eq!(parsed, Err("too_large"), "the parser's cap refuses the 72-million-element hyper-latent");
    for (what, cap, bad, params, code, set_bytes) in cases {
        let (r, largest) = largest_in(|| decode_int(cap, &bad, params));
        let bound = FEW_KB.max(4 * set_bytes);
        assert!(largest <= bound, "{what}: refusing it allocated {largest} bytes at once, over {bound} ({r:?})");
        assert_eq!(r.map(|_| ()).map_err(|e| e.code()), Err(code), "{what}");
    }

    // Headers every cap admits, whose streams are 8 bytes and cannot carry what the
    // header declares (§2, §7): each is refused at the first symbol that runs its
    // stream out, having allocated only what it decoded. Before this, the 16-frame
    // clip of 4729 × 4729 ran for about a minute and reached gigabytes before it was
    // refused at its last symbol.
    use wai::int_payload::wire;
    let (zeros, ones) = ([0u8; 8], [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f]);
    let clip = |n_frames: u32, h: u32, w: u32, block: u32, cdf: Vec<i32>, rans: &[u8]| wire::write_wiv1(&wire::Wiv1 {
        mv_precision: 1, keyframe: 0, n_frames, h, w, block, q: 1, offset: 0, cdf,
        mvs: vec![(0, 0); n_frames as usize - 1], rans, wih1: None,
    });
    let wih1 = {
        // The tiny h_s doubles the hyper-latent's sides.
        let mut k = wire::parse_wih1(&ih).unwrap();
        (k.zh, k.zw, k.sh, k.sw) = (64, 64, 128, 128);
        k.z_rans = &ones[..];
        wire::write_wih1(&k)
    };
    let wia1 = {
        // The committed synthesis gives 16 samples per latent sample (512 → 8192), so
        // t_out is what a latent of 4096 presents.
        let mut a = wire::parse_wia1(&au).unwrap();
        a.lat_len = 1 << 12;
        a.t_out = 16 << 12;
        a.rans = &ones[..];
        wire::write_wia1(&a)
    };
    let wis1 = {
        // The splat g_s quadruples the latent's sides.
        let mut v = wire::parse_wis1(&sp).unwrap();
        (v.sh, v.sw, v.s) = (64, 64, 256);
        v.rans = &ones[..];
        wire::write_wis1(&v)
    };
    let starved: Vec<(&str, &str, Vec<u8>, &IntParams<'_>, usize)> = vec![
        // Small enough that a decode which ran on past the first exhausted symbol would
        // finish here quickly, and allocate a 256 × 256 frame.
        ("16 frames of 256 × 256", c::CAP_VIDEO_INT_MOTION, clip(16, 256, 256, 1 << 31, vec![0, 1, 65_536], &zeros), &none, 0),
        ("16 frames of 4729 × 4729", c::CAP_VIDEO_INT_MOTION, clip(16, 4729, 4729, 1 << 31, vec![0, 1, 65_536], &zeros), &none, 0),
        ("one frame of 4096 × 5461", c::CAP_VIDEO_INT_MOTION, clip(1, 4096, 5461, 1, vec![0, 65_534, 65_535, 65_536], &ones), &none, 0),
        ("a hyper-latent of 2 × 64 × 64", c::CAP_NEURAL_INT_HYPER, wih1, &hyp, hyp_bytes),
        ("an audio latent of 16 × 4096", c::CAP_AUDIO_INT_CODEC, wia1, &aud, model.len()),
        ("a splat latent of 32 × 64 × 64", c::CAP_SPLAT_INT_CODEC, wis1, &spl, gs.len()),
    ];
    for (what, cap, bad, params, set_bytes) in starved {
        assert!(wai::container::check_int_payload_framing(cap, &bad).is_ok(), "{what}: the payload parses");
        let (r, largest) = largest_in(|| decode_int(cap, &bad, params));
        let bound = FEW_KB.max(4 * set_bytes);
        assert!(largest <= bound, "{what}: refusing it allocated {largest} bytes at once, over {bound} ({r:?})");
        assert_eq!(r.map(|_| ()).map_err(|e| e.code()), Err("stream_exhausted"), "{what}");
    }

    // A WIA1 of 384 bytes whose stream is fully valid (8.4 million symbols of frequency
    // 65535/65536) for a latent of 16 × 524287 that presents one sample: its synthesis
    // would hold gigabytes for that one sample. The latent is longer than t_out needs,
    // so the plan refuses it (geometry_mismatch) before a symbol is decoded, and the
    // declared cost refuses it alike.
    let wia1 = {
        let mut a = wire::parse_wia1(&au).unwrap();
        let n = a.n_ch as usize;
        a.lat_len = 524_287;
        a.t_out = 1;
        a.offset = 0;
        a.median_q = vec![0; n];
        a.cdfs = vec![vec![0, 65_535, 65_536]; n];
        let rans = encode(n as u64 * 524_287, |_| (0, 65_535));
        let mut b = a.clone();
        b.rans = &rans;
        wire::write_wia1(&b)
    };
    assert!(wia1.len() <= 400, "the payload is {} bytes", wia1.len());
    assert!(wai::container::check_int_payload_framing(c::CAP_AUDIO_INT_CODEC, &wia1).is_ok(), "it parses");
    let (r, largest) = largest_in(|| decode_int(c::CAP_AUDIO_INT_CODEC, &wia1, &aud));
    let bound = FEW_KB.max(4 * model.len());
    assert!(largest <= bound, "the 384-byte WIA1: refusing it allocated {largest} bytes at once, over {bound}");
    assert_eq!(r.map(|_| ()).map_err(|e| e.code()), Err("geometry_mismatch"), "the 384-byte WIA1");
    assert_eq!(wai::int_payload::declared_cost(c::CAP_AUDIO_INT_CODEC, &wia1, &aud).map(|_| ()).map_err(|e| e.code()),
               Err("geometry_mismatch"), "the declared cost refuses it too");

    // The tools' JSON form, read by int_video::VideoJson (the converter's and the JSON
    // decode's reader): a header past the caps is refused before anything is sized
    // from it, on every target. A side of 2^31 made the frame size overflow and the
    // reader panic; sixteen frames of 4096 × 4096 at block 1, each with an empty list,
    // were padded to 2^28 motion vectors before any cap was checked.
    use serde_json::json;
    let js = |side: u64, nf: usize| json!({
        "block": 1, "q": 1, "offset": 0, "cdf": [0, 1, 65_536], "n_frames": nf, "h": side, "w": side,
        "mvs": vec![json!([]); nf], "rans_b64": "AAAAAAAAAAA="
    });
    for (what, j) in [("h = w = 2^31", js(1 << 31, 2)), ("2^28 padded motion vectors", js(4096, 16))] {
        let (r, largest) = largest_in(|| std::panic::catch_unwind(|| (
            wai::int_video::VideoJson::from_value(&j).is_err(),
            wai::int_payload::wiv1_from_json(&j, None).map(|_| ()).map_err(|e| e.code()),
            wai::int_video::decode(&j).is_err(),
        )));
        assert_eq!(r.expect("the JSON reader must not panic"), (true, Err("field_range"), true), "{what}");
        assert!(largest <= FEW_KB, "{what}: refusing it allocated {largest} bytes at once");
    }
}