From 8e88d32b9d3e093513f5ecda2351a5f691e9bfb8 Mon Sep 17 00:00:00 2001 From: Andreas Molzer Date: Mon, 27 Sep 2021 01:10:11 +0200 Subject: [PATCH 1/7] Add validation: chunks should have valid length --- src/encoder.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/encoder.rs b/src/encoder.rs index 76f50b5793..e9163b4a57 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -529,7 +529,18 @@ impl Writer { Ok(self) } + /// Write a raw chunk of PNG data. + /// + /// The chunk will have its CRC calculated and correctly. The data is not filtered in any way, + /// but the chunk needs to be short enough to have its length encoded correctly. pub fn write_chunk(&mut self, name: ChunkType, data: &[u8]) -> Result<()> { + use std::convert::TryFrom; + + if u32::try_from(data.len()).map_or(false, |length| length > i32::MAX as u32) { + let kind = FormatErrorKind::WrittenTooMuch(data.len() - i32::MAX as usize); + return Err(EncodingError::Format(kind.into())); + } + write_chunk(&mut self.w, name, data) } From 4b10bf8313c7a7db1aebc3cb8f18c0e922429547 Mon Sep 17 00:00:00 2001 From: Andreas Molzer Date: Mon, 27 Sep 2021 01:10:47 +0200 Subject: [PATCH 2/7] Simplify/Rework logic of frame tracking We no longer _enforce_ that the number of frames is exactly as deduced by animation control/separate default image. This isn't necessary for a correct chunk stream. It's not exactly correct chunk ordering though: > Multiple IDAT chunks shall be consecutive Instead, we leniently allow encoding of multiple images. When we run out of animation frames we simply switch back to IDAT. The behavior when no animation frames had been added should then be the one of version 0.16. --- src/encoder.rs | 155 +++++++++++++++++++++++++++++++++---------------- 1 file changed, 104 insertions(+), 51 deletions(-) diff --git a/src/encoder.rs b/src/encoder.rs index e9163b4a57..5803538810 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -138,7 +138,13 @@ impl From for EncodingError { } } -/// PNG Encoder +/// PNG Encoder. +/// +/// This configures the PNG format options such as animation chunks, palette use, color types, +/// auxiliary chunks etc. +/// +/// FIXME: Configuring APNG might be easier (less individual errors) if we had an _adapter_ which +/// borrows this mutably but guarantees that `info.frame_control` is not `None`. pub struct Encoder<'a, W: Write> { w: W, info: Info<'a>, @@ -162,24 +168,31 @@ impl<'a, W: Write> Encoder<'a, W> { /// /// `num_frames` controls how many frames the animation has, while /// `num_plays` controls how many times the animation should be - /// repeaded until it stops, if it's zero then it will repeat - /// inifinitely + /// repeated until it stops, if it's zero then it will repeat + /// infinitely. + /// + /// When this method is returns successfully then the images written will be encoded as fdAT + /// chunks, except for the first image that is still encoded as `IDAT`. You can control if the + /// first frame should be treated as an animation frame with [`set_sep_def_img()`]. /// /// This method returns an error if `num_frames` is 0. pub fn set_animated(&mut self, num_frames: u32, num_plays: u32) -> Result<()> { if num_frames == 0 { return Err(EncodingError::Format(FormatErrorKind::ZeroFrames.into())); } + let actl = AnimationControl { num_frames, num_plays, }; + let fctl = FrameControl { sequence_number: 0, width: self.info.width, height: self.info.height, ..Default::default() }; + self.info.animation_control = Some(actl); self.info.frame_control = Some(fctl); Ok(()) @@ -228,6 +241,9 @@ impl<'a, W: Write> Encoder<'a, W> { self.info.srgb = Some(rendering_intent); } + /// Start encoding by writing the header data. + /// + /// The remaining data can be supplied by methods on the returned [`Writer`]. pub fn write_header(self) -> Result> { Writer::new( self.w, @@ -407,6 +423,13 @@ impl<'a, W: Write> Encoder<'a, W> { } /// PNG writer +/// +/// Progresses through the image by writing images, frames, or raw individual chunks. This is +/// constructed through [`Encoder::write_header()`]. +/// +/// FIXME: Writing of animated chunks might be clearer if we had an _adapter_ that you would call +/// to guarantee the next image to be prefaced with a fcTL-chunk, and all other chunks would be +/// guaranteed to be `IDAT`/not affected by APNG's frame control. pub struct Writer { w: W, info: PartialInfo, @@ -414,6 +437,7 @@ pub struct Writer { adaptive_filter: AdaptiveFilterType, sep_def_img: bool, written: u64, + animation_written: u32, } /// Contains the subset of attributes of [Info] needed for [Writer] to function @@ -500,6 +524,7 @@ impl Writer { adaptive_filter, sep_def_img, written: 0, + animation_written: 0, } } @@ -556,20 +581,16 @@ impl Writer { } } - /// Writes the image data. - pub fn write_image_data(&mut self, data: &[u8]) -> Result<()> { - const MAX_IDAT_CHUNK_LEN: u32 = std::u32::MAX >> 1; - #[allow(non_upper_case_globals)] - const MAX_fdAT_CHUNK_LEN: u32 = (std::u32::MAX >> 1) - 4; + const MAX_IDAT_CHUNK_LEN: u32 = std::u32::MAX >> 1; + #[allow(non_upper_case_globals)] + const MAX_fdAT_CHUNK_LEN: u32 = (std::u32::MAX >> 1) - 4; + /// Writes the next image data. + pub fn write_image_data(&mut self, data: &[u8]) -> Result<()> { if self.info.color_type == ColorType::Indexed && !self.info.has_palette { return Err(EncodingError::Format(FormatErrorKind::NoPalette.into())); } - if self.written > self.max_frames() { - return Err(EncodingError::Format(FormatErrorKind::EndReached.into())); - } - let width: usize; let height: usize; if let Some(ref mut fctl) = self.info.frame_control { @@ -603,6 +624,7 @@ impl Writer { let bpp = self.info.bpp_in_prediction(); let filter_method = self.filter; let adaptive_method = self.adaptive_filter; + for line in data.chunks(in_len) { current.copy_from_slice(&line); let filter_type = filter(filter_method, adaptive_method, bpp, &prev, &mut current); @@ -611,33 +633,51 @@ impl Writer { prev = line; } let zlib_encoded = zlib.finish()?; - if self.sep_def_img || self.info.frame_control.is_none() { - self.sep_def_img = false; - for chunk in zlib_encoded.chunks(MAX_IDAT_CHUNK_LEN as usize) { - self.write_chunk(chunk::IDAT, &chunk)?; - } - } else if let Some(ref mut fctl) = self.info.frame_control { - fctl.encode(&mut self.w)?; - fctl.sequence_number = fctl.sequence_number.wrapping_add(1); - if self.written == 0 { - for chunk in zlib_encoded.chunks(MAX_IDAT_CHUNK_LEN as usize) { - self.write_chunk(chunk::IDAT, &chunk)?; - } - } else { - let buff_size = zlib_encoded.len().min(MAX_fdAT_CHUNK_LEN as usize); - let mut alldata = vec![0u8; 4 + buff_size]; - for chunk in zlib_encoded.chunks(MAX_fdAT_CHUNK_LEN as usize) { - alldata[..4].copy_from_slice(&fctl.sequence_number.to_be_bytes()); - alldata[4..][..chunk.len()].copy_from_slice(chunk); - write_chunk(&mut self.w, chunk::fdAT, &alldata[..4 + chunk.len()])?; - fctl.sequence_number = fctl.sequence_number.wrapping_add(1); + match self.info.frame_control { + None => { + self.write_zlib_encoded_idat(&zlib_encoded)?; + } + Some(_) if self.sep_def_img => { + self.sep_def_img = false; + self.write_zlib_encoded_idat(&zlib_encoded)?; + } + Some(ref mut fctl) => { + fctl.encode(&mut self.w)?; + fctl.sequence_number = fctl.sequence_number.wrapping_add(1); + self.animation_written += 1; + + // If the default image is the first frame of an animation, it's still an IDAT. + if self.written == 0 { + self.write_zlib_encoded_idat(&zlib_encoded)?; + } else { + let buff_size = zlib_encoded.len().min(Self::MAX_fdAT_CHUNK_LEN as usize); + let mut alldata = vec![0u8; 4 + buff_size]; + for chunk in zlib_encoded.chunks(Self::MAX_fdAT_CHUNK_LEN as usize) { + alldata[..4].copy_from_slice(&fctl.sequence_number.to_be_bytes()); + alldata[4..][..chunk.len()].copy_from_slice(chunk); + write_chunk(&mut self.w, chunk::fdAT, &alldata[..4 + chunk.len()])?; + fctl.sequence_number = fctl.sequence_number.wrapping_add(1); + } } } - } else { - unreachable!(); // this should be unreachable } + self.written += 1; + if let Some(actl) = self.info.animation_control { + if actl.num_frames <= self.animation_written { + // If we've written all animation frames, all following will be normal image chunks. + self.info.frame_control = None; + } + } + + Ok(()) + } + + fn write_zlib_encoded_idat(&mut self, zlib_encoded: &[u8]) -> Result<()> { + for chunk in zlib_encoded.chunks(Self::MAX_IDAT_CHUNK_LEN as usize) { + self.write_chunk(chunk::IDAT, &chunk)?; + } Ok(()) } @@ -1115,8 +1155,6 @@ pub struct StreamWriter<'a, W: Write> { line_len: usize, /// size of the frame (width * height * sample_size) to_write: usize, - /// Flag used to signal the end of the image - end: bool, width: u32, height: u32, @@ -1130,10 +1168,6 @@ pub struct StreamWriter<'a, W: Write> { impl<'a, W: Write> StreamWriter<'a, W> { fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> Result> { - if writer.max_frames() < writer.written { - return Err(EncodingError::Format(FormatErrorKind::EndReached.into())); - } - let PartialInfo { width, height, @@ -1159,7 +1193,6 @@ impl<'a, W: Write> StreamWriter<'a, W> { index: 0, prev_buf, curr_buf, - end: false, bpp, filter, width, @@ -1352,13 +1385,11 @@ impl<'a, W: Write> StreamWriter<'a, W> { } pub fn finish(mut self) -> Result<()> { - if !self.end { - let err = FormatErrorKind::MissingFrames.into(); - return Err(EncodingError::Format(err)); - } else if self.to_write > 0 { + if self.to_write > 0 { let err = FormatErrorKind::MissingData(self.to_write).into(); return Err(EncodingError::Format(err)); } + // TODO: call `writer.finish` somehow? self.flush()?; Ok(()) @@ -1381,7 +1412,6 @@ impl<'a, W: Write> StreamWriter<'a, W> { self.to_write = size; wrt.writer.written += 1; wrt.write_header()?; - self.end = wrt.writer.written + 1 == wrt.writer.max_frames(); // now it can be taken because the next statements cannot cause any errors let wrt = match self.writer.take() { @@ -1405,10 +1435,6 @@ impl<'a, W: Write> Write for StreamWriter<'a, W> { } if self.to_write == 0 { - if self.end { - let err = FormatErrorKind::EndReached.into(); - return Err(EncodingError::Format(err).into()); - } match self.writer.take() { Wrapper::Zlib(wrt) => match wrt.finish() { Ok(chunk) => self.writer = Wrapper::Chunk(chunk), @@ -1422,6 +1448,7 @@ impl<'a, W: Write> Write for StreamWriter<'a, W> { }; self.new_frame()?; } + let written = data.read(&mut self.curr_buf[..self.line_len][self.index..])?; self.index += written; self.to_write -= written; @@ -1875,6 +1902,32 @@ mod tests { Ok(()) } + #[test] + fn write_image_chunks_beyond_first() -> Result<()> { + use std::io::Cursor; + let width = 10; + let height = 10; + + let output = vec![0u8; 1024]; + let writer = Cursor::new(output); + + // Not an animation but we should still be able to write multiple images + // See issue: + // This is technically all valid png so there is no issue with correctness. + let mut encoder = Encoder::new(writer, width, height); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + let mut png_writer = encoder.write_header()?; + + for _ in 0..3 { + let correct_image_size = (width * height) as usize; + let image = vec![0u8; correct_image_size]; + png_writer.write_image_data(image.as_ref())?; + } + + Ok(()) + } + /// A Writer that only writes a few bytes at a time struct RandomChunkWriter { rng: R, @@ -1899,7 +1952,7 @@ mod tests { /// /// Since this only contains trait impls, there is no need to make this public, they are simply /// available when the mod is compiled as well. -impl crate::common::Compression { +impl Compression { fn to_options(self) -> deflate::CompressionOptions { match self { Compression::Default => deflate::CompressionOptions::default(), From e17523573f23848fe2ef54d80a532ba28af260c4 Mon Sep 17 00:00:00 2001 From: Andreas Molzer Date: Tue, 28 Sep 2021 22:27:36 +0200 Subject: [PATCH 3/7] Test new error check in chunk writer --- src/encoder.rs | 51 +++++++++++++++++++++++++++++++++++++------------- 1 file changed, 38 insertions(+), 13 deletions(-) diff --git a/src/encoder.rs b/src/encoder.rs index 5803538810..6c2ff7f7b4 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -561,7 +561,7 @@ impl Writer { pub fn write_chunk(&mut self, name: ChunkType, data: &[u8]) -> Result<()> { use std::convert::TryFrom; - if u32::try_from(data.len()).map_or(false, |length| length > i32::MAX as u32) { + if u32::try_from(data.len()).map_or(true, |length| length > i32::MAX as u32) { let kind = FormatErrorKind::WrittenTooMuch(data.len() - i32::MAX as usize); return Err(EncodingError::Format(kind.into())); } @@ -1500,7 +1500,7 @@ mod tests { use rand::{thread_rng, Rng}; use std::fs::File; - use std::io::Write; + use std::io::{Cursor, Write}; use std::{cmp, io}; #[test] @@ -1643,8 +1643,6 @@ mod tests { #[test] fn expect_error_on_wrong_image_len() -> Result<()> { - use std::io::Cursor; - let width = 10; let height = 10; @@ -1665,8 +1663,6 @@ mod tests { #[test] fn expect_error_on_empty_image() -> Result<()> { - use std::io::Cursor; - let output = vec![0u8; 1024]; let mut writer = Cursor::new(output); @@ -1684,8 +1680,6 @@ mod tests { #[test] fn expect_error_on_invalid_bit_depth_color_type_combination() -> Result<()> { - use std::io::Cursor; - let output = vec![0u8; 1024]; let mut writer = Cursor::new(output); @@ -1744,8 +1738,6 @@ mod tests { #[test] fn can_write_header_with_valid_bit_depth_color_type_combination() -> Result<()> { - use std::io::Cursor; - let output = vec![0u8; 1024]; let mut writer = Cursor::new(output); @@ -1839,7 +1831,7 @@ mod tests { encoder.set_filter(filter); encoder.write_header()?.write_image_data(&pixel)?; - let decoder = crate::Decoder::new(io::Cursor::new(buffer)); + let decoder = crate::Decoder::new(Cursor::new(buffer)); let mut reader = decoder.read_info()?; let info = reader.info(); assert_eq!(info.width, 4); @@ -1875,7 +1867,7 @@ mod tests { } encoder.write_header()?.write_image_data(&pixel)?; - let decoder = crate::Decoder::new(io::Cursor::new(buffer)); + let decoder = crate::Decoder::new(Cursor::new(buffer)); let mut reader = decoder.read_info()?; assert_eq!( reader.info().source_gamma, @@ -1904,7 +1896,6 @@ mod tests { #[test] fn write_image_chunks_beyond_first() -> Result<()> { - use std::io::Cursor; let width = 10; let height = 10; @@ -1928,6 +1919,40 @@ mod tests { Ok(()) } + #[test] + #[cfg(all(unix, not(target_pointer_width = "32")))] + fn exper_error_on_huge_chunk() -> Result<()> { + // Okay, so we want a proper 4 GB chunk but not actually spend the memory for reserving it. + // Let's rely on overcommit? Otherwise we got the rather dumb option of mmap-ing /dev/zero. + let empty = vec![0; 1usize << 31]; + let writer = Cursor::new(vec![0u8; 1024]); + + let mut encoder = Encoder::new(writer, 10, 10); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + let mut png_writer = encoder.write_header()?; + + assert!(png_writer.write_chunk(chunk::fdAT, &empty).is_err()); + Ok(()) + } + + #[test] + #[cfg(all(unix, not(target_pointer_width = "32")))] + fn exper_error_on_non_u32_chunk() -> Result<()> { + // Okay, so we want a proper 4 GB chunk but not actually spend the memory for reserving it. + // Let's rely on overcommit? Otherwise we got the rather dumb option of mmap-ing /dev/zero. + let empty = vec![0; 1usize << 32]; + let writer = Cursor::new(vec![0u8; 1024]); + + let mut encoder = Encoder::new(writer, 10, 10); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + let mut png_writer = encoder.write_header()?; + + assert!(png_writer.write_chunk(chunk::fdAT, &empty).is_err()); + Ok(()) + } + /// A Writer that only writes a few bytes at a time struct RandomChunkWriter { rng: R, From 57af3468176c1d57721b3d5f06ff5274bba8ed20 Mon Sep 17 00:00:00 2001 From: Andreas Molzer Date: Wed, 29 Sep 2021 00:44:20 +0200 Subject: [PATCH 4/7] Group options, add new validate option --- src/encoder.rs | 97 ++++++++++++++++++++++++++++++-------------------- 1 file changed, 59 insertions(+), 38 deletions(-) diff --git a/src/encoder.rs b/src/encoder.rs index 6c2ff7f7b4..fa9667912e 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -148,9 +148,26 @@ impl From for EncodingError { pub struct Encoder<'a, W: Write> { w: W, info: Info<'a>, + options: Options, +} + +/// Decoding options, internal type, forwarded to the Writer. +struct Options { filter: FilterType, adaptive_filter: AdaptiveFilterType, sep_def_img: bool, + validate_sequence: bool, +} + +impl Default for Options { + fn default() -> Options { + Options { + filter: FilterType::default(), + adaptive_filter: AdaptiveFilterType::default(), + sep_def_img: false, + validate_sequence: false, + } + } } impl<'a, W: Write> Encoder<'a, W> { @@ -158,9 +175,7 @@ impl<'a, W: Write> Encoder<'a, W> { Encoder { w, info: Info::with_size(width, height), - filter: FilterType::default(), - adaptive_filter: AdaptiveFilterType::default(), - sep_def_img: false, + options: Options::default(), } } @@ -198,9 +213,18 @@ impl<'a, W: Write> Encoder<'a, W> { Ok(()) } + /// Mark the first animated frame as a 'separate default image'. + /// + /// In APNG each animated frame is preceded by a special control chunk, `fcTL`. It's up to the + /// encoder to decide if the first image, the standard `IDAT` data, should be part of the + /// animation by emitting this chunk or by not doing so. A default image that is _not_ part of + /// the animation is often interpreted as a thumbnail. + /// + /// This method will return an error when animation control was not configured + /// (which is doing by calling [`set_animated`]). pub fn set_sep_def_img(&mut self, sep_def_img: bool) -> Result<()> { if self.info.animation_control.is_none() { - self.sep_def_img = sep_def_img; + self.options.sep_def_img = sep_def_img; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) @@ -245,14 +269,7 @@ impl<'a, W: Write> Encoder<'a, W> { /// /// The remaining data can be supplied by methods on the returned [`Writer`]. pub fn write_header(self) -> Result> { - Writer::new( - self.w, - PartialInfo::new(&self.info), - self.filter, - self.adaptive_filter, - self.sep_def_img, - ) - .init(&self.info) + Writer::new(self.w, PartialInfo::new(&self.info), self.options).init(&self.info) } /// Set the color of the encoded image. @@ -286,7 +303,7 @@ impl<'a, W: Write> Encoder<'a, W> { /// [`FilterType::Sub`]: enum.FilterType.html#variant.Sub /// [`FilterType::Paeth`]: enum.FilterType.html#variant.Paeth pub fn set_filter(&mut self, filter: FilterType) { - self.filter = filter; + self.options.filter = filter; } /// Set the adaptive filter type. @@ -298,7 +315,7 @@ impl<'a, W: Write> Encoder<'a, W> { /// /// [`AdaptiveFilterType::NonAdaptive`]: enum.AdaptiveFilterType.html pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) { - self.adaptive_filter = adaptive_filter; + self.options.adaptive_filter = adaptive_filter; } /// Set the fraction of time every frame is going to be displayed, in seconds. @@ -420,6 +437,20 @@ impl<'a, W: Write> Encoder<'a, W> { self.info.utf8_text.push(text_chunk); Ok(()) } + + /// Validate the written image sequence. + /// + /// When validation is turned on (it's turned off by default) then attempts to write more than + /// one `IDAT` image or images beyond the number of frames indicated in the animation control + /// chunk will fail and return an error result instead. Attempts to [finish][finish] the image + /// with missing frames will also return an error. + /// + /// [finish]: StreamWriter::finish + /// + /// (It's possible to circumvent these checks by writing raw chunks instead.) + pub fn validate_sequence(&mut self, validate: bool) { + self.options.validate_sequence = validate; + } } /// PNG writer @@ -433,9 +464,7 @@ impl<'a, W: Write> Encoder<'a, W> { pub struct Writer { w: W, info: PartialInfo, - filter: FilterType, - adaptive_filter: AdaptiveFilterType, - sep_def_img: bool, + options: Options, written: u64, animation_written: u32, } @@ -510,19 +539,11 @@ pub(crate) fn write_chunk(mut w: W, name: chunk::ChunkType, data: &[u8 } impl Writer { - fn new( - w: W, - info: PartialInfo, - filter: FilterType, - adaptive_filter: AdaptiveFilterType, - sep_def_img: bool, - ) -> Writer { + fn new(w: W, info: PartialInfo, options: Options) -> Writer { Writer { w, info, - filter, - adaptive_filter, - sep_def_img, + options, written: 0, animation_written: 0, } @@ -575,7 +596,7 @@ impl Writer { fn max_frames(&self) -> u64 { match self.info.animation_control { - Some(a) if self.sep_def_img => a.num_frames as u64 + 1, + Some(a) if self.options.sep_def_img => a.num_frames as u64 + 1, Some(a) => a.num_frames as u64, None => 1, } @@ -622,8 +643,8 @@ impl Writer { self.info.compression.clone().to_options(), ); let bpp = self.info.bpp_in_prediction(); - let filter_method = self.filter; - let adaptive_method = self.adaptive_filter; + let filter_method = self.options.filter; + let adaptive_method = self.options.adaptive_filter; for line in data.chunks(in_len) { current.copy_from_slice(&line); @@ -639,7 +660,7 @@ impl Writer { self.write_zlib_encoded_idat(&zlib_encoded)?; } Some(_) if self.sep_def_img => { - self.sep_def_img = false; + self.options.sep_def_img = false; self.write_zlib_encoded_idat(&zlib_encoded)?; } Some(ref mut fctl) => { @@ -690,7 +711,7 @@ impl Writer { /// [`FilterType::Sub`]: enum.FilterType.html#variant.Sub /// [`FilterType::Paeth`]: enum.FilterType.html#variant.Paeth pub fn set_filter(&mut self, filter: FilterType) { - self.filter = filter; + self.options.filter = filter; } /// Set the adaptive filter type for the following frames. @@ -702,7 +723,7 @@ impl Writer { /// /// [`AdaptiveFilterType::NonAdaptive`]: enum.AdaptiveFilterType.html pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) { - self.adaptive_filter = adaptive_filter; + self.options.adaptive_filter = adaptive_filter; } /// Set the fraction of time the following frames are going to be displayed, @@ -968,7 +989,7 @@ impl<'a, W: Write> ChunkWriter<'a, W> { // is 64 bits. const CAP: usize = std::u32::MAX as usize >> 1; let curr_chunk; - if writer.sep_def_img || writer.info.frame_control.is_none() || writer.written == 0 { + if writer.options.sep_def_img || writer.info.frame_control.is_none() || writer.written == 0 { curr_chunk = chunk::IDAT; } else { curr_chunk = chunk::fdAT; @@ -1013,7 +1034,7 @@ impl<'a, W: Write> ChunkWriter<'a, W> { let wrt = self.writer.deref_mut(); self.curr_chunk = match wrt.info.frame_control { - _ if wrt.sep_def_img => chunk::IDAT, + _ if wrt.options.sep_def_img => chunk::IDAT, None => chunk::IDAT, Some(ref mut fctl) => { fctl.encode(&mut wrt.w)?; @@ -1068,7 +1089,7 @@ impl<'a, W: Write> Write for ChunkWriter<'a, W> { if self.index == 0 { let wrt = self.writer.deref_mut(); // ??? maybe use self.curr_chunk == chunk::fdAT ??? - if !wrt.sep_def_img && wrt.info.frame_control.is_some() && wrt.written > 0 { + if !wrt.options.sep_def_img && wrt.info.frame_control.is_some() && wrt.written > 0 { let fctl = wrt.info.frame_control.as_mut().unwrap(); self.buffer[0..4].copy_from_slice(&fctl.sequence_number.to_be_bytes()); fctl.sequence_number += 1; @@ -1178,8 +1199,8 @@ impl<'a, W: Write> StreamWriter<'a, W> { let bpp = writer.info.bpp_in_prediction(); let in_len = writer.info.raw_row_length() - 1; - let filter = writer.filter; - let adaptive_filter = writer.adaptive_filter; + let filter = writer.options.filter; + let adaptive_filter = writer.options.adaptive_filter; let prev_buf = vec![0; in_len]; let curr_buf = vec![0; in_len]; From 93061ee2b1e329247c53c742e7c0a5bf256fc6bd Mon Sep 17 00:00:00 2001 From: Andreas Molzer Date: Wed, 6 Oct 2021 18:58:51 +0200 Subject: [PATCH 5/7] Simplify tracking of idat/fdAT/fctl We no longer mutate separate default image. It's just the first image that is IDAT, independent of frame control or not. --- src/encoder.rs | 38 ++++++++++++++++++++++++-------------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/src/encoder.rs b/src/encoder.rs index fa9667912e..99ff6f693c 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -659,8 +659,7 @@ impl Writer { None => { self.write_zlib_encoded_idat(&zlib_encoded)?; } - Some(_) if self.sep_def_img => { - self.options.sep_def_img = false; + Some(_) if self.should_skip_frame_control_on_default_image() => { self.write_zlib_encoded_idat(&zlib_encoded)?; } Some(ref mut fctl) => { @@ -695,6 +694,10 @@ impl Writer { Ok(()) } + fn should_skip_frame_control_on_default_image(&self) -> bool { + self.options.sep_def_img && self.written == 0 + } + fn write_zlib_encoded_idat(&mut self, zlib_encoded: &[u8]) -> Result<()> { for chunk in zlib_encoded.chunks(Self::MAX_IDAT_CHUNK_LEN as usize) { self.write_chunk(chunk::IDAT, &chunk)?; @@ -989,7 +992,7 @@ impl<'a, W: Write> ChunkWriter<'a, W> { // is 64 bits. const CAP: usize = std::u32::MAX as usize >> 1; let curr_chunk; - if writer.options.sep_def_img || writer.info.frame_control.is_none() || writer.written == 0 { + if writer.written == 0 { curr_chunk = chunk::IDAT; } else { curr_chunk = chunk::fdAT; @@ -1033,18 +1036,21 @@ impl<'a, W: Write> ChunkWriter<'a, W> { assert_eq!(self.index, 0, "Called when not flushed"); let wrt = self.writer.deref_mut(); - self.curr_chunk = match wrt.info.frame_control { - _ if wrt.options.sep_def_img => chunk::IDAT, - None => chunk::IDAT, + self.curr_chunk = if wrt.written == 0 { + chunk::IDAT + } else { + chunk::fdAT + }; + + match wrt.info.frame_control { + Some(_) if wrt.should_skip_frame_control_on_default_image() => {} Some(ref mut fctl) => { fctl.encode(&mut wrt.w)?; fctl.sequence_number += 1; - match wrt.written { - 0 => chunk::IDAT, - _ => chunk::fdAT, - } } - }; + _ => {} + } + Ok(()) } @@ -1073,6 +1079,7 @@ impl<'a, W: Write> ChunkWriter<'a, W> { self.curr_chunk, &self.buffer[..self.index], )?; + self.index = 0; } Ok(()) @@ -1085,11 +1092,13 @@ impl<'a, W: Write> Write for ChunkWriter<'a, W> { return Ok(0); } - // index == 0 means a chunk as been flushed out + // index == 0 means a chunk has been flushed out if self.index == 0 { let wrt = self.writer.deref_mut(); - // ??? maybe use self.curr_chunk == chunk::fdAT ??? - if !wrt.options.sep_def_img && wrt.info.frame_control.is_some() && wrt.written > 0 { + + // Prepare the next animated frame, if any. + let no_fctl = wrt.should_skip_frame_control_on_default_image(); + if wrt.info.frame_control.is_some() && !no_fctl { let fctl = wrt.info.frame_control.as_mut().unwrap(); self.buffer[0..4].copy_from_slice(&fctl.sequence_number.to_be_bytes()); fctl.sequence_number += 1; @@ -1109,6 +1118,7 @@ impl<'a, W: Write> Write for ChunkWriter<'a, W> { if self.index == self.buffer.len() { self.flush_inner()?; } + Ok(written) } From 5429a4d6be5526f1d9602aee3d5b36457c700f6c Mon Sep 17 00:00:00 2001 From: Andreas Molzer Date: Wed, 6 Oct 2021 22:09:29 +0200 Subject: [PATCH 6/7] Implement validation of image sequence This also provides a `finish` method on the main encoder, which can also catch IO errors happening inside writing of the IEND chunk. --- src/encoder.rs | 196 +++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 191 insertions(+), 5 deletions(-) diff --git a/src/encoder.rs b/src/encoder.rs index 99ff6f693c..c176dfbca7 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -223,7 +223,7 @@ impl<'a, W: Write> Encoder<'a, W> { /// This method will return an error when animation control was not configured /// (which is doing by calling [`set_animated`]). pub fn set_sep_def_img(&mut self, sep_def_img: bool) -> Result<()> { - if self.info.animation_control.is_none() { + if self.info.animation_control.is_some() { self.options.sep_def_img = sep_def_img; Ok(()) } else { @@ -594,11 +594,41 @@ impl Writer { text_chunk.encode(&mut self.w) } - fn max_frames(&self) -> u64 { + /// Check if we should allow writing another image. + fn validate_new_image(&self) -> Result<()> { + if !self.options.validate_sequence { + return Ok(()); + } + match self.info.animation_control { - Some(a) if self.options.sep_def_img => a.num_frames as u64 + 1, - Some(a) => a.num_frames as u64, - None => 1, + None => { + if self.written == 0 { + Ok(()) + } else { + Err(EncodingError::Format(FormatErrorKind::EndReached.into())) + } + } + Some(_) => { + if self.info.frame_control.is_some() { + Ok(()) + } else { + Err(EncodingError::Format(FormatErrorKind::EndReached.into())) + } + } + } + } + + fn validate_sequence_done(&self) -> Result<()> { + if !self.options.validate_sequence { + return Ok(()); + } + + if self.info.animation_control.is_some() && self.info.frame_control.is_some() { + Err(EncodingError::Format(FormatErrorKind::MissingFrames.into())) + } else if self.written == 0 { + Err(EncodingError::Format(FormatErrorKind::MissingFrames.into())) + } else { + Ok(()) } } @@ -612,6 +642,8 @@ impl Writer { return Err(EncodingError::Format(FormatErrorKind::NoPalette.into())); } + self.validate_new_image()?; + let width: usize; let height: usize; if let Some(ref mut fctl) = self.info.frame_control { @@ -926,6 +958,18 @@ impl Writer { pub fn into_stream_writer_with_size(self, size: usize) -> Result> { StreamWriter::new(ChunkOutput::Owned(self), size) } + + /// Consume the stream writer with validation. + /// + /// Unlike a simple drop this ensures that the final chunk was written correctly. When other + /// validation options (chunk sequencing) had been turned on in the configuration then it will + /// also do a check on their correctness _before_ writing the final chunk. + pub fn finish(self) -> Result<()> { + self.validate_sequence_done()?; + // Prevent the writing on Drop, we do it fallibly. + let mut finalize = mem::ManuallyDrop::new(self); + finalize.write_chunk(chunk::IEND, &[]) + } } impl Drop for Writer { @@ -1423,6 +1467,12 @@ impl<'a, W: Write> StreamWriter<'a, W> { // TODO: call `writer.finish` somehow? self.flush()?; + + match self.writer.take() { + Wrapper::Chunk(wrt) => wrt.writer.validate_sequence_done()?, + _ => unreachable!(), + } + Ok(()) } @@ -1435,6 +1485,8 @@ impl<'a, W: Write> StreamWriter<'a, W> { _ => unreachable!(), }; wrt.flush()?; + wrt.writer.validate_new_image()?; + if let Some(fctl) = self.fctl { wrt.set_fctl(fctl); } @@ -1950,6 +2002,140 @@ mod tests { Ok(()) } + #[test] + fn image_validate_sequence_without_animation() -> Result<()> { + let width = 10; + let height = 10; + + let output = vec![0u8; 1024]; + let writer = Cursor::new(output); + + let mut encoder = Encoder::new(writer, width, height); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + encoder.validate_sequence(true); + let mut png_writer = encoder.write_header()?; + + let correct_image_size = (width * height) as usize; + let image = vec![0u8; correct_image_size]; + png_writer.write_image_data(image.as_ref())?; + + assert!(png_writer.write_image_data(image.as_ref()).is_err()); + Ok(()) + } + + #[test] + fn image_validate_animation() -> Result<()> { + let width = 10; + let height = 10; + + let output = vec![0u8; 1024]; + let writer = Cursor::new(output); + let correct_image_size = (width * height) as usize; + let image = vec![0u8; correct_image_size]; + + let mut encoder = Encoder::new(writer, width, height); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + encoder.set_animated(1, 0)?; + encoder.validate_sequence(true); + let mut png_writer = encoder.write_header()?; + + png_writer.write_image_data(image.as_ref())?; + + Ok(()) + } + + #[test] + fn image_validate_animation2() -> Result<()> { + let width = 10; + let height = 10; + + let output = vec![0u8; 1024]; + let writer = Cursor::new(output); + let correct_image_size = (width * height) as usize; + let image = vec![0u8; correct_image_size]; + + let mut encoder = Encoder::new(writer, width, height); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + encoder.set_animated(2, 0)?; + encoder.validate_sequence(true); + let mut png_writer = encoder.write_header()?; + + png_writer.write_image_data(image.as_ref())?; + png_writer.write_image_data(image.as_ref())?; + png_writer.finish()?; + + Ok(()) + } + + #[test] + fn image_validate_animation_sep_def_image() -> Result<()> { + let width = 10; + let height = 10; + + let output = vec![0u8; 1024]; + let writer = Cursor::new(output); + let correct_image_size = (width * height) as usize; + let image = vec![0u8; correct_image_size]; + + let mut encoder = Encoder::new(writer, width, height); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + encoder.set_animated(1, 0)?; + encoder.set_sep_def_img(true)?; + encoder.validate_sequence(true); + let mut png_writer = encoder.write_header()?; + + png_writer.write_image_data(image.as_ref())?; + png_writer.write_image_data(image.as_ref())?; + png_writer.finish()?; + + Ok(()) + } + + #[test] + fn image_validate_missing_image() -> Result<()> { + let width = 10; + let height = 10; + + let output = vec![0u8; 1024]; + let writer = Cursor::new(output); + + let mut encoder = Encoder::new(writer, width, height); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + encoder.validate_sequence(true); + let png_writer = encoder.write_header()?; + + assert!(png_writer.finish().is_err()); + Ok(()) + } + + #[test] + fn image_validate_missing_animated_frame() -> Result<()> { + let width = 10; + let height = 10; + + let output = vec![0u8; 1024]; + let writer = Cursor::new(output); + let correct_image_size = (width * height) as usize; + let image = vec![0u8; correct_image_size]; + + let mut encoder = Encoder::new(writer, width, height); + encoder.set_depth(BitDepth::Eight); + encoder.set_color(ColorType::Grayscale); + encoder.set_animated(2, 0)?; + encoder.validate_sequence(true); + let mut png_writer = encoder.write_header()?; + + png_writer.write_image_data(image.as_ref())?; + assert!(png_writer.finish().is_err()); + + Ok(()) + } + #[test] #[cfg(all(unix, not(target_pointer_width = "32")))] fn exper_error_on_huge_chunk() -> Result<()> { From 57d1549ac4e5f4a77fcffd4aded6b06a68b11c5d Mon Sep 17 00:00:00 2001 From: Andreas Molzer Date: Thu, 7 Oct 2021 18:43:42 +0200 Subject: [PATCH 7/7] Add test case for issue 307 --- CHANGES.md | 3 +++ src/encoder.rs | 61 ++++++++++++++++++++++++++++++++++++++++++++------ 2 files changed, 57 insertions(+), 7 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 31b3fe0949..f98352c415 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,6 +1,9 @@ ## Unreleased * Added support for encoding and decoding tEXt/zTXt/iTXt chunks. +* Added ability to make frame sequence validation optional. +* Fix an issue where, after `flush` or `write` of an underlying writer + panicked, the library would abort. ## 0.17.1 diff --git a/src/encoder.rs b/src/encoder.rs index c176dfbca7..b3e9e42ad3 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -1470,7 +1470,7 @@ impl<'a, W: Write> StreamWriter<'a, W> { match self.writer.take() { Wrapper::Chunk(wrt) => wrt.writer.validate_sequence_done()?, - _ => unreachable!(), + _ => {} } Ok(()) @@ -1479,10 +1479,16 @@ impl<'a, W: Write> StreamWriter<'a, W> { /// Flushes the buffered chunk, checks if it was the last frame, /// writes the next frame header and gets the next frame scanline size /// and image size. + /// NOTE: This method must only be called when the writer is the variant Chunk(_) fn new_frame(&mut self) -> Result<()> { let wrt = match &mut self.writer { Wrapper::Chunk(wrt) => wrt, - _ => unreachable!(), + Wrapper::Unrecoverable => { + let err = FormatErrorKind::Unrecoverable.into(); + return Err(EncodingError::Format(err).into()); + } + Wrapper::Zlib(_) => unreachable!("never called on a half-finished frame"), + Wrapper::None => unreachable!(), }; wrt.flush()?; wrt.writer.validate_new_image()?; @@ -1497,11 +1503,14 @@ impl<'a, W: Write> StreamWriter<'a, W> { wrt.write_header()?; // now it can be taken because the next statements cannot cause any errors - let wrt = match self.writer.take() { - Wrapper::Chunk(wrt) => wrt, + match self.writer.take() { + Wrapper::Chunk(wrt) => { + let encoder = ZlibEncoder::new(wrt, self.compression.to_options()); + self.writer = Wrapper::Zlib(encoder); + } _ => unreachable!(), }; - self.writer = Wrapper::Zlib(ZlibEncoder::new(wrt, self.compression.to_options())); + Ok(()) } } @@ -1527,8 +1536,11 @@ impl<'a, W: Write> Write for StreamWriter<'a, W> { } }, chunk @ Wrapper::Chunk(_) => self.writer = chunk, - Wrapper::None | Wrapper::Unrecoverable => unreachable!(), + Wrapper::Unrecoverable => unreachable!(), + Wrapper::None => unreachable!(), }; + + // Transition Wrapper::Chunk to Wrapper::Zlib. self.new_frame()?; } @@ -1549,23 +1561,33 @@ impl<'a, W: Write> Write for StreamWriter<'a, W> { Wrapper::Zlib(wrt) => wrt, _ => unreachable!(), }; + wrt.write_all(&[filter_type as u8])?; wrt.write_all(&self.curr_buf)?; mem::swap(&mut self.prev_buf, &mut self.curr_buf); self.index = 0; } + Ok(written) } fn flush(&mut self) -> io::Result<()> { match &mut self.writer { Wrapper::Zlib(wrt) => wrt.flush()?, - _ => unreachable!(), + Wrapper::Chunk(wrt) => wrt.flush()?, + // This handles both the case where we entered an unrecoverable state after zlib + // decoding failure and after a panic while we had taken the chunk/zlib reader. + Wrapper::Unrecoverable | Wrapper::None => { + let err = FormatErrorKind::Unrecoverable.into(); + return Err(EncodingError::Format(err).into()); + } } + if self.index > 0 { let err = FormatErrorKind::WrittenTooMuch(self.index).into(); return Err(EncodingError::Format(err).into()); } + Ok(()) } } @@ -2136,6 +2158,31 @@ mod tests { Ok(()) } + #[test] + fn issue_307_stream_validation() -> Result<()> { + let output = vec![0u8; 1024]; + let mut cursor = Cursor::new(output); + + let encoder = Encoder::new(&mut cursor, 1, 1); // Create a 1-pixel image + let mut writer = encoder.write_header()?; + let mut stream = writer.stream_writer()?; + + let written = stream.write(&[1, 2, 3, 4])?; + assert_eq!(written, 1); + stream.finish()?; + drop(writer); + + { + cursor.set_position(0); + let mut decoder = Decoder::new(cursor).read_info().expect("A valid image"); + let mut buffer = [0u8; 1]; + decoder.next_frame(&mut buffer[..]).expect("Valid read"); + assert_eq!(buffer, [1]); + } + + Ok(()) + } + #[test] #[cfg(all(unix, not(target_pointer_width = "32")))] fn exper_error_on_huge_chunk() -> Result<()> {