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

//! What reading a manifest allocates (SPEC §2, *Reader limits*; SECURITY.md, *Manifest
//! parsing*). A counting global allocator records every byte the reader asks for, and
//! each case requires the peak to stay within the bound the reference states: 128 bytes
//! for each value the manifest holds, one for each byte of its text, and, for a `WAI2`,
//! one [`RenditionMeta`] for each entry it lists, with 64 KiB for fixed costs. With the
//! value limit (2^22) and the length limit (2^27) that is at most 640 MiB for a `WAI1`
//! manifest, whatever its shape.
//!
//! The cases: a manifest at the length limit made of small values, refused for its
//! values before anything is built from it (the reader that parsed such a manifest into a
//! JSON tree took some 5 GB); one at the length limit within the value limit, which is
//! read; manifests of small values in each shape at the value limit, read, and one value
//! past it, refused; the largest manifest a conforming staged layer has, a `prev` with a
//! different producer for each of 2^20 units; and a `WAI2` listing as many renditions as
//! its table can hold, read, and one more, refused.
//!
//! The allocator counts what is asked of it, so the figures are the same on every 64-bit
//! target. A reallocation is counted as the new block allocated before the old one is
//! freed, which is the most it can hold at once. One test runs at a time (`ONE`), so the
//! count is this reader's alone.

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::time::Instant;

use wai::container::{
    check_manifest_text, RenditionMeta, Wai, WaiMulti, MANIFEST_MAX_LEN, MANIFEST_MAX_VALUES,
};

struct Counting;
static NOW: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

fn grow(n: usize) {
    let now = NOW.fetch_add(n, Relaxed) + n;
    PEAK.fetch_max(now, Relaxed);
}

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(l) };
        if !p.is_null() {
            grow(l.size());
        }
        p
    }
    unsafe fn alloc_zeroed(&self, l: Layout) -> *mut u8 {
        let p = unsafe { System.alloc_zeroed(l) };
        if !p.is_null() {
            grow(l.size());
        }
        p
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        unsafe { System.dealloc(p, l) };
        NOW.fetch_sub(l.size(), Relaxed);
    }
    unsafe fn realloc(&self, p: *mut u8, l: Layout, new: usize) -> *mut u8 {
        let q = unsafe { System.realloc(p, l, new) };
        if !q.is_null() {
            grow(new);
            NOW.fetch_sub(l.size(), Relaxed);
        }
        q
    }
}

#[global_allocator]
static COUNTING: Counting = Counting;

static ONE: Mutex<()> = Mutex::new(());

/// The most the reference may allocate reading a manifest of `len` bytes holding
/// `values` values and listing `renditions` `WAI2` entries (SPEC §2; SECURITY.md).
fn bound(len: usize, values: usize, renditions: usize) -> usize {
    128 * values + len + renditions * std::mem::size_of::<RenditionMeta>() + (1 << 16)
}

/// What `f` allocates at its peak beyond what was held when it started, with its result
/// (which the peak includes) and the seconds it took.
fn peak_of<T>(f: impl FnOnce() -> T) -> (T, usize, f64) {
    let base = NOW.load(Relaxed);
    PEAK.store(base, Relaxed);
    let t = Instant::now();
    let r = f();
    (r, PEAK.load(Relaxed) - base, t.elapsed().as_secs_f64())
}

fn envelope(magic: &[u8; 4], m: &str, rows: usize) -> Vec<u8> {
    let mut b = magic.to_vec();
    b.extend_from_slice(&(m.len() as u32).to_le_bytes());
    b.extend_from_slice(m.as_bytes());
    if magic == b"WAI2" {
        b.extend_from_slice(&(rows as u16).to_le_bytes());
        for _ in 0..rows {
            b.extend_from_slice(&[0; 8]);
        }
    } else {
        b.extend_from_slice(&0u32.to_le_bytes());
    }
    b
}

/// A `WAI1` manifest with a member no revision defines, `pad` (a reader ignores it, SPEC
/// §8), whose value is `open`, `n` elements written by `elem`, and `close`: with `pad`,
/// 10 values and the elements'.
const HEAD: &str = r#"{"wai":"1.0","media":"text","intent":"replicate","model_requirement":{"capability":"wai.text.zstd"},"conditioning":{"kind":"zstd"},"target":null,"pad":"#;
const FIXED: usize = 10;

fn padded(n: usize, open: &str, close: &str, mut elem: impl FnMut(usize, &mut String)) -> String {
    let mut m = String::with_capacity(HEAD.len() + 2 + n * 8);
    m.push_str(HEAD);
    m.push_str(open);
    for i in 0..n {
        if i > 0 {
            m.push(',');
        }
        elem(i, &mut m);
    }
    m.push_str(close);
    m.push('}');
    m
}

