mikronova/openzl-sys-rs

tests/alloc_budget.rs

//! The allocation budget (`with_alloc_budget`): a hostile frame cannot make
//! OpenZL allocate past it, an honest frame decodes under it, and a thread
//! without one is unchanged.

use openzl_sys_rs::*;

/// The heimdall fuzz frame with its element count (10-byte varint at 0x11) set to `n`.
fn constant_frame_with_count(n: u64) -> Vec<u8> {
    let mut frame = include_bytes!("fixtures/regress-constant-stream-alloc.zl").to_vec();
    let mut v = n;
    for i in 0..10 {
        frame[0x11 + i] = (v & 0x7f) as u8 | if i < 9 { 0x80 } else { 0 };
        v >>= 7;
    }
    frame
}

fn honest_frame(len: usize) -> (Vec<u8>, Vec<u8>) {
    let src: Vec<u8> = (0..len as u32).map(|x| ((x * 31 + x / 7) % 253) as u8).collect();
    let mut c = ZlCCtx::new().unwrap();
    c.set_parameter(ZL_CParam_ZL_CParam_formatVersion, 25).unwrap();
    c.set_parameter(ZL_CParam_ZL_CParam_compressionLevel, 3).unwrap();
    let mut frame = vec![0u8; zl_compress_bound(src.len())];
    let n = c.compress(&mut frame, &src).unwrap();
    frame.truncate(n);
    (src, frame)
}

#[test]
fn a_hostile_frame_is_refused_at_the_budget() {
    // 2^24 elements: big enough to be refused by a 1 MiB budget, small enough
    // that the test is harmless if the budget is not enforced.
    let frame = constant_frame_with_count(1 << 24);
    let mut out = [0u8; 16];
    let (got, spent) = with_alloc_budget(1 << 20, || zl_decompress(&mut out, &frame));
    assert!(got.is_err(), "the frame is refused");
    assert!(spent.exceeded, "the refusal was the budget's: {spent:?}");
    assert!(spent.used <= spent.limit, "nothing was charged past the limit: {spent:?}");
}

#[test]
fn an_honest_frame_decodes_under_a_budget_and_is_charged() {
    let (src, frame) = honest_frame(1 << 20);
    let mut out = vec![0u8; src.len()];
    let (got, spent) = with_alloc_budget(64 << 20, || zl_decompress(&mut out, &frame));
    assert_eq!(got.unwrap(), src.len());
    assert_eq!(out, src);
    assert!(!spent.exceeded, "{spent:?}");
    assert!(spent.used > 0, "OpenZL's allocations are charged: {spent:?}");
}

#[test]
fn a_thread_without_a_budget_is_unchanged() {
    let (src, frame) = honest_frame(1 << 20);
    let mut out = vec![0u8; src.len()];
    // Spend a budget first: it must not linger on the thread.
    let _ = with_alloc_budget(1, || ());
    assert_eq!(zl_decompress(&mut out, &frame).unwrap(), src.len());
    assert_eq!(out, src);
}

#[test]
fn budgets_nest_and_restore() {
    let (src, frame) = honest_frame(1 << 20);
    let hostile = constant_frame_with_count(1 << 24);
    let mut out = vec![0u8; src.len()];
    let ((inner, ok), outer) = with_alloc_budget(64 << 20, || {
        let mut small = [0u8; 16];
        let (inner, _) = with_alloc_budget(1 << 20, || zl_decompress(&mut small, &hostile));
        (inner, zl_decompress(&mut out, &frame))
    });
    assert!(inner.is_err());
    assert_eq!(ok.unwrap(), src.len());
    assert!(!outer.exceeded, "the inner refusal is not the outer budget's: {outer:?}");
    assert!(outer.used > 0, "the outer budget is back in force after the inner one: {outer:?}");
}

#[test]
fn the_budget_is_restored_on_unwind() {
    let hostile = constant_frame_with_count(1 << 24);
    let r = std::panic::catch_unwind(|| with_alloc_budget(1 << 20, || panic!("boom")));
    assert!(r.is_err());
    // No budget is left behind: the frame fails on its own terms, not the budget's.
    let mut out = [0u8; 16];
    let (_, spent) = with_alloc_budget(usize::MAX, || zl_decompress(&mut out, &hostile));
    assert!(!spent.exceeded);
}