use borrow::Cow; use io::{Read, Write}; use ops::{Deref, DerefMut}; use std::{borrow, error, fmt, io, mem, ops, result}; use crc32fast::Hasher as Crc32; use flate2::write::ZlibEncoder; use crate::chunk::{self, ChunkType}; use crate::common::{ AnimationControl, BitDepth, BlendOp, BytesPerPixel, ColorType, Compression, DisposeOp, FrameControl, Info, ParameterError, ParameterErrorKind, PixelDimensions, ScaledFloat, }; use crate::filter::{filter, AdaptiveFilterType, FilterType}; use crate::text_metadata::{ EncodableTextChunk, ITXtChunk, TEXtChunk, TextEncodingError, ZTXtChunk, }; use crate::traits::WriteBytesExt; pub type Result = result::Result; #[derive(Debug)] pub enum EncodingError { IoError(io::Error), Format(FormatError), Parameter(ParameterError), LimitsExceeded, } #[derive(Debug)] pub struct FormatError { inner: FormatErrorKind, } #[derive(Debug)] enum FormatErrorKind { ZeroWidth, ZeroHeight, InvalidColorCombination(BitDepth, ColorType), NoPalette, // TODO: wait, what? WrittenTooMuch(usize), NotAnimated, OutOfBounds, EndReached, ZeroFrames, MissingFrames, MissingData(usize), Unrecoverable, BadTextEncoding(TextEncodingError), } impl error::Error for EncodingError { fn cause(&self) -> Option<&(dyn error::Error + 'static)> { match self { EncodingError::IoError(err) => Some(err), _ => None, } } } impl fmt::Display for EncodingError { fn fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error> { use self::EncodingError::*; match self { IoError(err) => write!(fmt, "{}", err), Format(desc) => write!(fmt, "{}", desc), Parameter(desc) => write!(fmt, "{}", desc), LimitsExceeded => write!(fmt, "Limits are exceeded."), } } } impl fmt::Display for FormatError { fn fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error> { use FormatErrorKind::*; match self.inner { ZeroWidth => write!(fmt, "Zero width not allowed"), ZeroHeight => write!(fmt, "Zero height not allowed"), ZeroFrames => write!(fmt, "Zero frames not allowed"), InvalidColorCombination(depth, color) => write!( fmt, "Invalid combination of bit-depth '{:?}' and color-type '{:?}'", depth, color ), NoPalette => write!(fmt, "can't write indexed image without palette"), WrittenTooMuch(index) => write!(fmt, "wrong data size, got {} bytes too many", index), NotAnimated => write!(fmt, "not an animation"), OutOfBounds => write!( fmt, "the dimension and position go over the frame boundaries" ), EndReached => write!(fmt, "all the frames have been already written"), MissingFrames => write!(fmt, "there are still frames to be written"), MissingData(n) => write!(fmt, "there are still {} bytes to be written", n), Unrecoverable => write!( fmt, "a previous error put the writer into an unrecoverable state" ), BadTextEncoding(tee) => match tee { TextEncodingError::Unrepresentable => write!( fmt, "The text metadata cannot be encoded into valid ISO 8859-1" ), TextEncodingError::InvalidKeywordSize => write!(fmt, "Invalid keyword size"), TextEncodingError::CompressionError => { write!(fmt, "Unable to compress text metadata") } }, } } } impl From for EncodingError { fn from(err: io::Error) -> EncodingError { EncodingError::IoError(err) } } impl From for io::Error { fn from(err: EncodingError) -> io::Error { io::Error::new(io::ErrorKind::Other, err.to_string()) } } // Private impl. impl From for FormatError { fn from(kind: FormatErrorKind) -> Self { FormatError { inner: kind } } } impl From for EncodingError { fn from(tee: TextEncodingError) -> Self { EncodingError::Format(FormatError { inner: FormatErrorKind::BadTextEncoding(tee), }) } } /// 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>, options: Options, } /// Decoding options, internal type, forwarded to the Writer. #[derive(Default)] struct Options { filter: FilterType, adaptive_filter: AdaptiveFilterType, sep_def_img: bool, validate_sequence: bool, } impl<'a, W: Write> Encoder<'a, W> { pub fn new(w: W, width: u32, height: u32) -> Encoder<'static, W> { Encoder { w, info: Info::with_size(width, height), options: Options::default(), } } /// Specify that the image is animated. /// /// `num_frames` controls how many frames the animation has, while /// `num_plays` controls how many times the animation should be /// 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 [`Encoder::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(()) } /// 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 done by calling [`Encoder::set_animated`]). pub fn set_sep_def_img(&mut self, sep_def_img: bool) -> Result<()> { if self.info.animation_control.is_some() { self.options.sep_def_img = sep_def_img; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Sets the raw byte contents of the PLTE chunk. This method accepts /// both borrowed and owned byte data. pub fn set_palette>>(&mut self, palette: T) { self.info.palette = Some(palette.into()); } /// Sets the raw byte contents of the tRNS chunk. This method accepts /// both borrowed and owned byte data. pub fn set_trns>>(&mut self, trns: T) { self.info.trns = Some(trns.into()); } /// Set the display gamma of the source system on which the image was generated or last edited. pub fn set_source_gamma(&mut self, source_gamma: ScaledFloat) { self.info.source_gamma = Some(source_gamma); } /// Set the chromaticities for the source system's display channels (red, green, blue) and the whitepoint /// of the source system on which the image was generated or last edited. pub fn set_source_chromaticities( &mut self, source_chromaticities: super::SourceChromaticities, ) { self.info.source_chromaticities = Some(source_chromaticities); } /// Mark the image data as conforming to the SRGB color space with the specified rendering intent. /// /// Matching source gamma and chromaticities chunks are added automatically. /// Any manually specified source gamma or chromaticities will be ignored. pub fn set_srgb(&mut self, rendering_intent: super::SrgbRenderingIntent) { 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, PartialInfo::new(&self.info), self.options).init(&self.info) } /// Set the color of the encoded image. /// /// These correspond to the color types in the png IHDR data that will be written. The length /// of the image data that is later supplied must match the color type, otherwise an error will /// be emitted. pub fn set_color(&mut self, color: ColorType) { self.info.color_type = color; } /// Set the indicated depth of the image data. pub fn set_depth(&mut self, depth: BitDepth) { self.info.bit_depth = depth; } /// Set compression parameters. /// /// Accepts a `Compression` or any type that can transform into a `Compression`. Notably `deflate::Compression` and /// `deflate::CompressionOptions` which "just work". pub fn set_compression(&mut self, compression: Compression) { self.info.compression = compression; } /// Set the used filter type. /// /// The default filter is [`FilterType::Sub`] which provides a basic prediction algorithm for /// sample values based on the previous. For a potentially better compression ratio, at the /// cost of more complex processing, try out [`FilterType::Paeth`]. /// /// [`FilterType::Sub`]: enum.FilterType.html#variant.Sub /// [`FilterType::Paeth`]: enum.FilterType.html#variant.Paeth pub fn set_filter(&mut self, filter: FilterType) { self.options.filter = filter; } /// Set the adaptive filter type. /// /// Adaptive filtering attempts to select the best filter for each line /// based on heuristics which minimize the file size for compression rather /// than use a single filter for the entire image. The default method is /// [`AdaptiveFilterType::NonAdaptive`]. /// /// [`AdaptiveFilterType::NonAdaptive`]: enum.AdaptiveFilterType.html pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) { self.options.adaptive_filter = adaptive_filter; } /// Set the fraction of time every frame is going to be displayed, in seconds. /// /// *Note that this parameter can be set for each individual frame after /// [`Encoder::write_header`] is called. (see [`Writer::set_frame_delay`])* /// /// If the denominator is 0, it is to be treated as if it were 100 /// (that is, the numerator then specifies 1/100ths of a second). /// If the the value of the numerator is 0 the decoder should render the next frame /// as quickly as possible, though viewers may impose a reasonable lower bound. /// /// The default value is 0 for both the numerator and denominator. /// /// This method will return an error if the image is not animated. /// (see [`set_animated`]) /// /// [`write_header`]: struct.Encoder.html#method.write_header /// [`set_animated`]: struct.Encoder.html#method.set_animated /// [`Writer::set_frame_delay`]: struct.Writer#method.set_frame_delay pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { fctl.delay_den = denominator; fctl.delay_num = numerator; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the blend operation for every frame. /// /// The blend operation specifies whether the frame is to be alpha blended /// into the current output buffer content, or whether it should completely /// replace its region in the output buffer. /// /// *Note that this parameter can be set for each individual frame after /// [`write_header`] is called. (see [`Writer::set_blend_op`])* /// /// See the [`BlendOp`] documentation for the possible values and their effects. /// /// *Note that for the first frame the two blend modes are functionally /// equivalent due to the clearing of the output buffer at the beginning /// of each play.* /// /// The default value is [`BlendOp::Source`]. /// /// This method will return an error if the image is not animated. /// (see [`set_animated`]) /// /// [`BlendOP`]: enum.BlendOp.html /// [`BlendOP::Source`]: enum.BlendOp.html#variant.Source /// [`write_header`]: struct.Encoder.html#method.write_header /// [`set_animated`]: struct.Encoder.html#method.set_animated /// [`Writer::set_blend_op`]: struct.Writer#method.set_blend_op pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { fctl.blend_op = op; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the dispose operation for every frame. /// /// The dispose operation specifies how the output buffer should be changed /// at the end of the delay (before rendering the next frame) /// /// *Note that this parameter can be set for each individual frame after /// [`write_header`] is called (see [`Writer::set_dispose_op`])* /// /// See the [`DisposeOp`] documentation for the possible values and their effects. /// /// *Note that if the first frame uses [`DisposeOp::Previous`] /// it will be treated as [`DisposeOp::Background`].* /// /// The default value is [`DisposeOp::None`]. /// /// This method will return an error if the image is not animated. /// (see [`set_animated`]) /// /// [`DisposeOp`]: ../common/enum.BlendOp.html /// [`DisposeOp::Previous`]: ../common/enum.BlendOp.html#variant.Previous /// [`DisposeOp::Background`]: ../common/enum.BlendOp.html#variant.Background /// [`DisposeOp::None`]: ../common/enum.BlendOp.html#variant.None /// [`write_header`]: struct.Encoder.html#method.write_header /// [`set_animated`]: struct.Encoder.html#method.set_animated /// [`Writer::set_dispose_op`]: struct.Writer#method.set_dispose_op pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { fctl.dispose_op = op; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } pub fn set_pixel_dims(&mut self, pixel_dims: Option) { self.info.pixel_dims = pixel_dims } /// Convenience function to add tEXt chunks to [`Info`] struct pub fn add_text_chunk(&mut self, keyword: String, text: String) -> Result<()> { let text_chunk = TEXtChunk::new(keyword, text); self.info.uncompressed_latin1_text.push(text_chunk); Ok(()) } /// Convenience function to add zTXt chunks to [`Info`] struct pub fn add_ztxt_chunk(&mut self, keyword: String, text: String) -> Result<()> { let text_chunk = ZTXtChunk::new(keyword, text); self.info.compressed_latin1_text.push(text_chunk); Ok(()) } /// Convenience function to add iTXt chunks to [`Info`] struct /// /// This function only sets the `keyword` and `text` field of the iTXt chunk. /// To set the other fields, create a [`ITXtChunk`] directly, and then encode it to the output stream. pub fn add_itxt_chunk(&mut self, keyword: String, text: String) -> Result<()> { let text_chunk = ITXtChunk::new(keyword, text); 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 /// /// 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 { /// The underlying writer. w: W, /// The local version of the `Info` struct. info: PartialInfo, /// Global encoding options. options: Options, /// The total number of image frames, counting all consecutive IDAT and fdAT chunks. images_written: u64, /// The total number of animation frames, that is equivalent to counting fcTL chunks. animation_written: u32, /// A flag to note when the IEND chunk was already added. /// This is only set on code paths that drop `Self` to control the destructor. iend_written: bool, } /// Contains the subset of attributes of [Info] needed for [Writer] to function struct PartialInfo { width: u32, height: u32, bit_depth: BitDepth, color_type: ColorType, frame_control: Option, animation_control: Option, compression: Compression, has_palette: bool, } impl PartialInfo { fn new(info: &Info) -> Self { PartialInfo { width: info.width, height: info.height, bit_depth: info.bit_depth, color_type: info.color_type, frame_control: info.frame_control, animation_control: info.animation_control, compression: info.compression, has_palette: info.palette.is_some(), } } fn bpp_in_prediction(&self) -> BytesPerPixel { // Passthrough self.to_info().bpp_in_prediction() } fn raw_row_length(&self) -> usize { // Passthrough self.to_info().raw_row_length() } fn raw_row_length_from_width(&self, width: u32) -> usize { // Passthrough self.to_info().raw_row_length_from_width(width) } /// Converts this partial info to an owned Info struct, /// setting missing values to their defaults fn to_info(&self) -> Info<'static> { Info { width: self.width, height: self.height, bit_depth: self.bit_depth, color_type: self.color_type, frame_control: self.frame_control, animation_control: self.animation_control, compression: self.compression, ..Default::default() } } } const DEFAULT_BUFFER_LENGTH: usize = 4 * 1024; pub(crate) fn write_chunk(mut w: W, name: chunk::ChunkType, data: &[u8]) -> Result<()> { w.write_be(data.len() as u32)?; w.write_all(&name.0)?; w.write_all(data)?; let mut crc = Crc32::new(); crc.update(&name.0); crc.update(data); w.write_be(crc.finalize())?; Ok(()) } impl Writer { fn new(w: W, info: PartialInfo, options: Options) -> Writer { Writer { w, info, options, images_written: 0, animation_written: 0, iend_written: false, } } fn init(mut self, info: &Info<'_>) -> Result { if self.info.width == 0 { return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into())); } if self.info.height == 0 { return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into())); } if self .info .color_type .is_combination_invalid(self.info.bit_depth) { return Err(EncodingError::Format( FormatErrorKind::InvalidColorCombination(self.info.bit_depth, self.info.color_type) .into(), )); } self.w.write_all(&[137, 80, 78, 71, 13, 10, 26, 10])?; // PNG signature info.encode(&mut self.w)?; 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(true, |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) } pub fn write_text_chunk(&mut self, text_chunk: &T) -> Result<()> { text_chunk.encode(&mut self.w) } /// 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 { None => { if self.images_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()) || self.images_written == 0 { Err(EncodingError::Format(FormatErrorKind::MissingFrames.into())) } else { Ok(()) } } 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())); } self.validate_new_image()?; let width: usize; let height: usize; if let Some(ref mut fctl) = self.info.frame_control { width = fctl.width as usize; height = fctl.height as usize; } else { width = self.info.width as usize; height = self.info.height as usize; } let in_len = self.info.raw_row_length_from_width(width as u32) - 1; let data_size = in_len * height; if data_size != data.len() { return Err(EncodingError::Parameter( ParameterErrorKind::ImageBufferSize { expected: data_size, actual: data.len(), } .into(), )); } let prev = vec![0; in_len]; let mut prev = prev.as_slice(); let bpp = self.info.bpp_in_prediction(); let filter_method = self.options.filter; let adaptive_method = self.options.adaptive_filter; let zlib_encoded = match self.info.compression { Compression::Fast => { let mut compressor = fdeflate::Compressor::new(std::io::Cursor::new(Vec::new()))?; let mut current = vec![0; in_len + 1]; for line in data.chunks(in_len) { let filter_type = filter( filter_method, adaptive_method, bpp, prev, line, &mut current[1..], ); current[0] = filter_type as u8; compressor.write_data(¤t)?; prev = line; } let compressed = compressor.finish()?.into_inner(); if compressed.len() > fdeflate::StoredOnlyCompressor::<()>::compressed_size((in_len + 1) * height) { // Write uncompressed data since the result from fast compression would take // more space than that. // // We always use FilterType::NoFilter here regardless of the filter method // requested by the user. Doing filtering again would only add performance // cost for both encoding and subsequent decoding, without improving the // compression ratio. let mut compressor = fdeflate::StoredOnlyCompressor::new(std::io::Cursor::new(Vec::new()))?; for line in data.chunks(in_len) { compressor.write_data(&[0])?; compressor.write_data(line)?; } compressor.finish()?.into_inner() } else { compressed } } _ => { let mut current = vec![0; in_len]; let mut zlib = ZlibEncoder::new(Vec::new(), self.info.compression.to_options()); for line in data.chunks(in_len) { let filter_type = filter( filter_method, adaptive_method, bpp, prev, line, &mut current, ); zlib.write_all(&[filter_type as u8])?; zlib.write_all(¤t)?; prev = line; } zlib.finish()? } }; match self.info.frame_control { None => { self.write_zlib_encoded_idat(&zlib_encoded)?; } Some(_) if self.should_skip_frame_control_on_default_image() => { 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.images_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); } } } } self.increment_images_written(); Ok(()) } fn increment_images_written(&mut self) { self.images_written = self.images_written.saturating_add(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; } } } fn write_iend(&mut self) -> Result<()> { self.iend_written = true; self.write_chunk(chunk::IEND, &[]) } fn should_skip_frame_control_on_default_image(&self) -> bool { self.options.sep_def_img && self.images_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)?; } Ok(()) } /// Set the used filter type for the following frames. /// /// The default filter is [`FilterType::Sub`] which provides a basic prediction algorithm for /// sample values based on the previous. For a potentially better compression ratio, at the /// cost of more complex processing, try out [`FilterType::Paeth`]. /// /// [`FilterType::Sub`]: enum.FilterType.html#variant.Sub /// [`FilterType::Paeth`]: enum.FilterType.html#variant.Paeth pub fn set_filter(&mut self, filter: FilterType) { self.options.filter = filter; } /// Set the adaptive filter type for the following frames. /// /// Adaptive filtering attempts to select the best filter for each line /// based on heuristics which minimize the file size for compression rather /// than use a single filter for the entire image. The default method is /// [`AdaptiveFilterType::NonAdaptive`]. /// /// [`AdaptiveFilterType::NonAdaptive`]: enum.AdaptiveFilterType.html pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) { self.options.adaptive_filter = adaptive_filter; } /// Set the fraction of time the following frames are going to be displayed, /// in seconds /// /// If the denominator is 0, it is to be treated as if it were 100 /// (that is, the numerator then specifies 1/100ths of a second). /// If the the value of the numerator is 0 the decoder should render the next frame /// as quickly as possible, though viewers may impose a reasonable lower bound. /// /// This method will return an error if the image is not animated. pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { fctl.delay_den = denominator; fctl.delay_num = numerator; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the dimension of the following frames. /// /// This function will return an error when: /// - The image is not an animated; /// /// - The selected dimension, considering also the current frame position, /// goes outside the image boundaries; /// /// - One or both the width and height are 0; /// // ??? TODO ??? // - The next frame is the default image pub fn set_frame_dimension(&mut self, width: u32, height: u32) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { if Some(width) > self.info.width.checked_sub(fctl.x_offset) || Some(height) > self.info.height.checked_sub(fctl.y_offset) { return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into())); } else if width == 0 { return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into())); } else if height == 0 { return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into())); } fctl.width = width; fctl.height = height; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the position of the following frames. /// /// An error will be returned if: /// - The image is not animated; /// /// - The selected position, considering also the current frame dimension, /// goes outside the image boundaries; /// // ??? TODO ??? // - The next frame is the default image pub fn set_frame_position(&mut self, x: u32, y: u32) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { if Some(x) > self.info.width.checked_sub(fctl.width) || Some(y) > self.info.height.checked_sub(fctl.height) { return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into())); } fctl.x_offset = x; fctl.y_offset = y; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the frame dimension to occupy all the image, starting from /// the current position. /// /// To reset the frame to the full image size [`reset_frame_position`] /// should be called first. /// /// This method will return an error if the image is not animated. /// /// [`reset_frame_position`]: struct.Writer.html#method.reset_frame_position pub fn reset_frame_dimension(&mut self) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { fctl.width = self.info.width - fctl.x_offset; fctl.height = self.info.height - fctl.y_offset; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the frame position to (0, 0). /// /// Equivalent to calling [`set_frame_position(0, 0)`]. /// /// This method will return an error if the image is not animated. /// /// [`set_frame_position(0, 0)`]: struct.Writer.html#method.set_frame_position pub fn reset_frame_position(&mut self) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { fctl.x_offset = 0; fctl.y_offset = 0; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the blend operation for the following frames. /// /// The blend operation specifies whether the frame is to be alpha blended /// into the current output buffer content, or whether it should completely /// replace its region in the output buffer. /// /// See the [`BlendOp`] documentation for the possible values and their effects. /// /// *Note that for the first frame the two blend modes are functionally /// equivalent due to the clearing of the output buffer at the beginning /// of each play.* /// /// This method will return an error if the image is not animated. /// /// [`BlendOP`]: enum.BlendOp.html pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { fctl.blend_op = op; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the dispose operation for the following frames. /// /// The dispose operation specifies how the output buffer should be changed /// at the end of the delay (before rendering the next frame) /// /// See the [`DisposeOp`] documentation for the possible values and their effects. /// /// *Note that if the first frame uses [`DisposeOp::Previous`] /// it will be treated as [`DisposeOp::Background`].* /// /// This method will return an error if the image is not animated. /// /// [`DisposeOp`]: ../common/enum.BlendOp.html /// [`DisposeOp::Previous`]: ../common/enum.BlendOp.html#variant.Previous /// [`DisposeOp::Background`]: ../common/enum.BlendOp.html#variant.Background pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> { if let Some(ref mut fctl) = self.info.frame_control { fctl.dispose_op = op; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Create a stream writer. /// /// This allows you to create images that do not fit in memory. The default /// chunk size is 4K, use `stream_writer_with_size` to set another chunk /// size. /// /// This borrows the writer which allows for manually appending additional /// chunks after the image data has been written. pub fn stream_writer(&mut self) -> Result> { self.stream_writer_with_size(DEFAULT_BUFFER_LENGTH) } /// Create a stream writer with custom buffer size. /// /// See [`stream_writer`]. /// /// [`stream_writer`]: #fn.stream_writer pub fn stream_writer_with_size(&mut self, size: usize) -> Result> { StreamWriter::new(ChunkOutput::Borrowed(self), size) } /// Turn this into a stream writer for image data. /// /// This allows you to create images that do not fit in memory. The default /// chunk size is 4K, use `stream_writer_with_size` to set another chunk /// size. pub fn into_stream_writer(self) -> Result> { self.into_stream_writer_with_size(DEFAULT_BUFFER_LENGTH) } /// Turn this into a stream writer with custom buffer size. /// /// See [`into_stream_writer`]. /// /// [`into_stream_writer`]: #fn.into_stream_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(mut self) -> Result<()> { self.validate_sequence_done()?; self.write_iend()?; self.w.flush()?; // Explicitly drop `self` just for clarity. drop(self); Ok(()) } } impl Drop for Writer { fn drop(&mut self) { if !self.iend_written { let _ = self.write_iend(); } } } enum ChunkOutput<'a, W: Write> { Borrowed(&'a mut Writer), Owned(Writer), } // opted for deref for practical reasons impl<'a, W: Write> Deref for ChunkOutput<'a, W> { type Target = Writer; fn deref(&self) -> &Self::Target { match self { ChunkOutput::Borrowed(writer) => writer, ChunkOutput::Owned(writer) => writer, } } } impl<'a, W: Write> DerefMut for ChunkOutput<'a, W> { fn deref_mut(&mut self) -> &mut Self::Target { match self { ChunkOutput::Borrowed(writer) => writer, ChunkOutput::Owned(writer) => writer, } } } /// This writer is used between the actual writer and the /// ZlibEncoder and has the job of packaging the compressed /// data into a PNG chunk, based on the image metadata /// /// Currently the way it works is that the specified buffer /// will hold one chunk at the time and buffer the incoming /// data until `flush` is called or the maximum chunk size /// is reached. /// /// The maximum chunk is the smallest between the selected buffer size /// and `u32::MAX >> 1` (`0x7fffffff` or `2147483647` dec) /// /// When a chunk has to be flushed the length (that is now known) /// and the CRC will be written at the correct locations in the chunk. struct ChunkWriter<'a, W: Write> { writer: ChunkOutput<'a, W>, buffer: Vec, /// keeps track of where the last byte was written index: usize, curr_chunk: ChunkType, } impl<'a, W: Write> ChunkWriter<'a, W> { fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> ChunkWriter<'a, W> { // currently buf_len will determine the size of each chunk // the len is capped to the maximum size every chunk can hold // (this wont ever overflow an u32) // // TODO (maybe): find a way to hold two chunks at a time if `usize` // is 64 bits. const CAP: usize = std::u32::MAX as usize >> 1; let curr_chunk = if writer.images_written == 0 { chunk::IDAT } else { chunk::fdAT }; ChunkWriter { writer, buffer: vec![0; CAP.min(buf_len)], index: 0, curr_chunk, } } /// Returns the size of each scanline for the next frame /// paired with the size of the whole frame /// /// This is used by the `StreamWriter` to know when the scanline ends /// so it can filter compress it and also to know when to start /// the next one fn next_frame_info(&self) -> (usize, usize) { let wrt = self.writer.deref(); let width: usize; let height: usize; if let Some(fctl) = wrt.info.frame_control { width = fctl.width as usize; height = fctl.height as usize; } else { width = wrt.info.width as usize; height = wrt.info.height as usize; } let in_len = wrt.info.raw_row_length_from_width(width as u32) - 1; let data_size = in_len * height; (in_len, data_size) } /// NOTE: this bypasses the internal buffer so the flush method should be called before this /// in the case there is some data left in the buffer when this is called, it will panic fn write_header(&mut self) -> Result<()> { assert_eq!(self.index, 0, "Called when not flushed"); let wrt = self.writer.deref_mut(); self.curr_chunk = if wrt.images_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; } _ => {} } Ok(()) } /// Set the `FrameControl` for the following frame /// /// It will ignore the `sequence_number` of the parameter /// as it is updated internally. fn set_fctl(&mut self, f: FrameControl) { if let Some(ref mut fctl) = self.writer.info.frame_control { // Ignore the sequence number *fctl = FrameControl { sequence_number: fctl.sequence_number, ..f }; } else { panic!("This function must be called on an animated PNG") } } /// Flushes the current chunk fn flush_inner(&mut self) -> io::Result<()> { if self.index > 0 { // flush the chunk and reset everything write_chunk( &mut self.writer.w, self.curr_chunk, &self.buffer[..self.index], )?; self.index = 0; } Ok(()) } } impl<'a, W: Write> Write for ChunkWriter<'a, W> { fn write(&mut self, mut data: &[u8]) -> io::Result { if data.is_empty() { return Ok(0); } // index == 0 means a chunk has been flushed out if self.index == 0 { let wrt = self.writer.deref_mut(); // 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; self.index = 4; } } // Cap the buffer length to the maximum number of bytes that can't still // be added to the current chunk let written = data.len().min(self.buffer.len() - self.index); data = &data[..written]; self.buffer[self.index..][..written].copy_from_slice(data); self.index += written; // if the maximum data for this chunk as been reached it needs to be flushed if self.index == self.buffer.len() { self.flush_inner()?; } Ok(written) } fn flush(&mut self) -> io::Result<()> { self.flush_inner() } } impl Drop for ChunkWriter<'_, W> { fn drop(&mut self) { let _ = self.flush(); } } // TODO: find a better name // /// This enum is used to be allow the `StreamWriter` to keep /// its inner `ChunkWriter` without wrapping it inside a /// `ZlibEncoder`. This is used in the case that between the /// change of state that happens when the last write of a frame /// is performed an error occurs, which obviously has to be returned. /// This creates the problem of where to store the writer before /// exiting the function, and this is where `Wrapper` comes in. /// /// Unfortunately the `ZlibWriter` can't be used because on the /// write following the error, `finish` would be called and that /// would write some data even if 0 bytes where compressed. /// /// If the `finish` function fails then there is nothing much to /// do as the `ChunkWriter` would get lost so the `Unrecoverable` /// variant is used to signal that. enum Wrapper<'a, W: Write> { Chunk(ChunkWriter<'a, W>), Zlib(ZlibEncoder>), Unrecoverable, /// This is used in-between, should never be matched None, } impl<'a, W: Write> Wrapper<'a, W> { /// Like `Option::take` this returns the `Wrapper` contained /// in `self` and replaces it with `Wrapper::None` fn take(&mut self) -> Wrapper<'a, W> { let mut swap = Wrapper::None; mem::swap(self, &mut swap); swap } } /// Streaming PNG writer /// /// This may silently fail in the destructor, so it is a good idea to call /// [`finish`](#method.finish) or [`flush`] before dropping. /// /// [`flush`]: https://doc.rust-lang.org/stable/std/io/trait.Write.html#tymethod.flush pub struct StreamWriter<'a, W: Write> { /// The option here is needed in order to access the inner `ChunkWriter` in-between /// each frame, which is needed for writing the fcTL chunks between each frame writer: Wrapper<'a, W>, prev_buf: Vec, curr_buf: Vec, /// Amount of data already written index: usize, /// length of the current scanline line_len: usize, /// size of the frame (width * height * sample_size) to_write: usize, width: u32, height: u32, bpp: BytesPerPixel, filter: FilterType, adaptive_filter: AdaptiveFilterType, fctl: Option, compression: Compression, } impl<'a, W: Write> StreamWriter<'a, W> { fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> Result> { let PartialInfo { width, height, frame_control: fctl, compression, .. } = writer.info; let bpp = writer.info.bpp_in_prediction(); let in_len = writer.info.raw_row_length() - 1; 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]; let mut chunk_writer = ChunkWriter::new(writer, buf_len); let (line_len, to_write) = chunk_writer.next_frame_info(); chunk_writer.write_header()?; let zlib = ZlibEncoder::new(chunk_writer, compression.to_options()); Ok(StreamWriter { writer: Wrapper::Zlib(zlib), index: 0, prev_buf, curr_buf, bpp, filter, width, height, adaptive_filter, line_len, to_write, fctl, compression, }) } /// Set the used filter type for the next frame. /// /// The default filter is [`FilterType::Sub`] which provides a basic prediction algorithm for /// sample values based on the previous. For a potentially better compression ratio, at the /// cost of more complex processing, try out [`FilterType::Paeth`]. /// /// [`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; } /// Set the adaptive filter type for the next frame. /// /// Adaptive filtering attempts to select the best filter for each line /// based on heuristics which minimize the file size for compression rather /// than use a single filter for the entire image. The default method is /// [`AdaptiveFilterType::NonAdaptive`]. /// /// [`AdaptiveFilterType::NonAdaptive`]: enum.AdaptiveFilterType.html pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) { self.adaptive_filter = adaptive_filter; } /// Set the fraction of time the following frames are going to be displayed, /// in seconds /// /// If the denominator is 0, it is to be treated as if it were 100 /// (that is, the numerator then specifies 1/100ths of a second). /// If the the value of the numerator is 0 the decoder should render the next frame /// as quickly as possible, though viewers may impose a reasonable lower bound. /// /// This method will return an error if the image is not animated. pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> { if let Some(ref mut fctl) = self.fctl { fctl.delay_den = denominator; fctl.delay_num = numerator; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the dimension of the following frames. /// /// This function will return an error when: /// - The image is not an animated; /// /// - The selected dimension, considering also the current frame position, /// goes outside the image boundaries; /// /// - One or both the width and height are 0; /// pub fn set_frame_dimension(&mut self, width: u32, height: u32) -> Result<()> { if let Some(ref mut fctl) = self.fctl { if Some(width) > self.width.checked_sub(fctl.x_offset) || Some(height) > self.height.checked_sub(fctl.y_offset) { return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into())); } else if width == 0 { return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into())); } else if height == 0 { return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into())); } fctl.width = width; fctl.height = height; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the position of the following frames. /// /// An error will be returned if: /// - The image is not animated; /// /// - The selected position, considering also the current frame dimension, /// goes outside the image boundaries; /// pub fn set_frame_position(&mut self, x: u32, y: u32) -> Result<()> { if let Some(ref mut fctl) = self.fctl { if Some(x) > self.width.checked_sub(fctl.width) || Some(y) > self.height.checked_sub(fctl.height) { return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into())); } fctl.x_offset = x; fctl.y_offset = y; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the frame dimension to occupy all the image, starting from /// the current position. /// /// To reset the frame to the full image size [`reset_frame_position`] /// should be called first. /// /// This method will return an error if the image is not animated. /// /// [`reset_frame_position`]: struct.Writer.html#method.reset_frame_position pub fn reset_frame_dimension(&mut self) -> Result<()> { if let Some(ref mut fctl) = self.fctl { fctl.width = self.width - fctl.x_offset; fctl.height = self.height - fctl.y_offset; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the frame position to (0, 0). /// /// Equivalent to calling [`set_frame_position(0, 0)`]. /// /// This method will return an error if the image is not animated. /// /// [`set_frame_position(0, 0)`]: struct.Writer.html#method.set_frame_position pub fn reset_frame_position(&mut self) -> Result<()> { if let Some(ref mut fctl) = self.fctl { fctl.x_offset = 0; fctl.y_offset = 0; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the blend operation for the following frames. /// /// The blend operation specifies whether the frame is to be alpha blended /// into the current output buffer content, or whether it should completely /// replace its region in the output buffer. /// /// See the [`BlendOp`] documentation for the possible values and their effects. /// /// *Note that for the first frame the two blend modes are functionally /// equivalent due to the clearing of the output buffer at the beginning /// of each play.* /// /// This method will return an error if the image is not animated. /// /// [`BlendOP`]: enum.BlendOp.html pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> { if let Some(ref mut fctl) = self.fctl { fctl.blend_op = op; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } /// Set the dispose operation for the following frames. /// /// The dispose operation specifies how the output buffer should be changed /// at the end of the delay (before rendering the next frame) /// /// See the [`DisposeOp`] documentation for the possible values and their effects. /// /// *Note that if the first frame uses [`DisposeOp::Previous`] /// it will be treated as [`DisposeOp::Background`].* /// /// This method will return an error if the image is not animated. /// /// [`DisposeOp`]: ../common/enum.BlendOp.html /// [`DisposeOp::Previous`]: ../common/enum.BlendOp.html#variant.Previous /// [`DisposeOp::Background`]: ../common/enum.BlendOp.html#variant.Background pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> { if let Some(ref mut fctl) = self.fctl { fctl.dispose_op = op; Ok(()) } else { Err(EncodingError::Format(FormatErrorKind::NotAnimated.into())) } } pub fn finish(mut self) -> Result<()> { 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()?; if let Wrapper::Chunk(wrt) = self.writer.take() { wrt.writer.validate_sequence_done()?; } Ok(()) } /// 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, Wrapper::Unrecoverable => { let err = FormatErrorKind::Unrecoverable.into(); return Err(EncodingError::Format(err)); } Wrapper::Zlib(_) => unreachable!("never called on a half-finished frame"), Wrapper::None => unreachable!(), }; wrt.flush()?; wrt.writer.validate_new_image()?; if let Some(fctl) = self.fctl { wrt.set_fctl(fctl); } let (scansize, size) = wrt.next_frame_info(); self.line_len = scansize; self.to_write = size; wrt.write_header()?; wrt.writer.increment_images_written(); // now it can be taken because the next statements cannot cause any errors match self.writer.take() { Wrapper::Chunk(wrt) => { let encoder = ZlibEncoder::new(wrt, self.compression.to_options()); self.writer = Wrapper::Zlib(encoder); } _ => unreachable!(), }; Ok(()) } } impl<'a, W: Write> Write for StreamWriter<'a, W> { fn write(&mut self, mut data: &[u8]) -> io::Result { if let Wrapper::Unrecoverable = self.writer { let err = FormatErrorKind::Unrecoverable.into(); return Err(EncodingError::Format(err).into()); } if data.is_empty() { return Ok(0); } if self.to_write == 0 { match self.writer.take() { Wrapper::Zlib(wrt) => match wrt.finish() { Ok(chunk) => self.writer = Wrapper::Chunk(chunk), Err(err) => { self.writer = Wrapper::Unrecoverable; return Err(err); } }, chunk @ Wrapper::Chunk(_) => self.writer = chunk, Wrapper::Unrecoverable => unreachable!(), Wrapper::None => unreachable!(), }; // Transition Wrapper::Chunk to Wrapper::Zlib. self.new_frame()?; } let written = data.read(&mut self.curr_buf[..self.line_len][self.index..])?; self.index += written; self.to_write -= written; if self.index == self.line_len { // TODO: reuse this buffer between rows. let mut filtered = vec![0; self.curr_buf.len()]; let filter_type = filter( self.filter, self.adaptive_filter, self.bpp, &self.prev_buf, &self.curr_buf, &mut filtered, ); // This can't fail as the other variant is used only to allow the zlib encoder to finish let wrt = match &mut self.writer { Wrapper::Zlib(wrt) => wrt, _ => unreachable!(), }; wrt.write_all(&[filter_type as u8])?; wrt.write_all(&filtered)?; 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()?, 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(()) } } impl Drop for StreamWriter<'_, W> { fn drop(&mut self) { let _ = self.flush(); } } #[cfg(test)] mod tests { use super::*; use crate::Decoder; use rand::{thread_rng, Rng}; use std::fs::File; use std::io::{Cursor, Write}; use std::{cmp, io}; #[test] fn roundtrip() { // More loops = more random testing, but also more test wait time for _ in 0..10 { for path in glob::glob("tests/pngsuite/*.png") .unwrap() .map(|r| r.unwrap()) { if path.file_name().unwrap().to_str().unwrap().starts_with('x') { // x* files are expected to fail to decode continue; } eprintln!("{}", path.display()); // Decode image let decoder = Decoder::new(File::open(path).unwrap()); let mut reader = decoder.read_info().unwrap(); let mut buf = vec![0; reader.output_buffer_size()]; let info = reader.next_frame(&mut buf).unwrap(); // Encode decoded image let mut out = Vec::new(); { let mut wrapper = RandomChunkWriter { rng: thread_rng(), w: &mut out, }; let mut encoder = Encoder::new(&mut wrapper, info.width, info.height); encoder.set_color(info.color_type); encoder.set_depth(info.bit_depth); if let Some(palette) = &reader.info().palette { encoder.set_palette(palette.clone()); } let mut encoder = encoder.write_header().unwrap(); encoder.write_image_data(&buf).unwrap(); } // Decode encoded decoded image let decoder = Decoder::new(&*out); let mut reader = decoder.read_info().unwrap(); let mut buf2 = vec![0; reader.output_buffer_size()]; reader.next_frame(&mut buf2).unwrap(); // check if the encoded image is ok: assert_eq!(buf, buf2); } } } #[test] fn roundtrip_stream() { // More loops = more random testing, but also more test wait time for _ in 0..10 { for path in glob::glob("tests/pngsuite/*.png") .unwrap() .map(|r| r.unwrap()) { if path.file_name().unwrap().to_str().unwrap().starts_with('x') { // x* files are expected to fail to decode continue; } // Decode image let decoder = Decoder::new(File::open(path).unwrap()); let mut reader = decoder.read_info().unwrap(); let mut buf = vec![0; reader.output_buffer_size()]; let info = reader.next_frame(&mut buf).unwrap(); // Encode decoded image let mut out = Vec::new(); { let mut wrapper = RandomChunkWriter { rng: thread_rng(), w: &mut out, }; let mut encoder = Encoder::new(&mut wrapper, info.width, info.height); encoder.set_color(info.color_type); encoder.set_depth(info.bit_depth); if let Some(palette) = &reader.info().palette { encoder.set_palette(palette.clone()); } let mut encoder = encoder.write_header().unwrap(); let mut stream_writer = encoder.stream_writer().unwrap(); let mut outer_wrapper = RandomChunkWriter { rng: thread_rng(), w: &mut stream_writer, }; outer_wrapper.write_all(&buf).unwrap(); } // Decode encoded decoded image let decoder = Decoder::new(&*out); let mut reader = decoder.read_info().unwrap(); let mut buf2 = vec![0; reader.output_buffer_size()]; reader.next_frame(&mut buf2).unwrap(); // check if the encoded image is ok: assert_eq!(buf, buf2); } } } #[test] fn image_palette() -> Result<()> { for &bit_depth in &[1u8, 2, 4, 8] { // Do a reference decoding, choose a fitting palette image from pngsuite let path = format!("tests/pngsuite/basn3p0{}.png", bit_depth); let decoder = Decoder::new(File::open(&path).unwrap()); let mut reader = decoder.read_info().unwrap(); let mut decoded_pixels = vec![0; reader.output_buffer_size()]; let info = reader.info(); assert_eq!( info.width as usize * info.height as usize * usize::from(bit_depth), decoded_pixels.len() * 8 ); let info = reader.next_frame(&mut decoded_pixels).unwrap(); let indexed_data = decoded_pixels; let palette = reader.info().palette.as_ref().unwrap(); let mut out = Vec::new(); { let mut encoder = Encoder::new(&mut out, info.width, info.height); encoder.set_depth(BitDepth::from_u8(bit_depth).unwrap()); encoder.set_color(ColorType::Indexed); encoder.set_palette(palette.as_ref()); let mut writer = encoder.write_header().unwrap(); writer.write_image_data(&indexed_data).unwrap(); } // Decode re-encoded image let decoder = Decoder::new(&*out); let mut reader = decoder.read_info().unwrap(); let mut redecoded = vec![0; reader.output_buffer_size()]; reader.next_frame(&mut redecoded).unwrap(); // check if the encoded image is ok: assert_eq!(indexed_data, redecoded); } Ok(()) } #[test] fn expect_error_on_wrong_image_len() -> Result<()> { let width = 10; let height = 10; let output = vec![0u8; 1024]; let writer = Cursor::new(output); let mut encoder = Encoder::new(writer, width as u32, height as u32); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::Rgb); let mut png_writer = encoder.write_header()?; let correct_image_size = width * height * 3; let image = vec![0u8; correct_image_size + 1]; let result = png_writer.write_image_data(image.as_ref()); assert!(result.is_err()); Ok(()) } #[test] fn expect_error_on_empty_image() -> Result<()> { let output = vec![0u8; 1024]; let mut writer = Cursor::new(output); let encoder = Encoder::new(&mut writer, 0, 0); assert!(encoder.write_header().is_err()); let encoder = Encoder::new(&mut writer, 100, 0); assert!(encoder.write_header().is_err()); let encoder = Encoder::new(&mut writer, 0, 100); assert!(encoder.write_header().is_err()); Ok(()) } #[test] fn expect_error_on_invalid_bit_depth_color_type_combination() -> Result<()> { let output = vec![0u8; 1024]; let mut writer = Cursor::new(output); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::One); encoder.set_color(ColorType::Rgb); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::One); encoder.set_color(ColorType::GrayscaleAlpha); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::One); encoder.set_color(ColorType::Rgba); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Two); encoder.set_color(ColorType::Rgb); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Two); encoder.set_color(ColorType::GrayscaleAlpha); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Two); encoder.set_color(ColorType::Rgba); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Four); encoder.set_color(ColorType::Rgb); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Four); encoder.set_color(ColorType::GrayscaleAlpha); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Four); encoder.set_color(ColorType::Rgba); assert!(encoder.write_header().is_err()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Sixteen); encoder.set_color(ColorType::Indexed); assert!(encoder.write_header().is_err()); Ok(()) } #[test] fn can_write_header_with_valid_bit_depth_color_type_combination() -> Result<()> { let output = vec![0u8; 1024]; let mut writer = Cursor::new(output); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::One); encoder.set_color(ColorType::Grayscale); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::One); encoder.set_color(ColorType::Indexed); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Two); encoder.set_color(ColorType::Grayscale); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Two); encoder.set_color(ColorType::Indexed); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Four); encoder.set_color(ColorType::Grayscale); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Four); encoder.set_color(ColorType::Indexed); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::Grayscale); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::Rgb); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::Indexed); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::GrayscaleAlpha); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::Rgba); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Sixteen); encoder.set_color(ColorType::Grayscale); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Sixteen); encoder.set_color(ColorType::Rgb); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Sixteen); encoder.set_color(ColorType::GrayscaleAlpha); assert!(encoder.write_header().is_ok()); let mut encoder = Encoder::new(&mut writer, 1, 1); encoder.set_depth(BitDepth::Sixteen); encoder.set_color(ColorType::Rgba); assert!(encoder.write_header().is_ok()); Ok(()) } #[test] fn all_filters_roundtrip() -> io::Result<()> { let pixel: Vec<_> = (0..48).collect(); let roundtrip = |filter: FilterType| -> io::Result<()> { let mut buffer = vec![]; let mut encoder = Encoder::new(&mut buffer, 4, 4); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::Rgb); encoder.set_filter(filter); encoder.write_header()?.write_image_data(&pixel)?; let decoder = crate::Decoder::new(Cursor::new(buffer)); let mut reader = decoder.read_info()?; let info = reader.info(); assert_eq!(info.width, 4); assert_eq!(info.height, 4); let mut dest = vec![0; pixel.len()]; reader.next_frame(&mut dest)?; assert_eq!(dest, pixel, "Deviation with filter type {:?}", filter); Ok(()) }; roundtrip(FilterType::NoFilter)?; roundtrip(FilterType::Sub)?; roundtrip(FilterType::Up)?; roundtrip(FilterType::Avg)?; roundtrip(FilterType::Paeth)?; Ok(()) } #[test] fn some_gamma_roundtrip() -> io::Result<()> { let pixel: Vec<_> = (0..48).collect(); let roundtrip = |gamma: Option| -> io::Result<()> { let mut buffer = vec![]; let mut encoder = Encoder::new(&mut buffer, 4, 4); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::Rgb); encoder.set_filter(FilterType::Avg); if let Some(gamma) = gamma { encoder.set_source_gamma(gamma); } encoder.write_header()?.write_image_data(&pixel)?; let decoder = crate::Decoder::new(Cursor::new(buffer)); let mut reader = decoder.read_info()?; assert_eq!( reader.info().source_gamma, gamma, "Deviation with gamma {:?}", gamma ); let mut dest = vec![0; pixel.len()]; let info = reader.next_frame(&mut dest)?; assert_eq!(info.width, 4); assert_eq!(info.height, 4); Ok(()) }; roundtrip(None)?; roundtrip(Some(ScaledFloat::new(0.35)))?; roundtrip(Some(ScaledFloat::new(0.45)))?; roundtrip(Some(ScaledFloat::new(0.55)))?; roundtrip(Some(ScaledFloat::new(0.7)))?; roundtrip(Some(ScaledFloat::new(1.0)))?; roundtrip(Some(ScaledFloat::new(2.5)))?; Ok(()) } #[test] fn write_image_chunks_beyond_first() -> Result<()> { 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(()) } #[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] 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] fn stream_filtering() -> Result<()> { let output = vec![0u8; 1024]; let mut cursor = Cursor::new(output); let mut encoder = Encoder::new(&mut cursor, 8, 8); encoder.set_color(ColorType::Rgba); encoder.set_filter(FilterType::Paeth); let mut writer = encoder.write_header()?; let mut stream = writer.stream_writer()?; for _ in 0..8 { let written = stream.write(&[1; 32])?; assert_eq!(written, 32); } stream.finish()?; drop(writer); { cursor.set_position(0); let mut decoder = Decoder::new(cursor).read_info().expect("A valid image"); let mut buffer = [0u8; 256]; decoder.next_frame(&mut buffer[..]).expect("Valid read"); assert_eq!(buffer, [1; 256]); } 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(()) } #[test] fn finish_drops_inner_writer() -> Result<()> { struct NoWriter<'flag>(&'flag mut bool); impl Write for NoWriter<'_> { fn write(&mut self, buf: &[u8]) -> io::Result { Ok(buf.len()) } fn flush(&mut self) -> io::Result<()> { Ok(()) } } impl Drop for NoWriter<'_> { fn drop(&mut self) { *self.0 = true; } } let mut flag = false; { let mut encoder = Encoder::new(NoWriter(&mut flag), 10, 10); encoder.set_depth(BitDepth::Eight); encoder.set_color(ColorType::Grayscale); let mut writer = encoder.write_header()?; writer.write_image_data(&[0; 100])?; writer.finish()?; } assert!(flag, "PNG finished but writer was not dropped"); Ok(()) } /// A Writer that only writes a few bytes at a time struct RandomChunkWriter { rng: R, w: W, } impl Write for RandomChunkWriter { fn write(&mut self, buf: &[u8]) -> io::Result { // choose a random length to write let len = cmp::min(self.rng.gen_range(1..50), buf.len()); self.w.write(&buf[0..len]) } fn flush(&mut self) -> io::Result<()> { self.w.flush() } } } /// Mod to encapsulate the converters depending on the `deflate` crate. /// /// 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 Compression { fn to_options(self) -> flate2::Compression { #[allow(deprecated)] match self { Compression::Default => flate2::Compression::default(), Compression::Fast => flate2::Compression::fast(), Compression::Best => flate2::Compression::best(), #[allow(deprecated)] Compression::Huffman => flate2::Compression::none(), #[allow(deprecated)] Compression::Rle => flate2::Compression::none(), } } }