/// Read `m` as a `WAI1` envelope; the peak and time, and whether it was read.
fn read_wai1(m: &str) -> (Result<(), String>, usize, f64) {
    let env = envelope(b"WAI1", m, 0);
    peak_of(|| Wai::from_bytes(&env).map(|_| ()).map_err(|e| e.to_string()))
}

fn report(name: &str, len: usize, values: usize, peak: usize, secs: f64, bound: usize) {
    eprintln!(
        "{name}: {len} bytes, {values} values: peak {peak} bytes ({:.2} per byte, {:.1} per value) in {secs:.2} s; bound {bound}",
        peak as f64 / len as f64,
        peak as f64 / values as f64
    );
    assert!(peak <= bound, "{name}: the read allocated {peak} bytes, past the bound {bound}");
}

/// A manifest at the length limit made of small values is refused from its value count,
/// having built nothing: the first pass holds no value, and stops at the value past the
/// limit, some 8 MB into the text.
#[test]
fn a_manifest_of_small_values_at_the_length_limit_is_refused_before_it_is_built() {
    let _one = ONE.lock().unwrap_or_else(|e| e.into_inner());
    // Zeros, the last of them a number of ones long enough to bring the text to exactly
    // 2^27 bytes.
    let k = (MANIFEST_MAX_LEN - HEAD.len() - 3) / 2;
    let mut m = padded(k, "[", "]", |_, m| m.push('0'));
    let short = MANIFEST_MAX_LEN - m.len();
    let last = m.len() - 3;
    m.replace_range(last..last + 1, &"1".repeat(short + 1));
    assert_eq!(m.len(), MANIFEST_MAX_LEN);
    let (r, peak, secs) = read_wai1(&m);
    let e = r.expect_err("refused");
    assert!(e.starts_with("manifest_too_many_values:"), "{e}");
    report("2^27 bytes of zeros (refused)", m.len(), MANIFEST_MAX_VALUES + 1, peak, secs, 1 << 16);
}

/// A manifest at the length limit holding exactly the value limit is read within the
/// bound: 2^22 values, the most a reader must accept, each a string, to 2^27 bytes.
#[test]
fn a_manifest_at_both_limits_is_read_within_the_bound() {
    let _one = ONE.lock().unwrap_or_else(|e| e.into_inner());
    let n = MANIFEST_MAX_VALUES - FIXED;
    // Each string `"…"` and its comma: share out what is left of 2^27 bytes.
    let room = MANIFEST_MAX_LEN - HEAD.len() - 3 - (n - 1) - 2 * n;
    let (q, r) = (room / n, room % n);
    let m = padded(n, "[", "]", |i, m| {
        m.push('"');
        m.extend(std::iter::repeat_n('s', q + usize::from(i < r)));
        m.push('"');
    });
    assert_eq!(m.len(), MANIFEST_MAX_LEN);
    assert_eq!(check_manifest_text(m.as_bytes()).unwrap(), MANIFEST_MAX_VALUES);
    let (r, peak, secs) = read_wai1(&m);
    r.expect("read");
    report("2^27 bytes, 2^22 strings", m.len(), MANIFEST_MAX_VALUES, peak, secs, bound(m.len(), MANIFEST_MAX_VALUES, 0));
}

