diff --git a/src/read/mod.rs b/src/read/mod.rs index 50bcd7b2..40cba663 100644 --- a/src/read/mod.rs +++ b/src/read/mod.rs @@ -966,31 +966,7 @@ impl<'data> CompressedData<'data> { match self.format { CompressionFormat::None => Ok(Cow::Borrowed(self.data)), #[cfg(feature = "compression")] - CompressionFormat::Zlib => { - use core::convert::TryInto; - let size = self - .uncompressed_size - .try_into() - .ok() - .read_error("Uncompressed data size is too large.")?; - let mut decompressed = Vec::new(); - decompressed - .try_reserve_exact(size) - .ok() - .read_error("Uncompressed data allocation failed")?; - let mut decompress = flate2::Decompress::new(true); - decompress - .decompress_vec( - self.data, - &mut decompressed, - flate2::FlushDecompress::Finish, - ) - .ok() - .read_error("Invalid zlib compressed data")?; - Ok(Cow::Owned(decompressed)) - } - #[cfg(feature = "compression")] - CompressionFormat::Zstandard => { + CompressionFormat::Zlib | CompressionFormat::Zstandard => { use core::convert::TryInto; use std::io::Read; let size = self @@ -1003,13 +979,36 @@ impl<'data> CompressedData<'data> { .try_reserve_exact(size) .ok() .read_error("Uncompressed data allocation failed")?; - let mut decoder = ruzstd::StreamingDecoder::new(self.data) - .ok() - .read_error("Invalid zstd compressed data")?; - decoder - .read_to_end(&mut decompressed) - .ok() - .read_error("Invalid zstd compressed data")?; + + match self.format { + CompressionFormat::Zlib => { + let mut decompress = flate2::Decompress::new(true); + decompress + .decompress_vec( + self.data, + &mut decompressed, + flate2::FlushDecompress::Finish, + ) + .ok() + .read_error("Invalid zlib compressed data")?; + } + CompressionFormat::Zstandard => { + let mut decoder = ruzstd::StreamingDecoder::new(self.data) + .ok() + .read_error("Invalid zstd compressed data")?; + decoder + .read_to_end(&mut decompressed) + .ok() + .read_error("Invalid zstd compressed data")?; + } + _ => unreachable!(), + } + if size != decompressed.len() { + return Err(Error( + "Uncompressed data size does not match compression header", + )); + } + Ok(Cow::Owned(decompressed)) } _ => Err(Error("Unsupported compressed data.")),