From 42ee595e15b993a839965f76db81ea407c4d1a8f Mon Sep 17 00:00:00 2001
From: Jonathan Behrens <fintelia@gmail.com>
Date: Mon, 24 Oct 2022 23:00:56 -0700
Subject: [PATCH] Add PngDecoder::with_limits

---
 Cargo.toml               |  2 +-
 src/codecs/png.rs        | 18 ++++++++++++++----
 src/io/free_functions.rs |  7 ++++---
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6e9f1d9075..01f71c9bc1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -36,7 +36,7 @@ num-rational = { version = "0.4", default-features = false }
 num-traits = "0.2.0"
 gif = { version = "0.11.1", optional = true }
 jpeg = { package = "jpeg-decoder", version = "0.2.1", default-features = false, optional = true }
-png = { version = "0.17.0", optional = true }
+png = { version = "0.17.6", optional = true }
 scoped_threadpool = { version = "0.1", optional = true }
 tiff = { version = "0.7.1", optional = true }
 ravif = { version = "0.8.0", optional = true }
diff --git a/src/codecs/png.rs b/src/codecs/png.rs
index 63df3304d2..ae3032bf78 100644
--- a/src/codecs/png.rs
+++ b/src/codecs/png.rs
@@ -20,6 +20,7 @@ use crate::error::{
     ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind,
 };
 use crate::image::{AnimationDecoder, ImageDecoder, ImageEncoder, ImageFormat};
+use crate::io::Limits;
 use crate::{DynamicImage, GenericImage, ImageBuffer, Luma, LumaA, Rgb, Rgba, RgbaImage};
 
 // http://www.w3.org/TR/PNG-Structure.html
@@ -117,10 +118,19 @@ pub struct PngDecoder<R: Read> {
 impl<R: Read> PngDecoder<R> {
     /// Creates a new decoder that decodes from the stream ```r```
     pub fn new(r: R) -> ImageResult<PngDecoder<R>> {
-        let limits = png::Limits {
-            bytes: usize::max_value(),
-        };
-        let mut decoder = png::Decoder::new_with_limits(r, limits);
+        Self::with_limits(r, Limits::default())
+    }
+
+    /// Creates a new decoder that decodes from the stream ```r``` with the given limits.
+    pub fn with_limits(r: R, limits: Limits) -> ImageResult<PngDecoder<R>> {
+        limits.check_support(&crate::io::LimitSupport::default())?;
+
+        let max_bytes = usize::try_from(limits.max_alloc.unwrap_or(u64::MAX)).unwrap_or(usize::MAX);
+        let mut decoder = png::Decoder::new_with_limits(r, png::Limits { bytes: max_bytes });
+
+        let info = decoder.read_header_info().map_err(ImageError::from_png)?;
+        limits.check_dimensions(info.width, info.height)?;
+
         // By default the PNG decoder will scale 16 bpc to 8 bpc, so custom
         // transformations must be set. EXPAND preserves the default behavior
         // expanding bpc < 8 to 8 bpc.
diff --git a/src/io/free_functions.rs b/src/io/free_functions.rs
index fcb492fe33..741558ab92 100644
--- a/src/io/free_functions.rs
+++ b/src/io/free_functions.rs
@@ -45,6 +45,7 @@ pub(crate) trait DecoderVisitor {
 pub(crate) fn load_decoder<R: BufRead + Seek, V: DecoderVisitor>(
     r: R,
     format: ImageFormat,
+    limits: super::Limits,
     visitor: V,
 ) -> ImageResult<V::Result> {
     #[allow(unreachable_patterns)]
@@ -53,7 +54,7 @@ pub(crate) fn load_decoder<R: BufRead + Seek, V: DecoderVisitor>(
         #[cfg(feature = "avif-decoder")]
         image::ImageFormat::Avif => visitor.visit_decoder(avif::AvifDecoder::new(r)?),
         #[cfg(feature = "png")]
-        image::ImageFormat::Png => visitor.visit_decoder(png::PngDecoder::new(r)?),
+        image::ImageFormat::Png => visitor.visit_decoder(png::PngDecoder::with_limits(r, limits)?),
         #[cfg(feature = "gif")]
         image::ImageFormat::Gif => visitor.visit_decoder(gif::GifDecoder::new(r)?),
         #[cfg(feature = "jpeg")]
@@ -107,7 +108,7 @@ pub(crate) fn load_inner<R: BufRead + Seek>(
         }
     }
 
-    load_decoder(r, format, LoadVisitor(limits))
+    load_decoder(r, format, limits.clone(), LoadVisitor(limits))
 }
 
 pub(crate) fn image_dimensions_impl(path: &Path) -> ImageResult<(u32, u32)> {
@@ -131,7 +132,7 @@ pub(crate) fn image_dimensions_with_format_impl<R: BufRead + Seek>(
         }
     }
 
-    load_decoder(buffered_read, format, DimVisitor)
+    load_decoder(buffered_read, format, super::Limits::default(), DimVisitor)
 }
 
 #[allow(unused_variables)]