diff --git a/crates/clawhdf5-format/src/filters.rs b/crates/clawhdf5-format/src/filters.rs index 31c024a..ac17725 100644 --- a/crates/clawhdf5-format/src/filters.rs +++ b/crates/clawhdf5-format/src/filters.rs @@ -31,6 +31,26 @@ pub fn decompress_chunk( decompress_chunk_masked(compressed, pipeline, chunk_size, element_size, 0) } +/// Upper bound on the output of filter `filter_id` applied (in the write +/// direction) to `input` bytes. 0 means "unknown" and stays unknown. +/// +/// Shuffle preserves the size and Fletcher32 appends a 4-byte checksum. Any +/// other filter is a codec whose output can exceed its input on +/// incompressible data (deflate's stored blocks, LZ4's and zstd's literal +/// runs, codec headers); `n + n/8 + 64` covers every supported codec's worst +/// case while still bounding a decompression bomb to a small multiple of the +/// chunk. +fn filter_output_bound(filter_id: u16, input: usize) -> usize { + if input == 0 { + return 0; + } + match filter_id { + FILTER_SHUFFLE => input, + FILTER_FLETCHER32 => input.saturating_add(4), + _ => input.saturating_add(input / 8).saturating_add(64), + } +} + /// Whether bit `index` of a chunk's filter mask says filter `index` was /// skipped when the chunk was written. fn filter_skipped(filter_mask: u32, index: usize) -> bool { @@ -48,8 +68,12 @@ pub fn all_filters_skipped(pipeline: &FilterPipeline, filter_mask: u32) -> bool /// optional filter that declined, or a direct chunk write), so only that /// filter is skipped here; the others are still undone, in reverse order. /// -/// `chunk_size` is the chunk's decoded size (0 if unknown); it caps each -/// stage's output. +/// `chunk_size` is the chunk's decoded size (0 if unknown). Each stage's +/// output is capped at what the filters before it (in write order) can have +/// produced from `chunk_size` bytes — e.g. a Fletcher32 checksum placed +/// before deflate (NetCDF-4's ordering) makes deflate's output 4 bytes +/// larger than the chunk — so the decompression-bomb limit stays tight +/// without rejecting valid pipelines. pub fn decompress_chunk_masked( compressed: &[u8], pipeline: &FilterPipeline, @@ -57,12 +81,23 @@ pub fn decompress_chunk_masked( element_size: u32, filter_mask: u32, ) -> Result, FormatError> { + // bounds[i]: the most bytes that entered filter i on the write side, and + // so the most that undoing filter i may produce. + let mut bounds = Vec::with_capacity(pipeline.filters.len()); + let mut size = chunk_size; + for (i, filter) in pipeline.filters.iter().enumerate() { + bounds.push(size); + if !filter_skipped(filter_mask, i) { + size = filter_output_bound(filter.filter_id, size); + } + } + let mut data = compressed.to_vec(); for (i, filter) in pipeline.filters.iter().enumerate().rev() { if filter_skipped(filter_mask, i) { continue; } - let bound = chunk_size; + let bound = bounds[i]; data = match filter.filter_id { FILTER_SHUFFLE => shuffle_decompress(&data, element_size as usize)?, // `bound` caps the decoded size so these decoders can't be forced @@ -944,13 +979,13 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result, Forma if element_size <= 1 { return Ok(data.to_vec()); } - if !data.len().is_multiple_of(element_size) { - return Err(FormatError::FilterError( - "shuffle: data length not a multiple of element size".into(), - )); - } + // Like libhdf5, only whole elements are shuffled; trailing bytes (e.g. a + // Fletcher32 checksum appended before the shuffle) are stored as-is. + let whole = data.len() - data.len() % element_size; + let (data, tail) = data.split_at(whole); let num_elements = data.len() / element_size; - let mut result = vec![0u8; data.len()]; + let mut result = vec![0u8; whole]; + result.reserve_exact(tail.len()); // The shuffled stream is `element_size` byte planes of `num_elements` // bytes each; un-shuffling interleaves them. This is on the read path of @@ -981,6 +1016,7 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result, Forma } } } + result.extend_from_slice(tail); Ok(result) } @@ -996,19 +1032,19 @@ fn shuffle_compress(data: &[u8], element_size: usize) -> Result, FormatE if element_size <= 1 { return Ok(data.to_vec()); } - if !data.len().is_multiple_of(element_size) { - return Err(FormatError::FilterError( - "shuffle: data length not a multiple of element size".into(), - )); - } + // Trailing bytes that don't make a whole element are left in place, as + // libhdf5 does. + let whole = data.len() - data.len() % element_size; + let (data, tail) = data.split_at(whole); let num_elements = data.len() / element_size; - let mut result = vec![0u8; data.len()]; + let mut result = vec![0u8; whole]; match element_size { 4 => shuffle_compress_4(data, num_elements, &mut result), 8 => shuffle_compress_general(data, num_elements, element_size, &mut result), _ => shuffle_compress_general(data, num_elements, element_size, &mut result), } + result.extend_from_slice(tail); Ok(result) } @@ -1503,6 +1539,51 @@ mod tests { ); } + #[test] + #[cfg(feature = "deflate")] + fn fletcher32_ahead_of_deflate_stays_bounded() { + // NetCDF-4 order: the checksum is appended before shuffle and deflate, + // so deflate decodes chunk + 4 bytes. + let pipeline = FilterPipeline { + version: 2, + filters: vec![ + filter(FILTER_FLETCHER32), + filter(FILTER_SHUFFLE), + filter(FILTER_DEFLATE), + ], + }; + let data: Vec = (0..400).map(|i| (i * 13 % 251) as u8).collect(); + let n = data.len(); + let stored = compress_chunk(&data, &pipeline, 8).unwrap(); + assert_eq!(decompress_chunk(&stored, &pipeline, n, 8).unwrap(), data); + + // The cap still bites: a stream that inflates past chunk + 4 bytes + // is rejected rather than allocated. + let only_deflate = FilterPipeline { + version: 2, + filters: vec![filter(FILTER_DEFLATE)], + }; + let oversized = compress_chunk(&vec![0u8; n + 5], &only_deflate, 8).unwrap(); + let err = decompress_chunk(&oversized, &pipeline, n, 8).unwrap_err(); + assert!( + matches!(err, FormatError::DecompressionError(_)), + "expected a size-limit error, got {err:?}" + ); + // A bomb is still stopped near the chunk size. + let bomb = compress_chunk(&vec![0u8; 64 * n], &only_deflate, 8).unwrap(); + assert!(decompress_chunk(&bomb, &pipeline, n, 8).is_err()); + } + + #[test] + fn shuffle_leaves_a_partial_trailing_element_in_place() { + // libhdf5 shuffles whole elements and copies the remainder as-is. + let data: Vec = (0..20).collect(); + let shuffled = shuffle_compress(&data, 8).unwrap(); + assert_eq!(&shuffled[16..], &data[16..]); + assert_eq!(&shuffled[..4], &[0, 8, 1, 9]); + assert_eq!(shuffle_decompress(&shuffled, 8).unwrap(), data); + } + #[test] #[cfg(feature = "deflate")] fn pipeline_compress_decompress_roundtrip() { diff --git a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs index 58fe290..deda24f 100644 --- a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs +++ b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs @@ -173,3 +173,67 @@ with h5py.File("{p}", "w") as f: let want: Vec = (1006..1026i32).flat_map(i32::to_le_bytes).collect(); assert_eq!(part, want); } + +// --------------------------------------------------------------------------- +// Size-changing filters ahead of a codec +// --------------------------------------------------------------------------- + +/// Fletcher32 placed before the compressor (NetCDF-4's ordering) makes the +/// codec's decoded output 4 bytes larger than the chunk. The decompression +/// cap was the chunk size for every stage, so these files failed with +/// "deflate: output exceeds size limit". +#[test] +fn h5py_fletcher32_before_deflate_reads() { + skip_if_no_python!(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("fletcher_first.h5"); + let p = path.display().to_string(); + run_python(&format!( + r#" +import h5py, numpy as np +arr = np.sin(np.arange(5000) / 50.0) +grid = np.arange(37 * 21, dtype=" = (0..5000).map(|i| (f64::from(i) / 50.0).sin()).collect(); + for name in ["fl_shuf_gzip", "fl_gzip", "shuf_fl_gzip"] { + let values = file.dataset(name).unwrap().read_f64().unwrap(); + assert_eq!(values.len(), arr.len(), "{name}"); + for (i, (v, w)) in values.iter().zip(&arr).enumerate() { + assert!((v - w).abs() < 1e-12, "{name}[{i}]: {v} vs {w}"); + } + } + let grid: Vec = (0..37 * 21).collect(); + assert_eq!( + file.dataset("grid_fl_shuf_gzip") + .unwrap() + .read_i32() + .unwrap(), + grid + ); +}