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

//! A `wai.neural.int_mlicpp` payload whose header the model cannot fit is refused
//! before the decode allocates anything sized from it.
//!
//! The header caps each tensor it sizes directly, but what the networks allocate
//! depends on the model as well: a small, well-formed header can declare a hyper-latent
//! that `h_s` would expand far past the cap, or an image or latent that does not match
//! what the networks produce. The decode derives every network's shapes from its layer
//! geometry first (`int_hyper_synth::net_shape`) and refuses such a payload with
//! nothing allocated. This file checks that directly: it installs a global allocator
//! that records the largest single allocation, and asserts that refusing each
//! malformed payload allocates no more than a few kilobytes. On a 32-bit target, where
//! the allocation that check prevents would abort the process rather than fail, the
//! same test shows the decode returns an error instead.
//!
//! The model is loaded the same way: a damaged compact network (`WHM1`) is refused
//! without the parser reserving for what its counts claim.
//!
//! Its own test binary with a single test, so that no other test's allocations land in
//! the window it measures.
#![cfg(feature = "neural_int")]

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

use wai::container::Wai;
use wai::int_hyper_synth::{IntLayer, IntModel};
use wai::int_mlicpp::{decode, Header, IntMlicError, IntMlicModel};

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

/// A header field of the committed payload replaced (offsets as in `parse_header`).
fn with(p: &[u8], off: usize, bytes: &[u8]) -> Vec<u8> {
    let mut p = p.to_vec();
    p[off..off + bytes.len()].copy_from_slice(bytes);
    p
}

const FEW_KB: usize = 16 << 10;

#[test]
fn a_header_the_model_cannot_fit_is_refused_before_anything_is_allocated() {
    let model = IntMlicModel::from_bundle(&read("model.wmb")).expect("model");
    let pl = Wai::from_bytes(&read("stream.wai")).expect("envelope").payload;

    // The counter sees the decode's allocations: the committed payload's largest is at
    // least its 64 × 64 RGB image.
    let (img, largest) = largest_in(|| decode(&model, &pl));
    img.expect("the committed payload decodes");
    assert!(largest >= 64 * 64 * 3, "the counter saw a {largest}-byte allocation at most");

    // A compact network claiming 2^32 − 1 layers in 20 bytes: refused as truncated,
    // with nothing reserved for the layers it claims (about 480 GB of them).
    let whm1 = [b"WHM1".as_slice(), &12u32.to_le_bytes(), &16u32.to_le_bytes(), &24u32.to_le_bytes(),
                &u32::MAX.to_le_bytes()].concat();
    let (r, largest) = largest_in(|| IntModel::from_bytes(&whm1).err());
    assert!(largest <= FEW_KB, "parsing a damaged network reserved {largest} bytes");
    assert_eq!(r, Some("model truncated"));

    let u16le = |v: u16| v.to_le_bytes();
    let u32le = |v: u32| v.to_le_bytes();
    let cases: [(&str, Vec<u8>, &str); 3] = [
        // z [4, 4000, 4000]: 64 Mi elements, within the header's cap, but h_s would
        // expand it to [16, 16000, 16000] where the latent is [16, 8, 8].
        ("an inflated hyper-latent", with(&with(&pl, 14, &u16le(4000)), 16, &u16le(4000)), "h_s gives"),
        // A 1000 × 1000 latent where h_s gives 8 × 8.
        ("an inflated latent", with(&with(&pl, 20, &u16le(1000)), 22, &u16le(1000)), "h_s gives"),
        // A 4000 × 4000 image where g_s gives 64 × 64.
        ("an inflated image", with(&with(&pl, 4, &u32le(4000)), 8, &u32le(4000)), "g_s gives"),
    ];
    for (what, bad, why) in cases {
        let (r, largest) = largest_in(|| decode(&model, &bad));
        // Allocation first: a decoder can also refuse these later, once a stream runs
        // out, but only after allocating and filling what the header asked for.
        assert!(largest <= FEW_KB, "{what}: refusing it allocated {largest} bytes at once ({r:?})");
        match r {
            Err(IntMlicError::Model(e)) => assert!(e.contains(why), "{what}: {e}"),
            other => panic!("{what}: expected a model refusal, got {other:?}"),
        }
    }

    // A model whose synthesis, on the committed (well-formed) header's own shapes,
    // would form a buffer far past the cap on the way to the right output: a stride-2^20
    // transposed convolution spreads the 8 × 8 latent over a plane of about 7.3 million
    // squared, and a stride-2^20 convolution brings it back to 8 × 8 before the rest of
    // g_s. Its output shape is the image's, so only the size check can refuse it.
    let s = 1usize << 20;
    let mut wide = IntMlicModel::from_bundle(&read("model.wmb")).expect("model");
    let unit = |out_ch| (vec![1i64; 8 * out_ch], vec![1i64; out_ch], vec![0i64; out_ch]);
    let (weight, mult, bias) = unit(8);
    wide.g_s.layers.insert(0, IntLayer::ConvT { out_ch: 8, k: 1, stride: s, pad: 0, opad: 0, weight, mult, bias });
    let (weight, mult, bias) = unit(8);
    wide.g_s.layers.insert(1, IntLayer::Conv { out_ch: 8, k: 1, stride: s, pad: 0, weight, mult, bias });
    let (r, largest) = largest_in(|| decode(&wide, &pl));
    assert!(largest <= FEW_KB, "a synthesis wider than the cap: refusing it allocated {largest} bytes at once ({r:?})");
    match r {
        // On a 64-bit target the plane's size is formed, and is over the cap; on a
        // 32-bit one it does not fit `usize`, and the geometry is refused as such.
        Err(IntMlicError::Model(e)) => assert!(
            e.contains("g_s") && (e.contains(&format!("over the {}-element cap", Header::MAX_TENSOR)) || e.contains("overflows")),
            "{e}"),
        other => panic!("a synthesis wider than the cap: expected a model refusal, got {other:?}"),
    }
}