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

//! A container's reader reserves no more entries than its bytes could hold.
//!
//! Each extension container (world, audio scene, reel, score, haptic, splat4d) reads a
//! `u32` count of entries from its input and used to reserve that many before reading
//! one: a few bytes could ask the allocator for gigabytes. The coverage-guided fuzzer
//! (`fuzz/`, target `containers`) found it in the audio scene's object table, asking
//! for 48 GiB.
//!
//! For each container a committed vector is rewritten so that the section holding a
//! count declares 2^32 - 1 entries and holds nothing after the count. The reader must
//! refuse it having asked for no single allocation past a few megabytes. The allocator
//! here records the largest request, and refuses any past 256 MiB, so a reader that
//! still sizes its reservation from the count fails here instead of reserving.
#![cfg(all(
    feature = "world",
    feature = "audio_scene",
    feature = "film",
    feature = "score",
    feature = "haptic",
    feature = "splat4d"
))]

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

/// The system allocator, recording the largest single request.
struct Largest;
static LARGEST: AtomicUsize = AtomicUsize::new(0);
const REFUSE_PAST: usize = 256 << 20;

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

#[global_allocator]
static A: Largest = Largest;

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

/// The sections of a container in the layout every one of these shares: a magic, a
/// `u16` count, a table of `(kind u8, offset u32, length u32)`, then the bodies.
fn sections(b: &[u8]) -> Vec<(u8, Vec<u8>)> {
    let n = u16::from_le_bytes([b[4], b[5]]) as usize;
    let body = &b[6 + 9 * n..];
    (0..n)
        .map(|i| {
            let e = &b[6 + 9 * i..6 + 9 * (i + 1)];
            let off = u32::from_le_bytes(e[1..5].try_into().unwrap()) as usize;
            let len = u32::from_le_bytes(e[5..9].try_into().unwrap()) as usize;
            (e[0], body[off..off + len].to_vec())
        })
        .collect()
}

/// The container rebuilt with section `kind`'s body replaced by a count of 2^32 - 1.
fn declaring_too_many(b: &[u8], kind: u8) -> Vec<u8> {
    let mut secs = sections(b);
    let s = secs.iter_mut().find(|(k, _)| *k == kind).expect("the vector has the section");
    s.1 = u32::MAX.to_le_bytes().to_vec();
    let mut out = b[..4].to_vec();
    out.extend_from_slice(&(secs.len() as u16).to_le_bytes());
    let mut off = 0u32;
    for (k, body) in &secs {
        out.push(*k);
        out.extend_from_slice(&off.to_le_bytes());
        out.extend_from_slice(&(body.len() as u32).to_le_bytes());
        off += body.len().div_ceil(8) as u32 * 8;
    }
    for (_, body) in &secs {
        out.extend_from_slice(body);
        out.resize(out.len() + (body.len().div_ceil(8) * 8 - body.len()), 0);
    }
    out
}

#[test]
fn each_container_refuses_a_count_its_bytes_cannot_hold_without_reserving_it() {
    let cases: [(&str, &str, u8, fn(&[u8]) -> bool); 7] = [
        ("world assets", "wai-web/demo/vectors/gallery.wwld", 0x03, |b| wai::world::WorldContainer::from_bytes(b).is_err()),
        ("audio scene objects", "wai-web/demo/audio/moving.wscn", 0x02, |b| wai::audio_scene::SpatialScene::from_bytes(b).is_err()),
        ("audio scene assets", "wai-web/demo/audio/moving.wscn", 0x03, |b| wai::audio_scene::SpatialScene::from_bytes(b).is_err()),
        ("reel shots", "film-conformance/vectors/coverage.wrel", 0x03, |b| wai::film::Reel::from_bytes(b).is_err()),
        ("score notes", "score-conformance/vectors/arp.wscr", 0x02, |b| wai::score::Score::from_bytes(b).is_err()),
        ("haptic channels", "haptic-conformance/vectors/ramp.whap", 0x02, |b| wai::haptic::HapticSignal::from_bytes(b).is_err()),
        ("splat4d tracks", "volumetric-conformance/vectors/morph.w4ds", 0x02, |b| wai::splat4d::Splat4dClipData::from_bytes(b).is_err()),
    ];
    for (what, rel, kind, refused) in cases {
        let original = vector(rel);
        let bytes = if sections(&original).iter().any(|(k, _)| *k == kind) {
            declaring_too_many(&original, kind)
        } else {
            // The vector carries no such section (the audio scene's assets are optional):
            // add one, as the last section.
            let mut with = original.clone();
            let n = u16::from_le_bytes([with[4], with[5]]);
            with[4..6].copy_from_slice(&(n + 1).to_le_bytes());
            let secs = sections(&original);
            let end: usize = secs.iter().map(|(_, b)| b.len().div_ceil(8) * 8).sum();
            let mut entry = vec![kind];
            entry.extend_from_slice(&(end as u32).to_le_bytes());
            entry.extend_from_slice(&4u32.to_le_bytes());
            let table_end = 6 + 9 * n as usize;
            with.splice(table_end..table_end, entry);
            let body_start = table_end + 9;
            with.truncate(body_start + end);
            with.resize(body_start + end, 0);
            with.extend_from_slice(&u32::MAX.to_le_bytes());
            with
        };
        LARGEST.store(0, Relaxed);
        assert!(refused(&bytes), "{what}: 2^32 - 1 entries in four bytes were accepted");
        let largest = LARGEST.load(Relaxed);
        assert!(largest <= 8 << 20, "{what}: the reader asked for {largest} bytes at once");
    }
}