From 4989d5f83a4a1aaaf7b1fd1f33f7b4db1d3404d3 Mon Sep 17 00:00:00 2001 From: "Sergey \"Shnatsel\" Davidoff" Date: Sun, 21 Jan 2024 01:00:04 +0000 Subject: [PATCH] Support limits in PNG animation decoding (#2112) --- src/codecs/gif.rs | 18 +++++----- src/codecs/png.rs | 86 +++++++++++++++++++++++++++++++++++++---------- src/io/mod.rs | 18 +++++++++- 3 files changed, 94 insertions(+), 28 deletions(-) diff --git a/src/codecs/gif.rs b/src/codecs/gif.rs index 8b6f5ee26f..afa8b5b53b 100644 --- a/src/codecs/gif.rs +++ b/src/codecs/gif.rs @@ -263,19 +263,17 @@ impl Iterator for GifFrameIterator { type Item = ImageResult; fn next(&mut self) -> Option> { - fn limits_reserve_buffer(limits: &mut Limits, width: u32, height: u32) -> ImageResult<()> { - limits.check_dimensions(width, height)?; - // cannot overflow because width/height are u16 in the actual GIF format, - // so this cannot exceed 64GB which easily fits into a u64 - let in_memory_size = width as u64 * height as u64 * 4; - limits.reserve(in_memory_size) - } + // The iterator always produces RGBA8 images + const COLOR_TYPE: ColorType = ColorType::Rgba8; // Allocate the buffer for the previous frame. // This is done here and not in the constructor because // the constructor cannot return an error when the allocation limit is exceeded. if self.non_disposed_frame.is_none() { - if let Err(e) = limits_reserve_buffer(&mut self.limits, self.width, self.height) { + if let Err(e) = self + .limits + .reserve_buffer(self.width, self.height, COLOR_TYPE) + { return Some(Err(e)); } self.non_disposed_frame = Some(ImageBuffer::from_pixel( @@ -308,7 +306,7 @@ impl Iterator for GifFrameIterator { let mut local_limits = self.limits.clone(); // Check the allocation we're about to perform against the limits - if let Err(e) = limits_reserve_buffer(&mut local_limits, frame.width, frame.height) { + if let Err(e) = local_limits.reserve_buffer(frame.width, frame.height, COLOR_TYPE) { return Some(Err(e)); } // Allocate the buffer now that the limits allowed it @@ -382,7 +380,7 @@ impl Iterator for GifFrameIterator { frame_buffer } else { // Check limits before allocating the buffer - if let Err(e) = limits_reserve_buffer(&mut local_limits, self.width, self.height) { + if let Err(e) = local_limits.reserve_buffer(self.width, self.height, COLOR_TYPE) { return Some(Err(e)); } ImageBuffer::from_fn(self.width, self.height, |x, y| { diff --git a/src/codecs/png.rs b/src/codecs/png.rs index 87b28ca773..d92ee44970 100644 --- a/src/codecs/png.rs +++ b/src/codecs/png.rs @@ -112,12 +112,13 @@ impl Read for PngReader { pub struct PngDecoder { color_type: ColorType, reader: png::Reader, + limits: Limits, } impl PngDecoder { /// Creates a new decoder that decodes from the stream ```r``` pub fn new(r: R) -> ImageResult> { - Self::with_limits(r, Limits::default()) + Self::with_limits(r, Limits::no_limits()) } /// Creates a new decoder that decodes from the stream ```r``` with the given limits. @@ -191,7 +192,11 @@ impl PngDecoder { } }; - Ok(PngDecoder { color_type, reader }) + Ok(PngDecoder { + color_type, + reader, + limits, + }) } /// Returns the gamma value of the image or None if no gamma value is indicated. @@ -287,6 +292,16 @@ impl<'a, R: 'a + Read> ImageDecoder<'a> for PngDecoder { let width = self.reader.info().width; self.reader.output_line_size(width) as u64 } + + fn set_limits(&mut self, limits: Limits) -> ImageResult<()> { + limits.check_support(&crate::io::LimitSupport::default())?; + let info = self.reader.info(); + limits.check_dimensions(info.width, info.height)?; + self.limits = limits; + // TODO: add `png::Reader::change_limits()` and call it here + // to also constrain the internal buffer allocations in the PNG crate + Ok(()) + } } /// An [`AnimationDecoder`] adapter of [`PngDecoder`]. @@ -299,9 +314,9 @@ impl<'a, R: 'a + Read> ImageDecoder<'a> for PngDecoder { pub struct ApngDecoder { inner: PngDecoder, /// The current output buffer. - current: RgbaImage, + current: Option, /// The previous output buffer, used for dispose op previous. - previous: RgbaImage, + previous: Option, /// The dispose op of the current frame. dispose: DisposeOp, /// The number of image still expected to be able to load. @@ -312,7 +327,6 @@ pub struct ApngDecoder { impl ApngDecoder { fn new(inner: PngDecoder) -> Self { - let (width, height) = inner.dimensions(); let info = inner.reader.info(); let remaining = match info.animation_control() { // The expected number of fcTL in the remaining image. @@ -324,9 +338,8 @@ impl ApngDecoder { let has_thumbnail = info.frame_control.is_none(); ApngDecoder { inner, - // TODO: should we delay this allocation? At least if we support limits we should. - current: RgbaImage::new(width, height), - previous: RgbaImage::new(width, height), + current: None, + previous: None, dispose: DisposeOp::Background, remaining, has_thumbnail, @@ -337,6 +350,24 @@ impl ApngDecoder { /// Decode one subframe and overlay it on the canvas. fn mix_next_frame(&mut self) -> Result, ImageError> { + // The iterator always produces RGBA8 images + const COLOR_TYPE: ColorType = ColorType::Rgba8; + + // Allocate the buffers, honoring the memory limits + let (width, height) = self.inner.dimensions(); + { + let limits = &mut self.inner.limits; + if self.previous.is_none() { + limits.reserve_buffer(width, height, COLOR_TYPE)?; + self.previous = Some(RgbaImage::new(width, height)); + } + + if self.current.is_none() { + limits.reserve_buffer(width, height, COLOR_TYPE)?; + self.current = Some(RgbaImage::new(width, height)); + } + } + // Remove this image from remaining. self.remaining = match self.remaining.checked_sub(1) { None => return Ok(None), @@ -349,34 +380,52 @@ impl ApngDecoder { // Skip the thumbnail that is not part of the animation. if self.has_thumbnail { - self.has_thumbnail = false; + // Clone the limits so that our one-off allocation that's destroyed after this scope doesn't persist + let mut limits = self.inner.limits.clone(); + limits.reserve_usize(self.inner.reader.output_buffer_size())?; let mut buffer = vec![0; self.inner.reader.output_buffer_size()]; + // TODO: add `png::Reader::change_limits()` and call it here + // to also constrain the internal buffer allocations in the PNG crate self.inner .reader .next_frame(&mut buffer) .map_err(ImageError::from_png)?; + self.has_thumbnail = false; } self.animatable_color_type()?; + // We've initialized them earlier in this function + let previous = self.previous.as_mut().unwrap(); + let current = self.current.as_mut().unwrap(); + // Dispose of the previous frame. match self.dispose { DisposeOp::None => { - self.previous.clone_from(&self.current); + previous.clone_from(current); } DisposeOp::Background => { - self.previous.clone_from(&self.current); - self.current + previous.clone_from(current); + current .pixels_mut() .for_each(|pixel| *pixel = Rgba([0, 0, 0, 0])); } DisposeOp::Previous => { - self.current.clone_from(&self.previous); + current.clone_from(previous); } } + // The allocations from now on are not going to persist, + // and will be destroyed at the end of the scope. + // Clone the limits so that any changes to them die with the allocations. + let mut limits = self.inner.limits.clone(); + // Read next frame data. - let mut buffer = vec![0; self.inner.reader.output_buffer_size()]; + let raw_frame_size = self.inner.reader.output_buffer_size(); + limits.reserve_usize(raw_frame_size)?; + let mut buffer = vec![0; raw_frame_size]; + // TODO: add `png::Reader::change_limits()` and call it here + // to also constrain the internal buffer allocations in the PNG crate self.inner .reader .next_frame(&mut buffer) @@ -404,6 +453,7 @@ impl ApngDecoder { }; // Turn the data into an rgba image proper. + limits.reserve_buffer(width, height, COLOR_TYPE)?; let source = match self.inner.color_type { ColorType::L8 => { let image = ImageBuffer::, _>::from_raw(width, height, buffer).unwrap(); @@ -424,17 +474,19 @@ impl ApngDecoder { } _ => unreachable!("Invalid png color"), }; + // We've converted the raw frame to RGBA8 and disposed of the original allocation + limits.free_usize(raw_frame_size); match blend { BlendOp::Source => { - self.current + current .copy_from(&source, px, py) .expect("Invalid png image not detected in png"); } BlendOp::Over => { // TODO: investigate speed, speed-ups, and bounds-checks. for (x, y, p) in source.enumerate_pixels() { - self.current.get_pixel_mut(x + px, y + py).blend(p); + current.get_pixel_mut(x + px, y + py).blend(p); } } } @@ -442,7 +494,7 @@ impl ApngDecoder { // Ok, we can proceed with actually remaining images. self.remaining = remaining; // Return composited output buffer. - Ok(Some(&self.current)) + Ok(Some(self.current.as_ref().unwrap())) } fn animatable_color_type(&self) -> Result<(), ImageError> { diff --git a/src/io/mod.rs b/src/io/mod.rs index 8fbc6e24a6..2c48a948ee 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -2,7 +2,7 @@ use std::convert::TryFrom; -use crate::{error, ImageError, ImageResult}; +use crate::{error, ColorType, ImageError, ImageResult}; pub(crate) mod free_functions; mod reader; @@ -141,6 +141,22 @@ impl Limits { } } + /// This function acts identically to [`reserve`], but accepts the width, height and color type + /// used to create an [`ImageBuffer`] and does all the math for you. + pub fn reserve_buffer( + &mut self, + width: u32, + height: u32, + color_type: ColorType, + ) -> ImageResult<()> { + self.check_dimensions(width, height)?; + let in_memory_size = (width as u64) + .saturating_mul(height as u64) + .saturating_mul(color_type.bytes_per_pixel().into()); + self.reserve(in_memory_size)?; + Ok(()) + } + /// This function increases the `max_alloc` limit with amount. Should only be used /// together with [`reserve`]. ///