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

//! Crafted ZFP filter parameters and streams cannot make the decoder panic
//! or allocate out of proportion to the chunk it decodes.
//!
//! The field size comes from the filter's `cd_values` (up to 2^48 values),
//! not from the chunk: the decoder allocates the output only when it matches
//! the chunk's size (or, when that is unknown, is within the 256 MiB
//! ceiling), and only when the stream is long enough to hold a bit per
//! block. Peak heap use is measured with a counting global allocator; the
//! tests share it, so each holds `SERIAL` for its whole run.
#![cfg(feature = "zfp")]

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

use clawhdf5_format::filters_zfp::zfp_decompress;

struct Counting;

static CURRENT: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);
static SERIAL: Mutex<()> = Mutex::new(());

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 alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let p = unsafe { System.alloc_zeroed(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).saturating_sub(base))
}

fn lock() -> std::sync::MutexGuard<'static, ()> {
    SERIAL.lock().unwrap_or_else(|e| e.into_inner())
}

/// What decoding may hold: the output (at most the limit, and at most a
/// 4-D block of doubles, 2 KiB, per bit of input), and a little more.
fn bound(max_output: usize, input: &[u8]) -> usize {
    let limit = if max_output == 0 {
        256 << 20
    } else {
        max_output
    };
    limit.min(input.len() * 8 * 2048) + 4096
}

/// An LSB-first bit writer.
#[derive(Default)]
struct Bits {
    v: Vec<u8>,
    n: usize,
}

impl Bits {
    fn put(&mut self, x: u64, bits: usize) {
        for i in 0..bits {
            if self.n.is_multiple_of(8) {
                self.v.push(0);
            }
            if (x >> i) & 1 == 1 {
                *self.v.last_mut().unwrap() |= 1 << (self.n % 8);
            }
            self.n += 1;
        }
    }
}

/// H5Z-ZFP `cd_values`: a version word and a zfp header for a field of
/// `ztype` (0 int32, 1 int64, 2 float, 3 double) and sizes `n` (fastest
/// first), with `mode` (12 bits, or 64 when `long`).
fn cd_values(ztype: u64, n: &[u64], mode: u64, long: bool) -> Vec<u32> {
    let mut b = Bits::default();
    for c in b"zfp" {
        b.put(*c as u64, 8);
    }
    b.put(5, 8);
    let dims = n.len();
    let mut meta = 0u64;
    let bits = [48, 24, 16, 12][dims - 1];
    for &v in n.iter().rev() {
        meta = (meta << bits) + v - 1;
    }
    meta = (meta << 2) + dims as u64 - 1;
    meta = (meta << 2) + ztype;
    b.put(meta, 52);
    b.put(mode, if long { 64 } else { 12 });
    b.v.resize(b.v.len().div_ceil(4) * 4, 0);
    let mut cd = vec![0x1001_1111u32];
    cd.extend(
        b.v.chunks(4)
            .map(|w| u32::from_le_bytes(w.try_into().unwrap())),
    );
    cd
}

fn elem(ztype: u64) -> usize {
    if ztype & 1 == 0 { 4 } else { 8 }
}

/// A 1-D field of 2^32 doubles (32 GiB) in a 1-byte chunk: refused for
/// its size, with or without the chunk size known, before anything is
/// allocated.
#[test]
fn huge_fields_are_refused_without_allocating() {
    let _g = lock();
    for (n, ztype) in [
        (vec![1u64 << 32], 3),
        (vec![1 << 24, 1 << 24], 3),
        (vec![4096; 4], 1),
    ] {
        let cd = cd_values(ztype, &n, 2176, false);
        for max_output in [0usize, 1 << 20] {
            let (r, peak) = peak_during(|| zfp_decompress(&[0xff], &cd, max_output));
            assert!(r.is_err(), "{n:?}: decoded {:?} bytes", r.map(|v| v.len()));
            assert!(peak < 4096, "{n:?}: peak {peak} bytes");
        }
    }
}

/// A field of the chunk's size whose stream is too short for its blocks
/// is refused before the output is allocated.
#[test]
fn short_streams_are_refused_before_allocating() {
    let _g = lock();
    let n = [1u64 << 18];
    let cd = cd_values(2, &n, 2176, false);
    let size = (1 << 18) * 4;
    let (r, peak) = peak_during(|| zfp_decompress(&[0u8; 100], &cd, size));
    assert!(r.is_err());
    assert!(peak < 4096, "peak {peak} bytes");
    // A stream with a bit per block: all-zero blocks, which decode.
    let input = vec![0u8; (1 << 16) / 8];
    let out = zfp_decompress(&input, &cd, size).unwrap();
    assert_eq!(out, vec![0u8; size]);
}

