From 36cc87543dd8a90342aa9d922d9f9fef852df91c Mon Sep 17 00:00:00 2001 From: pommicket Date: Mon, 4 Sep 2023 12:46:52 -0400 Subject: fixed those bugs!! --- examples/small.rs | 2 +- src/lib.rs | 183 ++++++++++++++++++++++++++++++++++++----------------- test/smallish.png | Bin 0 -> 16687 bytes test/smallish0.png | Bin 0 -> 49425 bytes 4 files changed, 125 insertions(+), 60 deletions(-) create mode 100644 test/smallish.png create mode 100644 test/smallish0.png diff --git a/examples/small.rs b/examples/small.rs index fd86cb6..e6af172 100644 --- a/examples/small.rs +++ b/examples/small.rs @@ -1,5 +1,5 @@ fn main() { - let mut data = &include_bytes!("../test/small_rgb.png")[..]; + let mut data = &include_bytes!("small.png")[..]; let header = tiny_png::read_png_header(&mut data).unwrap(); let mut buf = vec![0; header.required_bytes()]; let result = tiny_png::read_png(&mut data, Some(&header), &mut buf); diff --git a/src/lib.rs b/src/lib.rs index 2f2ff1a..5869868 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,6 +25,8 @@ pub enum Error { BadFilter, BadPlteChunk, BadTrnsChunk, + BadNlen, + NoIdat, } impl From for Error { @@ -56,6 +58,8 @@ impl Display for Error { Self::BadFilter => write!(f, "bad PNG filter"), Self::BadPlteChunk => write!(f, "bad PLTE chunk"), Self::BadTrnsChunk => write!(f, "bad tRNS chunk"), + Self::NoIdat => write!(f, "missing IDAT chunk"), + Self::BadNlen => write!(f, "LEN doesn't match NLEN"), } } } @@ -122,29 +126,56 @@ impl<'a> Read for &'a [u8] { } } -struct BlockReader<'a, R: Read> { +struct IdatReader<'a, R: Read> { inner: &'a mut R, - bytes_left: usize, + bytes_left_in_block: usize, + palette: &'a mut [[u8; 4]; 256], + header: &'a ImageHeader, } -impl BlockReader<'_, R> { - fn read(&mut self, buf: &mut [u8]) -> Result<(), Error> { - if buf.len() > self.bytes_left { - return Err(Error::UnexpectedEob); +impl IdatReader<'_, R> { + fn read_partial(&mut self, buf: &mut [u8]) -> Result> { + if self.bytes_left_in_block >= buf.len() { + self.inner.read(buf)?; + self.bytes_left_in_block -= buf.len(); + Ok(buf.len()) + } else { + if self.bytes_left_in_block > 0 { + self.inner.read(&mut buf[..self.bytes_left_in_block])?; + } + let bytes_read = self.bytes_left_in_block; + + // CRC + self.inner.skip_bytes(4)?; + match read_non_idat_chunks(self.inner, self.header, self.palette)? { + None => Ok(bytes_read), + Some(n) => { + self.bytes_left_in_block = n; + Ok(self.read_partial(&mut buf[bytes_read..])? + bytes_read) + } + } } - self.inner.read(buf)?; - self.bytes_left -= buf.len(); - Ok(()) } - fn read_partial(&mut self, buf: &mut [u8]) -> Result> { - let count = min(self.bytes_left, buf.len()); - self.read(&mut buf[..count])?; - Ok(count) + fn read(&mut self, buf: &mut [u8]) -> Result<(), Error> { + let count = self.read_partial(buf)?; + if count == buf.len() { + Ok(()) + } else { + Err(Error::UnexpectedEob) + } } fn read_to_end(&mut self) -> Result<(), Error> { - self.inner.skip_bytes(self.bytes_left)?; + self.inner.skip_bytes(self.bytes_left_in_block)?; + // CRC + self.inner.skip_bytes(4)?; + loop { + match read_non_idat_chunks(self.inner, self.header, self.palette)? { + None => break, + Some(n) => self.inner.skip_bytes(n + 4)?, + } + } Ok(()) } } @@ -224,7 +255,7 @@ impl ImageHeader { pub fn color_type(&self) -> ColorType { self.color_type } - fn checked_required_bytes(&self) -> Option { + fn checked_decompressed_size(&self) -> Option { let row_bytes = 1 + usize::try_from(self.width()) .ok()? .checked_mul(usize::from(self.bit_depth() as u8))? @@ -233,9 +264,18 @@ impl ImageHeader { / 8; row_bytes.checked_mul(usize::try_from(self.height()).ok()?) } + + fn decompressed_size(&self) -> usize { + self.checked_decompressed_size().unwrap() + } + + fn checked_required_bytes(&self) -> Option { + self.checked_decompressed_size() + } pub fn required_bytes(&self) -> usize { self.checked_required_bytes().unwrap() } + pub fn bytes_per_scanline(&self) -> usize { (self.width() as usize * usize::from(self.bit_depth() as u8) @@ -249,20 +289,18 @@ impl ImageHeader { } /// number of bits to read in each [`Read::read`] call. -/// -/// don't change this to something bigger than `u32`, since we don't want to overread past the zlib checksum. type ReadBits = u32; /// number of bits to store in the [`BitReader`] buffer. type Bits = u64; struct BitReader<'a, R: Read> { - inner: BlockReader<'a, R>, + inner: IdatReader<'a, R>, bits: Bits, bits_left: u8, } -impl<'a, R: Read> From> for BitReader<'a, R> { - fn from(inner: BlockReader<'a, R>) -> Self { +impl<'a, R: Read> From> for BitReader<'a, R> { + fn from(inner: IdatReader<'a, R>) -> Self { Self { inner, bits: 0, @@ -313,6 +351,16 @@ impl BitReader<'_, R> { debug_assert!(count <= 16); self.read_bits(count).map(|x| x as u16) } + + fn read_aligned_bytes(&mut self, buf: &mut [u8]) -> Result<(), Error> { + debug_assert_eq!(self.bits_left % 8, 0); + let mut i = 0; + while self.bits_left > 0 && i < buf.len() { + buf[i] = self.read_bits_u8(8)?; + i += 1; + } + self.inner.read(&mut buf[i..]) + } } pub fn read_png_header(reader: &mut R) -> Result> { @@ -435,12 +483,14 @@ impl HuffmanTable { } else { // put it in a subtable. let main_table_entry = usize::from(code) & (HUFFMAN_MAIN_TABLE_SIZE - 1); - let mut subtable_index = self.main_table[main_table_entry] & 0x1ff; - if subtable_index == 0 { - subtable_index = self.subtables_used; - self.main_table[main_table_entry] = 0x200 | subtable_index; + let subtable_index = if self.main_table[main_table_entry] == 0 { + let i = self.subtables_used; + self.main_table[main_table_entry] = 0x200 | i; self.subtables_used += 1; - } + i + } else { + self.main_table[main_table_entry] & 0x1ff + }; let subtable = &mut self.subtables[usize::from(subtable_index)]; for i in 0..1u16 << (HUFFMAN_MAX_BITS - length) { subtable[usize::from( @@ -517,10 +567,10 @@ fn read_compressed_block( i += 1; } else if op == 16 { let rep = reader.read_bits_usize(2)? + 3; - if i + rep > total_code_lengths { + if i == 0 || i + rep > total_code_lengths { return Err(Error::BadHuffmanCodes); } - let l = code_lengths[i]; + let l = code_lengths[i - 1]; for _ in 0..rep { code_lengths[i] = l; i += 1; @@ -550,11 +600,9 @@ fn read_compressed_block( break; } } - let literal_length_code_lengths = &code_lengths[0..literal_length_code_lengths_count]; let distance_code_lengths = &code_lengths[literal_length_code_lengths_count..total_code_lengths]; - literal_length_table = HuffmanTable::from_code_lengths(literal_length_code_lengths); distance_table = HuffmanTable::from_code_lengths(distance_code_lengths); } else { @@ -584,17 +632,14 @@ fn read_compressed_block( match literal_length { 0..=255 => { // literal - //println!("lit {literal_length}"); writer.write_byte(literal_length as u8)?; } 256 => { // end of block - //println!("eob"); break; } _ => { // length + distance - //println!("{literal_length}"); let length = match literal_length { 257..=264 => literal_length - 254, 265..=284 => { @@ -626,7 +671,6 @@ fn read_compressed_block( } _ => return Err(Error::BadCode), }; - //println!("D {distance} L {length}"); writer.copy(usize::from(distance), usize::from(length))?; } } @@ -635,24 +679,30 @@ fn read_compressed_block( } fn read_idat( - mut reader: BlockReader<'_, R>, + reader: IdatReader<'_, R>, writer: &mut DecompressedDataWriter, ) -> Result<(), Error> { - let mut zlib_header = [0; 2]; - reader.read(&mut zlib_header)?; - let mut reader = BitReader::from(reader); - loop { + let _zlib_header = reader.read_bits(16); + let decompressed_size = reader.inner.header.decompressed_size(); + while writer.pos < decompressed_size { let bfinal = reader.read_bits(1)?; let btype = reader.read_bits(2)?; if btype == 0 { // uncompressed block - let len = reader.read_bits_usize(16)?; - reader.read_bits(16)?; // nlen - if writer.slice.len() < len { + reader.bits >>= reader.bits_left % 8; + reader.bits_left -= reader.bits_left % 8; + let len = reader.read_bits_u16(16)?; + let nlen = reader.read_bits_u16(16)?; + if len ^ nlen != 0xffff { + return Err(Error::BadNlen); + } + let len: usize = len.into(); + if len > writer.slice.len() - writer.pos { return Err(Error::TooMuchData); } - reader.inner.read(&mut writer.slice[..len])?; + reader.read_aligned_bytes(&mut writer.slice[writer.pos..writer.pos + len])?; + writer.pos += len; } else if btype == 1 || btype == 2 { // compressed block read_compressed_block(&mut reader, writer, btype == 2)?; @@ -761,17 +811,11 @@ impl ImageData<'_> { } } -pub fn read_png<'a, R: Read>( +fn read_non_idat_chunks( reader: &mut R, - header: Option<&ImageHeader>, - buf: &'a mut [u8], -) -> Result, Error> { - let header = match header { - None => read_png_header(reader)?, - Some(h) => *h, - }; - let mut writer = DecompressedDataWriter::from(buf); - let mut palette = [[0, 0, 0, 0]; 256]; + header: &ImageHeader, + palette: &mut [[u8; 4]; 256], +) -> Result, Error> { loop { let mut chunk_header = [0; 8]; reader.read(&mut chunk_header[..])?; @@ -790,13 +834,7 @@ pub fn read_png<'a, R: Read>( if &chunk_type == b"IEND" { break; } else if &chunk_type == b"IDAT" { - read_idat( - BlockReader { - inner: reader, - bytes_left: chunk_len + 4, - }, - &mut writer, - )?; + return Ok(Some(chunk_len)); } else if &chunk_type == b"PLTE" && header.color_type == ColorType::Indexed { if chunk_len > 256 * 3 || chunk_len % 3 != 0 { return Err(Error::BadPlteChunk); @@ -827,6 +865,33 @@ pub fn read_png<'a, R: Read>( return Err(Error::UnrecognizedChunk(chunk_type)); } } + Ok(None) +} + +pub fn read_png<'a, R: Read>( + reader: &mut R, + header: Option<&ImageHeader>, + buf: &'a mut [u8], +) -> Result, Error> { + let header = match header { + None => read_png_header(reader)?, + Some(h) => *h, + }; + let mut writer = DecompressedDataWriter::from(buf); + let mut palette = [[0, 0, 0, 0]; 256]; + let Some(idat_len) = read_non_idat_chunks(reader, &header, &mut palette)? else { + return Err(Error::NoIdat); + }; + read_idat( + IdatReader { + inner: reader, + bytes_left_in_block: idat_len, + header: &header, + palette: &mut palette, + }, + &mut writer, + )?; + let buf = writer.slice; apply_filters(&header, buf)?; Ok(ImageData { diff --git a/test/smallish.png b/test/smallish.png new file mode 100644 index 0000000..e6bb837 Binary files /dev/null and b/test/smallish.png differ diff --git a/test/smallish0.png b/test/smallish0.png new file mode 100644 index 0000000..b860c80 Binary files /dev/null and b/test/smallish0.png differ -- cgit v1.2.3