clawhdf5/crates/clawhdf5-format/tests/vl_heap_bounds.rs

//! Crafted files cannot make variable-length reads retain memory, or take
//! time, out of proportion to the file.
//!
//! `VlResolver` used to keep an owned copy of every object of every
//! collection it parsed, for the whole read. A file whose global heap
//! collections nest inside each other's object data — each element
//! pointing at a different one — then made retained memory O(K × file
//! size): a 744 KB file took 1.58 GB. The same nesting, with every
//! collection's object chain jumping to one shared run of tiny objects,
//! made the parse time O(K × M) as well. libhdf5 never writes overlapping
//! collections; they are now refused, and the cache holds only where
//! objects lie.
//!
//! Peak heap use is measured with a counting global allocator, so the
//! cases run one after another in a single test.

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

use clawhdf5_format::vl_data::VlResolver;

struct Counting;

static CURRENT: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(layout) };
        if !p.is_null() {
            let now = CURRENT.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
            PEAK.fetch_max(now, Ordering::Relaxed);
        }
        p
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) };
        CURRENT.fetch_sub(layout.size(), Ordering::Relaxed);
    }
}

#[global_allocator]
static ALLOC: Counting = Counting;

/// Bytes allocated at the peak of `f`, above what was live when it started.
fn peak_during<T>(f: impl FnOnce() -> T) -> (T, usize) {
    let base = CURRENT.load(Ordering::Relaxed);
    PEAK.store(base, Ordering::Relaxed);
    let out = f();
    (out, PEAK.load(Ordering::Relaxed) - base)
}

fn put_header(file: &mut [u8], at: usize, size: u64) {
    file[at..at + 4].copy_from_slice(b"GCOL");
    file[at + 4] = 1;
    file[at + 8..at + 16].copy_from_slice(&size.to_le_bytes());
}

fn put_object(file: &mut [u8], at: usize, index: u16, size: u64) {
    file[at..at + 2].copy_from_slice(&index.to_le_bytes());
    file[at + 2..at + 4].copy_from_slice(&1u16.to_le_bytes());
    file[at + 8..at + 16].copy_from_slice(&size.to_le_bytes());
}

fn element(length: u32, addr: u64, index: u32) -> Vec<u8> {
    let mut e = length.to_le_bytes().to_vec();
    e.extend_from_slice(&addr.to_le_bytes());
    e.extend_from_slice(&index.to_le_bytes());
    e
}

/// K collections 32 bytes apart, each running to the end of the file with
/// one object covering the rest of it (and so every later collection).
/// Element i is that object of collection i.
fn nested(k: usize) -> (Vec<u8>, Vec<u8>) {
    let base = 64;
    let end = base + 32 * k + 64;
    let mut file = vec![0u8; end];
    let mut raw = Vec::new();
    for i in 0..k {
        let at = base + 32 * i;
        put_header(&mut file, at, (end - at) as u64);
        let obj = (end - at - 32) as u64;
        put_object(&mut file, at + 16, 1, obj);
        raw.extend(element(obj as u32, at as u64, 1));
    }
    (file, raw)
}

/// K collections 32 bytes apart, each with a first object that jumps over
/// the later collections to one shared run of M empty objects, so parsing
/// every collection walks all M.
fn shared_tail(k: usize, m: usize) -> (Vec<u8>, Vec<u8>) {
    let base = 64;
    let tail = base + 32 * k + 32;
    let end = tail + 16 * m + 16;
    let mut file = vec![0u8; end];
    let mut raw = Vec::new();
    for i in 0..k {
        let at = base + 32 * i;
        put_header(&mut file, at, (end - at) as u64);
        let jump = (tail - at - 32) as u64;
        put_object(&mut file, at + 16, 1, jump);
        raw.extend(element(jump as u32, at as u64, 1));
    }
    for j in 0..m {
        put_object(&mut file, tail + 16 * j, (j % 65_000 + 2) as u16, 0);
    }
    (file, raw)
}

#[test]
fn overlapping_collections_are_refused_in_bounded_memory_and_time() {
    for (name, (file, raw)) in [
        ("nested", nested(2000)),
        ("shared tail", shared_tail(500, 10_000)),
    ] {
        let start = Instant::now();
        let (result, peak) = peak_during(|| {
            let mut r = VlResolver::new(&file, 8, 8);
            (r.string_bytes(&raw), r.sequences(&raw, 1).map(|s| s.len()))
        });
        let took = start.elapsed();
        // libhdf5 never writes overlapping collections, and refuses these
        // files; so do we, rather than returning what they claim.
        let (strings, sequences) = result;
        let e = strings.expect_err(name).to_string();
        assert!(e.contains("overlaps"), "{name}: {e}");
        assert!(sequences.is_err(), "{name}");
        // Measured before the fix: 129 MB ("nested", 64 KB file) and 350 MB
        // ("shared tail", 176 KB file) live at the peak; after, 97 KB and
        // 0.9 MB.
        assert!(
            peak < 4 * file.len() + (1 << 20),
            "{name}: peak {peak} bytes for a {}-byte file",
            file.len()
        );
        assert!(took < Duration::from_secs(5), "{name}: took {took:?}");
    }
}

/// Collections that do not overlap still read, however many elements point
/// into them, and the first object of a collection is returned for its
/// index (as before).
#[test]
fn separate_collections_still_read() {
    let mut file = vec![0u8; 64 + 3 * 64];
    let mut raw = Vec::new();
    for i in 0..3usize {
        let at = 64 + 64 * i;
        put_header(&mut file, at, 64);
        put_object(&mut file, at + 16, 1, 3);
        file[at + 32..at + 35].copy_from_slice(format!("s{i}!").as_bytes());
        raw.extend(element(3, at as u64, 1));
    }
    raw.extend(element(3, 64, 1));
    let mut r = VlResolver::new(&file, 8, 8);
    assert_eq!(r.strings(&raw).unwrap(), ["s0!", "s1!", "s2!", "s0!"]);
}