From 768be51145297357e5bca98e6bc3bc57a85c914b Mon Sep 17 00:00:00 2001 From: Patrick Marks Date: Sat, 5 Mar 2022 16:17:16 -0800 Subject: [PATCH] [de]compress_to_buffer methods to allow reusing buffers. Fixes #16 --- src/block/mod.rs | 154 +++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 136 insertions(+), 18 deletions(-) diff --git a/src/block/mod.rs b/src/block/mod.rs index 0ff1071ce..7e7ca2a35 100644 --- a/src/block/mod.rs +++ b/src/block/mod.rs @@ -33,6 +33,24 @@ pub enum CompressionMode { DEFAULT, } +/// Returns the size of the buffer that is guaranteed to hold the result of +/// compressing `uncompressed_size` bytes of in data. Returns std::io::Error +/// with ErrorKind::InvalidInput if input data is too long to be compressed by lz4. +pub fn compress_bound(uncompressed_size: usize) -> Result { + // 0 iff src too large + let compress_bound: i32 = unsafe { LZ4_compressBound(uncompressed_size as i32) }; + + if uncompressed_size > (i32::max_value() as usize) || compress_bound <= 0 { + return Err(Error::new( + ErrorKind::InvalidInput, + "Compression input too long.", + )); + } + + Ok(compress_bound as usize) +} + + /// Compresses the full src buffer using the specified CompressionMode, where None and Some(Default) /// are treated equally. If prepend_size is set, the source length will be prepended to the output /// buffer. @@ -42,7 +60,6 @@ pub enum CompressionMode { /// Returns std::io::Error with ErrorKind::InvalidInput if the src buffer is too long. /// Returns std::io::Error with ErrorKind::Other if the compression failed inside the C library. If /// this happens, the C api was not able to provide more information about the cause. -/// pub fn compress(src: &[u8], mode: Option, prepend_size: bool) -> Result> { // 0 iff src too large let compress_bound: i32 = unsafe { LZ4_compressBound(src.len() as i32) }; @@ -63,26 +80,51 @@ pub fn compress(src: &[u8], mode: Option, prepend_size: bool) - }) as usize ]; + let dec_size = compress_to_buffer(src, mode, prepend_size, &mut compressed)?; + compressed.truncate(dec_size as usize); + Ok(compressed) +} + +/// Compresses the full `src` buffer using the specified CompressionMode, where None and Some(Default) +/// are treated equally, writing compressed bytes to `buffer`. +/// +/// # Errors +/// Returns std::io::Error with ErrorKind::InvalidInput if the src buffer is too long. +/// Returns std::io::Error with ErrorKind::Other if the compression data does not find in `buffer`. +pub fn compress_to_buffer(src: &[u8], mode: Option, prepend_size: bool, buffer: &mut [u8]) -> Result { + + // check that src isn't too big for lz4 + let max_len: i32 = unsafe { LZ4_compressBound(src.len() as i32) }; + + if src.len() > (i32::max_value() as usize) || max_len <= 0 { + return Err(Error::new( + ErrorKind::InvalidInput, + "Compression input too long.", + )); + } + let dec_size; { let dst_buf = if prepend_size { let size = src.len() as u32; - compressed[0] = size as u8; - compressed[1] = (size >> 8) as u8; - compressed[2] = (size >> 16) as u8; - compressed[3] = (size >> 24) as u8; - &mut compressed[4..] + buffer[0] = size as u8; + buffer[1] = (size >> 8) as u8; + buffer[2] = (size >> 16) as u8; + buffer[3] = (size >> 24) as u8; + &mut buffer[4..] } else { - &mut compressed + buffer }; + let buf_len = dst_buf.len() as i32; + dec_size = match mode { Some(CompressionMode::HIGHCOMPRESSION(level)) => unsafe { LZ4_compress_HC( src.as_ptr() as *const c_char, dst_buf.as_mut_ptr() as *mut c_char, src.len() as i32, - compress_bound, + buf_len, level, ) }, @@ -91,7 +133,7 @@ pub fn compress(src: &[u8], mode: Option, prepend_size: bool) - src.as_ptr() as *const c_char, dst_buf.as_mut_ptr() as *mut c_char, src.len() as i32, - compress_bound, + buf_len, accel, ) }, @@ -100,7 +142,7 @@ pub fn compress(src: &[u8], mode: Option, prepend_size: bool) - src.as_ptr() as *const c_char, dst_buf.as_mut_ptr() as *mut c_char, src.len() as i32, - compress_bound, + buf_len, ) }, }; @@ -109,8 +151,50 @@ pub fn compress(src: &[u8], mode: Option, prepend_size: bool) - return Err(Error::new(ErrorKind::Other, "Compression failed")); } - compressed.truncate(if prepend_size { dec_size + 4 } else { dec_size } as usize); - Ok(compressed) + let written_size = if prepend_size { + dec_size + 4 + } else { + dec_size + }; + + Ok(written_size as usize) +} + +fn get_decompressed_size(src: &[u8], uncompressed_size: Option) -> Result { + let size; + + if let Some(s) = uncompressed_size { + size = s; + } else { + if src.len() < 4 { + return Err(Error::new( + ErrorKind::InvalidInput, + "Source buffer must at least contain size prefix.", + )); + } + size = + (src[0] as i32) | (src[1] as i32) << 8 | (src[2] as i32) << 16 | (src[3] as i32) << 24; + } + + if size < 0 { + return Err(Error::new( + ErrorKind::InvalidInput, + if uncompressed_size.is_some() { + "Size parameter must not be negative." + } else { + "Parsed size prefix in buffer must not be negative." + }, + )); + } + + if unsafe { LZ4_compressBound(size) } <= 0 { + return Err(Error::new( + ErrorKind::InvalidInput, + "Given size parameter is too big", + )); + } + + Ok(size as usize) } /// Decompresses the src buffer. If uncompressed_size is None, the source length will be read from @@ -123,7 +207,16 @@ pub fn compress(src: &[u8], mode: Option, prepend_size: bool) - /// Returns std::io::Error with ErrorKind::InvalidData if the decompression failed inside the C /// library. This is most likely due to malformed input. /// -pub fn decompress(mut src: &[u8], uncompressed_size: Option) -> Result> { +pub fn decompress(src: &[u8], uncompressed_size: Option) -> Result> { + let size = get_decompressed_size(src, uncompressed_size)?; + + let mut buffer = vec![0u8; size]; + + decompress_to_buffer(src, uncompressed_size, &mut buffer)?; + Ok(buffer) +} + +pub fn decompress_to_buffer(mut src: &[u8], uncompressed_size: Option, buffer: &mut [u8]) -> Result { let size; if let Some(s) = uncompressed_size { @@ -159,11 +252,17 @@ pub fn decompress(mut src: &[u8], uncompressed_size: Option) -> Result buffer.len() { + return Err(Error::new( + ErrorKind::InvalidInput, + "buffer isn't large enough to hold decompressed data" + )) + } + let dec_bytes = unsafe { LZ4_decompress_safe( src.as_ptr() as *const c_char, - decompressed.as_mut_ptr() as *mut c_char, + buffer.as_mut_ptr() as *mut c_char, src.len() as i32, size, ) @@ -176,13 +275,14 @@ pub fn decompress(mut src: &[u8], uncompressed_size: Option) -> Result