read: add check for uncompressed data size (#723)

gimli-rs · Aug 27, 2024 · 9c0b256 · 9c0b256
1 parent 3872c29
commit 9c0b256
Showing 1 changed file with 31 additions and 32 deletions.
diff --git 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.")),