/// Manifests of small values, in every shape that costs a reader the most for each
/// value, at exactly the value limit, are read within the bound; one value more is
/// refused (`manifest_too_many_values`).
#[test]
fn small_values_at_the_value_limit_are_read_within_the_bound_and_one_more_is_refused() {
    let _one = ONE.lock().unwrap_or_else(|e| e.into_inner());
    let n = MANIFEST_MAX_VALUES - FIXED;
    type Shape = (&'static str, usize, &'static str, &'static str, fn(usize, &mut String));
    let shapes: [Shape; 6] = [
        ("an array of zeros", n, "[", "]", |_, m| m.push('0')),
        ("an array of empty arrays", n, "[", "]", |_, m| m.push_str("[]")),
        ("an array of empty objects", n, "[", "]", |_, m| m.push_str("{}")),
        ("an array of one-element arrays", n / 2, "[", "]", |_, m| m.push_str("[0]")),
        ("an array of one-member objects", n / 2, "[", "]", |_, m| m.push_str(r#"{"a":0}"#)),
        ("one object of many members", n, "{", "}", |i, m| {
            m.push('"');
            m.push_str(&format!("{i:x}"));
            m.push_str("\":0");
        }),
    ];
    for (name, k, open, close, elem) in shapes {
        // A shape of two values an element takes one zero more when `n` is odd.
        let mut m = padded(k, open, close, elem);
        if k != n && n % 2 == 1 {
            m.insert_str(m.len() - 2, ",0");
        }
        let values = check_manifest_text(m.as_bytes()).unwrap();
        assert_eq!(values, MANIFEST_MAX_VALUES, "{name}");
        let (r, peak, secs) = read_wai1(&m);
        r.expect(name);
        report(name, m.len(), values, peak, secs, bound(m.len(), values, 0));

        m.insert_str(m.len() - 2, if open == "{" { r#","x":0"# } else { ",0" });
        let (r, peak, _) = read_wai1(&m);
        let e = r.expect_err("one value past the limit is refused");
        assert!(e.starts_with("manifest_too_many_values:"), "{name}: {e}");
        // The first pass holds the names of an object's members to find a repeat, so the
        // refusal of one object of many members holds them; it stays within the bound.
        assert!(peak <= bound(m.len(), MANIFEST_MAX_VALUES, 0), "{name}: the refusal allocated {peak} bytes");
    }
}

/// The largest manifest a conforming staged layer has (SPEC §2): a `prev` that names a
/// different producer for each of the 2^20 units a `WIR1` layer may cover, about 94 MB
/// and 3·2^20 values, is read within the bound, and its `refines` reads well formed.
#[test]
fn the_largest_conforming_layer_manifest_is_read_within_the_bound() {
    let _one = ONE.lock().unwrap_or_else(|e| e.into_inner());
    let (a, b) = ("ab".repeat(32), "cd".repeat(32));
    let units = 1usize << 20;
    let mut m = String::with_capacity(95 << 20);
    m.push_str(r#"{"wai":"1.0","media":"image","intent":"replicate","model_requirement":{"capability":"wai.image.int_refine","fallback":null},"conditioning":{"kind":"int_refine"},"refines":{"base":""#);
    m.push_str(&a);
    m.push_str(r#"","prev":["#);
    for i in 0..units {
        if i > 0 {
            m.push(',');
        }
        m.push_str(r#"{"producer":""#);
        m.push_str(if i % 2 == 0 { &a } else { &b });
        m.push_str(r#"","count":1}"#);
    }
    m.push_str(r#"],"layer":2},"target":null}"#);
    let values = check_manifest_text(m.as_bytes()).unwrap();
    assert_eq!(values, 3 * units + 14);
    let env = envelope(b"WAI1", &m, 0);
    let (w, peak, secs) = peak_of(|| Wai::from_bytes(&env).expect("read"));
    report("a layer naming 2^20 producers", m.len(), values, peak, secs, bound(m.len(), values, 0));
    let r = w.refines().expect("refines").expect("well formed");
    assert_eq!(r.units(units as u32), (0, units as u32));
    r.check_units(units as u32).expect("the runs cover the units");
}

/// A `WAI2` manifest listing the 65 535 renditions its table can hold is read within the
/// bound, which counts the entries it builds; a `renditions` of 65 536 is refused before
/// its entries are built.
#[test]
fn a_wai2_listing_its_tables_rows_is_read_within_the_bound() {
    let _one = ONE.lock().unwrap_or_else(|e| e.into_inner());
    let head = r#"{"wai":"1.1","media":"text","intent":"replicate","renditions":["#;
    let manifest = |rows: usize| {
        let mut m = String::from(head);
        for i in 0..rows {
            if i > 0 {
                m.push(',');
            }
            m.push_str(r#"{"capability":"wai.text.zstd","kind":"zstd"}"#);
        }
        m.push_str(r#"],"target":null}"#);
        m
    };
    let rows = u16::MAX as usize;
    let m = manifest(rows);
    let values = check_manifest_text(m.as_bytes()).unwrap();
    let env = envelope(b"WAI2", &m, rows);
    let (w, peak, secs) = peak_of(|| WaiMulti::from_bytes(&env).map(|w| w.manifest.renditions.len()));
    assert_eq!(w.expect("read"), rows);
    report("a WAI2 of 65535 renditions", m.len(), values, peak, secs, bound(m.len(), values, rows));

    let m = manifest(rows + 1);
    let e = WaiMulti::from_bytes(&envelope(b"WAI2", &m, 0)).err().expect("refused").to_string();
    assert!(e.contains("more than the 65535 rows a WAI2 table holds"), "{e}");
}