Repository: image-rs/image-tiff
Pull request #159 (2022-04-15 01:18:47)

Make LZWReader incrementally decompress input
Avoids a potentially large allocation. 


Patch:
diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs
index 8ba323b..8884c87 100644
--- a/src/decoder/mod.rs
+++ b/src/decoder/mod.rs
@@ -1032,14 +1032,7 @@ impl<R: Read + Seek> Decoder<R> {
         let byte_order = self.reader.byte_order;
         let compression_method = self.image().compression_method;
 
-        let reader = Self::create_reader(
-            &mut self.reader,
-            compression_method,
-            length,
-            buffer.len(),
-            buffer.byte_len(),
-            self.limits.intermediate_buffer_size,
-        )?;
+        let reader = Self::create_reader(&mut self.reader, compression_method, length)?;
 
         // Read into output buffer.
         {
@@ -1089,7 +1082,6 @@ impl<R: Read + Seek> Decoder<R> {
 
         let tile_attrs = self.image().tile_attributes.as_ref().unwrap();
         let (padding_right, padding_down) = tile_attrs.get_padding(tile);
-        let tile_samples = tile_attrs.tile_samples();
         let tile_length = tile_attrs.tile_length;
         let row_samples = tile_attrs.row_samples();
         let padding_right_samples = padding_right * self.image().bits_per_sample.len();
@@ -1100,14 +1092,8 @@ impl<R: Read + Seek> Decoder<R> {
         let compression_method = self.image().compression_method;
         let photometric_interpretation = self.image().photometric_interpretation;
 
-        let mut reader = Self::create_reader(
-            &mut self.reader,
-            compression_method,
-            compressed_length,
-            tile_samples,
-            byte_len,
-            self.limits.intermediate_buffer_size,
-        )?;
+        let mut reader =
+            Self::create_reader(&mut self.reader, compression_method, compressed_length)?;
 
         for row in 0..(tile_length - padding_down) {
             let buf = match &mut buffer {
@@ -1149,21 +1135,11 @@ impl<R: Read + Seek> Decoder<R> {
         reader: &'r mut SmartReader<R>,
         compression_method: CompressionMethod,
         compressed_length: u64,
-        samples: usize,  // Expected chunk length in samples
-        byte_len: usize, // Byte length of the samples in result buffer
-        intermediate_buffer_size: usize,
     ) -> TiffResult<Box<dyn Read + 'r>> {
         Ok(match compression_method {
             CompressionMethod::None => Box::new(reader),
             CompressionMethod::LZW => {
-                let clen = usize::try_from(compressed_length)?;
-
-                if samples * byte_len > intermediate_buffer_size || clen > intermediate_buffer_size
-                {
-                    return Err(TiffError::LimitsExceeded);
-                }
-
-                Box::new(LZWReader::new(reader, clen, samples * byte_len)?.1)
+                Box::new(LZWReader::new(reader, usize::try_from(compressed_length)?))
             }
             CompressionMethod::PackBits => {
                 Box::new(PackBitsReader::new(reader, usize::try_from(compressed_length)?)?.1)
diff --git a/src/decoder/stream.rs b/src/decoder/stream.rs
index bd247a4..75d45fd 100644
--- a/src/decoder/stream.rs
+++ b/src/decoder/stream.rs
@@ -1,6 +1,7 @@
 //! All IO functionality needed for TIFF decoding
 
-use std::io::{self, Read, Seek};
+use std::convert::TryFrom;
+use std::io::{self, BufRead, BufReader, Read, Seek, Take};
 
 /// Byte order of the TIFF file.
 #[derive(Clone, Copy, Debug)]
@@ -131,63 +132,53 @@ pub type DeflateReader<'r, R> = flate2::read::ZlibDecoder<&'r mut SmartReader<R>
 ///
 
 /// Reader that decompresses LZW streams
-pub struct LZWReader {
-    buffer: io::Cursor<Vec<u8>>,
+pub struct LZWReader<R: Read> {
+    reader: BufReader<Take<R>>,
+    decoder: weezl::decode::Decoder,
 }
 
-impl LZWReader {
+impl<R: Read> LZWReader<R> {
     /// Wraps a reader
-    pub fn new<R>(
-        reader: &mut SmartReader<R>,
-        compressed_length: usize,
-        max_uncompressed_length: usize,
-    ) -> io::Result<(usize, LZWReader)>
-    where
-        R: Read + Seek,
-    {
-        let mut compressed = vec![0; compressed_length as usize];
-        reader.read_exact(&mut compressed[..])?;
-        let mut uncompressed = vec![0; max_uncompressed_length];
-        let mut decoder = weezl::decode::Decoder::with_tiff_size_switch(weezl::BitOrder::Msb, 8);
-
-        let mut bytes_read = 0;
-        let mut bytes_written = 0;
+    pub fn new(reader: R, compressed_length: usize) -> LZWReader<R> {
+        Self {
+            reader: BufReader::with_capacity(
+                (32 * 1024).min(compressed_length),
+                reader.take(u64::try_from(compressed_length).unwrap()),
+            ),
+            decoder: weezl::decode::Decoder::with_tiff_size_switch(weezl::BitOrder::Msb, 8),
+        }
+    }
+}
 
-        while bytes_written < max_uncompressed_length {
-            let result = decoder.decode_bytes(
-                &compressed[bytes_read..],
-                &mut uncompressed[bytes_written..],
-            );
-            bytes_read += result.consumed_in;
-            bytes_written += result.consumed_out;
+impl<R: Read> Read for LZWReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        loop {
+            let result = self.decoder.decode_bytes(self.reader.fill_buf()?, buf);
+            self.reader.consume(result.consumed_in);
 
             match result.status {
-                Ok(weezl::LzwStatus::Ok) => {}
-                Ok(weezl::LzwStatus::Done) => break,
+                Ok(weezl::LzwStatus::Ok) => {
+                    if result.consumed_out == 0 {
+                        continue;
+                    } else {
+                        return Ok(result.consumed_out);
+                    }
+                }
                 Ok(weezl::LzwStatus::NoProgress) => {
+                    assert_eq!(result.consumed_in, 0);
+                    assert_eq!(result.consumed_out, 0);
+                    assert!(self.reader.buffer().is_empty());
                     return Err(io::Error::new(
                         io::ErrorKind::UnexpectedEof,
                         "no lzw end code found",
-                    ))
+                    ));
+                }
+                Ok(weezl::LzwStatus::Done) => {
+                    return Ok(result.consumed_out);
                 }
                 Err(err) => return Err(io::Error::new(io::ErrorKind::InvalidData, err)),
             }
         }
-
-        uncompressed.resize(bytes_written, 0);
-        Ok((
-            uncompressed.len(),
-            LZWReader {
-                buffer: io::Cursor::new(uncompressed),
-            },
-        ))
-    }
-}
-
-impl Read for LZWReader {
-    #[inline]
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        self.buffer.read(buf)
     }
 }
 


Test patch:
diff --git a/tests/decode_images.rs b/tests/decode_images.rs
index 5f32de1..ee6ccd1 100644
--- a/tests/decode_images.rs
+++ b/tests/decode_images.rs
@@ -396,7 +396,7 @@ fn oom() {
 }
 
 #[test]
-fn oom_2() {
+fn fuzzer_testcase4() {
     let image = [
         73, 73, 42, 0, 8, 0, 0, 0, 8, 0, 0, 1, 4, 0, 1, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 40, 1, 0, 0,
         0, 158, 0, 0, 251, 3, 1, 3, 0, 1, 0, 0, 0, 5, 0, 0, 0, 6, 1, 3, 0, 1, 0, 0, 0, 0, 0, 0, 0,
@@ -409,7 +409,7 @@ fn oom_2() {
     let err = decoder.read_image().unwrap_err();
 
     match err {
-        tiff::TiffError::LimitsExceeded => {}
+        tiff::TiffError::IoError(_) => {}
         unexpected => panic!("Unexpected error {}", unexpected),
     }
 }