Merge pull request #159 from fintelia/stream-lzw

Make LZWReader incrementally decompress input
This commit is contained in:
HeroicKatora authored and GitHub committed 2022-04-15 17:35:43 +02:00
commit 167088b33d
3 files changed
+41 -74

No files matched your search

+4 -28
View File
@@ -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)
+35 -44
View File
@@ -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);
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),
}
}
}
let mut bytes_read = 0;
let mut bytes_written = 0;
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)
}
}
+2 -2
View File
@@ -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),
}
}