diff --git a/crates/clawhdf5-format/src/filters_blosc2.rs b/crates/clawhdf5-format/src/filters_blosc2.rs index 5f332f8..d2b1ef0 100644 --- a/crates/clawhdf5-format/src/filters_blosc2.rs +++ b/crates/clawhdf5-format/src/filters_blosc2.rs @@ -451,7 +451,9 @@ struct Frame<'a> { offsets: Vec, } -fn parse_frame(buf: &[u8]) -> Result, FormatError> { +/// Parse a frame's header and decode its offsets chunk, holding no more +/// than an HDF5 chunk of `limit` bytes needs. +fn parse_frame(buf: &[u8], limit: usize) -> Result, FormatError> { if buf.len() < FRAME_HEADER_MINLEN { return Err(err("truncated frame header")); } @@ -497,29 +499,33 @@ fn parse_frame(buf: &[u8]) -> Result, FormatError> { .ok_or_else(|| err("chunks run past the frame"))?; let nbytes = usize::try_from(nbytes).map_err(|_| err("bad decoded size"))?; let chunksize = chunksize as usize; - // The offsets chunk follows the data chunks. + // The offsets chunk follows the data chunks: one `i64` per chunk. The + // frame header's sizes are the file's word, so they must not size it: + // it may be no larger than the HDF5 chunk (`limit`, at least 128 + // bytes), i.e. one Blosc2 chunk per 8 bytes of output. hdf5-blosc2 + // writes one chunk per frame; only python-blosc2 arrays of tiny + // chunks come near the cap. let off_src = &buf[data_end..]; - let expected = if nbytes == 0 { + let max_offsets = limit.max(128); + let off_len = if nbytes == 0 { 0 } else if chunksize > 0 { - nbytes.div_ceil(chunksize) + nbytes + .div_ceil(chunksize) + .checked_mul(8) + .filter(|&n| n <= max_offsets) + .ok_or_else(|| err("frame has more chunks than the HDF5 chunk can hold"))? } else { // Variable-size chunks: the offsets chunk tells how many. - usize::MAX - }; - let off_limit = if expected == usize::MAX { - 1 << 24 - } else { - expected - .checked_mul(8) - .ok_or_else(|| err("too many chunks"))? + max_offsets }; let offsets = if nbytes == 0 { Vec::new() } else { - blosc2_decompress_chunk(off_src, off_limit)? + blosc2_decompress_chunk(off_src, off_len)? }; - if !offsets.len().is_multiple_of(8) || (expected != usize::MAX && offsets.len() != off_limit) { + let expected = if chunksize > 0 { off_len } else { offsets.len() }; + if !offsets.len().is_multiple_of(8) || offsets.len() != expected { return Err(err("offsets chunk does not match the number of chunks")); } Ok(Frame { @@ -711,7 +717,7 @@ fn decode_frame( limit: usize, cd_shape: Option<&[usize]>, ) -> Result, FormatError> { - let frame = parse_frame(input)?; + let frame = parse_frame(input, limit)?; let meta = match frame.metalayer(b"b2nd")? { Some(m) => Some(m), None => frame.metalayer(b"caterva")?, @@ -1065,7 +1071,7 @@ mod tests { if w.is_err() { continue; } - let frame = parse_frame(&f).unwrap(); + let frame = parse_frame(&f, 1 << 20).unwrap(); let raw: [u8; 8] = frame.offsets[..8].try_into().unwrap(); let off = i64::from_le_bytes(raw); if off >= 0 { diff --git a/crates/clawhdf5-format/tests/blosc2_alloc_bounds.rs b/crates/clawhdf5-format/tests/blosc2_alloc_bounds.rs new file mode 100644 index 0000000..11e7812 --- /dev/null +++ b/crates/clawhdf5-format/tests/blosc2_alloc_bounds.rs @@ -0,0 +1,191 @@ +//! Crafted Blosc2 frames and chunks cannot make the decoder allocate out of +//! proportion to the HDF5 chunk it decodes. +//! +//! A frame's header, its offsets chunk and its chunk headers all declare +//! sizes, and the decoder used to allocate what they declared: a 173-byte +//! frame whose offsets chunk claimed 2 GiB was decoded in full before any +//! check failed. Every allocation is now bounded by the output limit (the +//! HDF5 chunk's size) and the input's length. +//! +//! Peak heap use is measured with a counting global allocator; the tests +//! share it, so each holds `SERIAL` for its whole run. +#![cfg(feature = "blosc2")] + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use clawhdf5_format::filters_blosc2::{blosc2_decompress, blosc2_decompress_chunk}; + +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(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)) +} + +/// What decoding one HDF5 chunk of `limit` bytes from `input` may hold at +/// once: the output, a few blocks of scratch (each no larger than the +/// output), the offsets table, and the Zstandard decoder's state. +fn bound(limit: usize, input: &[u8]) -> usize { + 6 * limit + 2 * input.len() + (1 << 20) +} + +fn lock() -> std::sync::MutexGuard<'static, ()> { + SERIAL.lock().unwrap_or_else(|e| e.into_inner()) +} + +/// A 32-byte (extended) Blosc2 chunk header. +fn chunk_header(ts: u8, nbytes: i32, blocksize: i32, cbytes: i32, special: u8) -> Vec { + let mut c = vec![5u8, 1, 0x05, ts]; + for v in [nbytes, blocksize, cbytes] { + c.extend_from_slice(&v.to_le_bytes()); + } + c.resize(32, 0); + c[31] = special << 4; + c +} + +/// A chunk of `nbytes` bytes that repeats one value (special type 3). +fn repeated(value: &[u8], nbytes: i32, blocksize: i32) -> Vec { + let mut c = chunk_header(value.len() as u8, nbytes, blocksize, 32 + value.len() as i32, 3); + c.extend_from_slice(value); + c +} + +/// A frame offset recording a special chunk of `kind` (1 zeros, 2 NaN). +fn special_offset(kind: u8) -> [u8; 8] { + (((0x80 | kind) as i64) << 56).to_le_bytes() +} + +/// A B2ND metalayer. +fn nd_meta(shape: &[i64], chunks: &[i32], blocks: &[i32]) -> Vec { + let n = shape.len() as u8; + let mut m = vec![0x95, 0, n, 0x90 | n]; + for s in shape { + m.push(0xd3); + m.extend_from_slice(&s.to_be_bytes()); + } + for dims in [chunks, blocks] { + m.push(0x90 | n); + for d in dims { + m.push(0xd2); + m.extend_from_slice(&d.to_be_bytes()); + } + } + m +} + +/// A contiguous frame: header (with a `b2nd` metalayer if given), the data +/// chunks, then the offsets chunk. +fn frame( + meta: Option<&[u8]>, + nbytes: i64, + typesize: i32, + chunksize: i32, + data: &[u8], + offsets: &[u8], +) -> Vec { + let mut h = vec![0u8; 91]; + h[0] = 0x9e; + h[1] = 0xa8; + h[2..10].copy_from_slice(b"b2frame\0"); + h[25] = 2; + match meta { + Some(m) => { + h.extend_from_slice(&[0xde, 0, 1, 0xa4]); + h.extend_from_slice(b"b2nd"); + let at = h.len() as i32 + 5; + h.push(0xd2); + h.extend_from_slice(&at.to_be_bytes()); + h.push(0xc6); + h.extend_from_slice(&(m.len() as u32).to_be_bytes()); + h.extend_from_slice(m); + } + None => h.extend_from_slice(&[0xde, 0, 0]), + } + let header_len = h.len() as i32; + h[11..15].copy_from_slice(&header_len.to_be_bytes()); + h[30..38].copy_from_slice(&nbytes.to_be_bytes()); + h[39..47].copy_from_slice(&(data.len() as i64).to_be_bytes()); + h[48..52].copy_from_slice(&typesize.to_be_bytes()); + h[58..62].copy_from_slice(&chunksize.to_be_bytes()); + h.extend_from_slice(data); + h.extend_from_slice(offsets); + let len = h.len() as u64; + h[16..24].copy_from_slice(&len.to_be_bytes()); + h +} + +/// The frame header's own sizes must not size the offsets chunk: a frame +/// declaring 32 Mi chunks of 4 bytes, whose offsets chunk (40 bytes) says +/// "one repeated offset, 256 MiB of them", made the decoder build all +/// 256 MiB of offsets for a 1 MiB HDF5 chunk and then return 4 bytes. +#[test] +fn offsets_chunk_is_bounded_by_the_output_limit() { + let _g = lock(); + let limit = 1 << 20; + let offsets_len: i32 = 256 << 20; + let nchunks = offsets_len as i64 / 8; + let offsets = repeated(&special_offset(1), offsets_len, 64 << 20); + let f = frame(None, nchunks * 4, 4, 4, &[], &offsets); + let (r, peak) = peak_during(|| blosc2_decompress(&f, limit)); + assert!(r.is_err(), "decoded {:?} bytes", r.map(|v| v.len())); + assert!( + peak <= bound(limit, &f), + "peak {peak} bytes for a {}-byte frame", + f.len() + ); + // The same frame with a variable chunk size (0): the offsets chunk + // alone says how many chunks there are. + let f = frame(None, nchunks * 4, 4, 0, &[], &offsets); + let (r, peak) = peak_during(|| blosc2_decompress(&f, limit)); + assert!(r.is_err()); + assert!(peak <= bound(limit, &f), "chunksize 0: peak {peak} bytes"); +} + +/// A legitimate frame of this shape (one chunk, its offset special) still +/// decodes. +#[test] +fn small_frames_still_decode() { + let _g = lock(); + let offsets = repeated(&special_offset(1), 8, 8); + let f = frame(None, 64, 4, 64, &[], &offsets); + assert_eq!(blosc2_decompress(&f, 64).unwrap(), vec![0; 64]); + let _ = blosc2_decompress_chunk; +}