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

//! Decoding a `wai.text.zstd` or `wai.text.xz` payload that would expand past the
//! length limit takes no more output buffer than the limit (SPEC §2, *Reader limits*).
//! Measured with a global allocator that counts the bytes live and their peak.
//!
//! Each payload here is a few kilobytes and decodes to four times the limit the call is
//! given. The call must refuse it, and the bytes it allocated on the way must stay
//! within the limit and a small allowance for its read buffer. A decode that grew its
//! buffer past the limit before checking, or that decoded everything and then measured
//! it, would not.
//!
//! The counter sees what this crate allocates: the output buffer and the read buffer.
//! The decoding library's own working memory is allocated by the library and is bounded
//! by the window limit, which `text_decode_bounds.rs` tests.
//!
//! Its own test binary with a single test, so that no other test's allocations land in
//! the windows it measures. Requires the `codecs` feature.
#![cfg(feature = "codecs")]

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

use wai::codecs::text::{xz_decode_within, zstd_decode_within};

/// The system allocator, counting the bytes live and their peak.
struct Count;
static CUR: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

fn grow(n: usize) {
    PEAK.fetch_max(CUR.fetch_add(n, Relaxed) + n, Relaxed);
}

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

#[global_allocator]
static ALLOC: Count = Count;

/// The most bytes live at once while `f` runs, above what was live when it started.
fn peak_of<T>(f: impl FnOnce() -> T) -> (T, usize) {
    let base = CUR.load(Relaxed);
    PEAK.store(base, Relaxed);
    let r = f();
    (r, PEAK.load(Relaxed).saturating_sub(base))
}

const LIMIT: usize = 8 << 20;
/// The read buffer, and room for the error's text.
const ALLOWANCE: usize = 1 << 20;

fn zeros_through<W: Write>(mut w: W, n: usize) -> W {
    let piece = vec![0u8; 1 << 20];
    let mut left = n;
    while left > 0 {
        let k = left.min(piece.len());
        w.write_all(&piece[..k]).unwrap();
        left -= k;
    }
    w
}

#[test]
fn a_payload_that_expands_past_the_limit_takes_no_more_buffer_than_the_limit() {
    let zstd_payload = zeros_through(zstd::stream::write::Encoder::new(Vec::new(), 1).unwrap(), 4 * LIMIT)
        .finish()
        .unwrap();
    let xz_payload = zeros_through(xz2::write::XzEncoder::new(Vec::new(), 0), 4 * LIMIT).finish().unwrap();
    assert!(zstd_payload.len() < LIMIT / 64 && xz_payload.len() < LIMIT / 64, "the payloads are small");

    let (r, peak) = peak_of(|| zstd_decode_within(&zstd_payload, LIMIT));
    let why = r.expect_err("zstd: the payload must be refused").0;
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");
    assert!(peak <= LIMIT + ALLOWANCE, "zstd: refusing took {peak} bytes of buffer, limit {LIMIT}");
    assert!(peak >= LIMIT / 2, "zstd: the refusal came before the output neared the limit ({peak} bytes): is the payload what it should be?");

    let (r, peak) = peak_of(|| xz_decode_within(&xz_payload, LIMIT));
    let why = r.expect_err("xz: the payload must be refused").0;
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");
    assert!(peak <= LIMIT + ALLOWANCE, "xz: refusing took {peak} bytes of buffer, limit {LIMIT}");
    assert!(peak >= LIMIT / 2, "xz: the refusal came before the output neared the limit ({peak} bytes): is the payload what it should be?");

    // At the limit the decode succeeds, in a buffer of the limit.
    let zstd_at = zeros_through(zstd::stream::write::Encoder::new(Vec::new(), 1).unwrap(), LIMIT).finish().unwrap();
    let (r, peak) = peak_of(|| zstd_decode_within(&zstd_at, LIMIT));
    assert_eq!(r.unwrap().len(), LIMIT);
    assert!(peak <= LIMIT + ALLOWANCE, "zstd: decoding to the limit took {peak} bytes of buffer, limit {LIMIT}");
}