From 007e7da4e09d1a5f74fff484f119acab2a3b1c11 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 5 Aug 2026 23:39:30 +0800 Subject: [PATCH] fix(encoding): bound general block decompression --- .../src/encodings/physical/block.rs | 155 ++++++++++++++++++ 1 file changed, 155 insertions(+) diff --git a/rust/lance-encoding/src/encodings/physical/block.rs b/rust/lance-encoding/src/encodings/physical/block.rs index 9e2eb94abd7..19c9e1b945e 100644 --- a/rust/lance-encoding/src/encodings/physical/block.rs +++ b/rust/lance-encoding/src/encodings/physical/block.rs @@ -307,6 +307,65 @@ mod zstd { Ok(()) } + fn decompress_exact( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + expected_bytes: usize, + ) -> Result<()> { + if input_buf.is_empty() { + if expected_bytes == 0 { + return Ok(()); + } + return Err(Error::invalid_input(format!( + "Zstd payload is empty, expected {expected_bytes} output bytes" + ))); + } + let compressed = if self.is_raw_stream_format(input_buf) { + input_buf + } else { + let declared = u64::from_le_bytes(input_buf[..8].try_into().map_err(|_| { + Error::invalid_input("Zstd payload is shorter than its 8-byte length prefix") + })?); + let declared = usize::try_from(declared).map_err(|_| { + Error::invalid_input("Zstd declared output length does not fit usize") + })?; + if declared != expected_bytes { + return Err(Error::invalid_input(format!( + "Zstd declares {declared} output bytes, expected {expected_bytes}" + ))); + } + &input_buf[8..] + }; + let start = output_buf.len(); + let end = start + .checked_add(expected_bytes) + .ok_or_else(|| Error::invalid_input("Zstd output buffer length overflows usize"))?; + output_buf + .try_reserve_exact(expected_bytes) + .map_err(|error| { + Error::invalid_input(format!( + "Zstd could not reserve {expected_bytes} output bytes: {error}" + )) + })?; + output_buf.resize(end, 0); + let decoded = decompress_to_buffer(compressed, &mut output_buf[start..]); + let decoded = match decoded { + Ok(decoded) => decoded, + Err(error) => { + output_buf.truncate(start); + return Err(error.into()); + } + }; + if decoded != expected_bytes { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd produced {decoded} bytes, expected {expected_bytes}" + ))); + } + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::Zstd, @@ -375,6 +434,56 @@ mod lz4 { Ok(()) } + fn decompress_exact( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + expected_bytes: usize, + ) -> Result<()> { + let declared = input_buf.get(..4).ok_or_else(|| { + Error::invalid_input("LZ4 payload is shorter than its 4-byte length prefix") + })?; + let declared = + u32::from_le_bytes(declared.try_into().map_err(|_| { + Error::invalid_input("LZ4 payload has an invalid length prefix") + })?) as usize; + if declared != expected_bytes { + return Err(Error::invalid_input(format!( + "LZ4 declares {declared} output bytes, expected {expected_bytes}" + ))); + } + let start = output_buf.len(); + let end = start + .checked_add(expected_bytes) + .ok_or_else(|| Error::invalid_input("LZ4 output buffer length overflows usize"))?; + output_buf + .try_reserve_exact(expected_bytes) + .map_err(|error| { + Error::invalid_input(format!( + "LZ4 could not reserve {expected_bytes} output bytes: {error}" + )) + })?; + output_buf.resize(end, 0); + let decoded = + ::lz4::block::decompress_to_buffer(input_buf, None, &mut output_buf[start..]); + let decoded = match decoded { + Ok(decoded) => decoded, + Err(error) => { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "LZ4 decompression error: {error}" + ))); + } + }; + if decoded != expected_bytes { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "LZ4 produced {decoded} bytes, expected {expected_bytes}" + ))); + } + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::Lz4, @@ -398,6 +507,22 @@ impl BufferCompressor for NoopBufferCompressor { Ok(()) } + fn decompress_exact( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + expected_bytes: usize, + ) -> Result<()> { + if input_buf.len() != expected_bytes { + return Err(Error::invalid_input(format!( + "Uncompressed payload has {} bytes, expected {expected_bytes}", + input_buf.len() + ))); + } + output_buf.extend_from_slice(input_buf); + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::None, @@ -727,6 +852,21 @@ mod tests { assert_eq!(input_data, decompressed_data.as_slice()); } + #[test] + fn test_zstd_exact_output_rejects_declared_length() { + let compressor = ZstdBufferCompressor::new(0); + let mut compressed = Vec::new(); + compressor.compress(b"bounded", &mut compressed).unwrap(); + compressed[..8].copy_from_slice(&u64::MAX.to_le_bytes()); + + let mut output = Vec::new(); + let error = compressor + .decompress_exact(&compressed, &mut output, 7) + .unwrap_err(); + assert!(error.to_string().contains("expected 7")); + assert!(output.is_empty()); + } + #[test] fn test_zstd_compress_decompress_multiple_times() { let compressor = ZstdBufferCompressor::new(0); @@ -829,6 +969,21 @@ mod tests { assert_eq!(input_data, decompressed_data.as_slice()); } + #[test] + fn test_lz4_exact_output_rejects_declared_length() { + let compressor = Lz4BufferCompressor::default(); + let mut compressed = Vec::new(); + compressor.compress(b"bounded", &mut compressed).unwrap(); + compressed[..4].copy_from_slice(&u32::MAX.to_le_bytes()); + + let mut output = Vec::new(); + let error = compressor + .decompress_exact(&compressed, &mut output, 7) + .unwrap_err(); + assert!(error.to_string().contains("declares")); + assert!(output.is_empty()); + } + #[test_log::test(tokio::test)] async fn test_lz4_compress_round_trip() { for data_type in &[