Support limits in PNG animation decoding (#2112)

src/codecs/gif.rs

@@ -263,19 +263,17 @@ impl<R: Read> Iterator for GifFrameIterator<R> {
     type Item = ImageResult<animation::Frame>;
 
     fn next(&mut self) -> Option<ImageResult<animation::Frame>> {
-        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<R: Read> Iterator for GifFrameIterator<R> {
         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<R: Read> Iterator for GifFrameIterator<R> {
             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| {

src/codecs/png.rs

@@ -112,12 +112,13 @@ impl<R: Read> Read for PngReader<R> {
 pub struct PngDecoder<R: Read> {
     color_type: ColorType,
     reader: png::Reader<R>,
+    limits: Limits,
 }
 
 impl<R: Read> PngDecoder<R> {
     /// Creates a new decoder that decodes from the stream ```r```
     pub fn new(r: R) -> ImageResult<PngDecoder<R>> {
-        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<R: Read> PngDecoder<R> {
             }
         };
 
-        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<R> {
         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<R> {
 pub struct ApngDecoder<R: Read> {
     inner: PngDecoder<R>,
     /// The current output buffer.
-    current: RgbaImage,
+    current: Option<RgbaImage>,
     /// The previous output buffer, used for dispose op previous.
-    previous: RgbaImage,
+    previous: Option<RgbaImage>,
     /// 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<R: Read> {
 
 impl<R: Read> ApngDecoder<R> {
     fn new(inner: PngDecoder<R>) -> 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<R: Read> ApngDecoder<R> {
         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<R: Read> ApngDecoder<R> {
 
     /// Decode one subframe and overlay it on the canvas.
     fn mix_next_frame(&mut self) -> Result<Option<&RgbaImage>, 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<R: Read> ApngDecoder<R> {
 
         // 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<R: Read> ApngDecoder<R> {
         };
 
         // 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::<Luma<_>, _>::from_raw(width, height, buffer).unwrap();
@@ -424,25 +474,27 @@ impl<R: Read> ApngDecoder<R> {
             }
             _ => 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);
                 }
             }
         }
 
         // 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> {

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`].
     ///