/// xorshift64*: deterministic, so a failure reproduces.
struct Rng(u64);

impl Rng {
    fn next(&mut self) -> u64 {
        let mut x = self.0;
        x ^= x >> 12;
        x ^= x << 25;
        x ^= x >> 27;
        self.0 = x;
        x.wrapping_mul(0x2545_F491_4F6C_DD1D)
    }

    fn below(&mut self, n: u64) -> u64 {
        self.next() % n.max(1)
    }
}

/// A mode word: one of the four short forms, or the 64-bit expert form
/// with parameters at and past their edges (maxbits below a float block's
/// exponent, minbits past maxbits, precision 0, minexp at the reversible
/// boundary).
fn mode(rng: &mut Rng) -> (u64, bool) {
    match rng.below(6) {
        0 => (rng.below(2048), false),
        1 => (2048 + rng.below(128), false),
        2 => (2176, false),
        3 => (2177 + rng.below(4094 - 2177 + 1), false),
        _ => {
            fn pick(rng: &mut Rng, v: [u64; 6]) -> u64 {
                v[rng.below(6) as usize]
            }
            let r = [rng.below(2000), rng.below(0x8000), rng.below(40)];
            let minbits = pick(rng, [0, 1, 11, r[0], r[1], 1]);
            let r = [rng.below(0x8000), rng.below(40)];
            let maxbits = pick(rng, [minbits, minbits + r[1], 7, 11, r[0], 0x7fff]);
            let maxprec = rng.below(0x80);
            let r = [rng.below(0x8000), rng.below(400)];
            let minexp = pick(
                rng,
                [
                    r[0],
                    16495 - 1074,
                    16495 - 1075,
                    16495 + r[1] - 200,
                    16495 - 1074,
                    r[0],
                ],
            );
            let m = ((((minexp << 7) + maxprec) << 15) + maxbits) << 15;
            ((m + minbits) << 12 | 0xfff, true)
        }
    }
}

/// Random fields, modes and streams (random bytes, runs of ones, mostly
/// zeros; of every length), decoded with the chunk size known and not,
/// and with header words that are random too.
#[test]
fn fuzzed_headers_and_streams_stay_within_the_allocation_bound() {
    let _g = lock();
    let mut rng = Rng(0x2f9);
    let mut decoded = 0;
    for i in 0..20_000 {
        let ztype = rng.below(4);
        let dims = 1 + rng.below(4) as usize;
        let max = [300, 40, 14, 8][dims - 1];
        let n: Vec<u64> = (0..dims).map(|_| 1 + rng.below(max)).collect();
        let (m, long) = mode(&mut rng);
        let mut cd = cd_values(ztype, &n, m, long);
        if rng.below(10) == 0 {
            let at = rng.below(cd.len() as u64) as usize;
            cd[at] ^= 1 << rng.below(32);
        }
        if rng.below(20) == 0 {
            cd.truncate(rng.below(cd.len() as u64 + 1) as usize);
        }
        let len = match rng.below(4) {
            0 => rng.below(8),
            1 => rng.below(300),
            _ => rng.below(20_000),
        } as usize;
        let input: Vec<u8> = match rng.below(3) {
            0 => (0..len).map(|_| rng.next() as u8).collect(),
            1 => (0..len)
                .map(|_| [0, 0xff, rng.next() as u8][rng.below(3) as usize])
                .collect(),
            _ => (0..len)
                .map(|_| [0, 0, 0, 1, 0x80, rng.next() as u8][rng.below(6) as usize])
                .collect(),
        };
        let size = n.iter().product::<u64>() as usize * elem(ztype);
        let max_output = if rng.below(4) == 0 { 0 } else { size };
        let (r, peak) = peak_during(|| zfp_decompress(&input, &cd, max_output));
        if let Ok(out) = &r {
            decoded += 1;
            if max_output != 0 {
                assert_eq!(out.len(), max_output, "iteration {i}");
            }
        }
        assert!(
            peak <= bound(max_output, &input),
            "iteration {i}: peak {peak} bytes for {n:?} from {} bytes ({:?})",
            input.len(),
            r.map(|v| v.len())
        );
    }
    // Most inputs are streams zfp decodes without running out.
    assert!(decoded > 5_000, "only {decoded} decoded");
}