Add max_uncompressed_size (#732)

Adds an option to skip files when the uncompressed IDAT (or the file
size itself) exceeds a certain size. This provides an effective way to
limit resource usage in constrained environments.

Closes #411.

Cargo.lock

@@ -418,12 +418,19 @@ dependencies = [
  "indexmap",
  "libdeflater",
  "log",
+ "parse-size",
  "rayon",
  "rgb",
  "rustc-hash",
  "zopfli",
 ]
 
+[[package]]
+name = "parse-size"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "487f2ccd1e17ce8c1bfab3a65c89525af41cfad4c8659021a1e9a2aacd73b89b"
+
 [[package]]
 name = "pkg-config"
 version = "0.3.32"

Cargo.toml

@@ -69,6 +69,10 @@ optional = true
 version = "4.5.48"
 features = ["wrap_help"]
 
+[dependencies.parse-size]
+optional = true
+version = "1.1.0"
+
 [target.'cfg(windows)'.dependencies.glob]
 optional = true
 version = "0.3.3"
@@ -80,7 +84,7 @@ features = ["png"]
 version = "0.25.8"
 
 [features]
-binary = ["dep:clap", "dep:glob", "dep:env_logger"]
+binary = ["dep:clap", "dep:glob", "dep:env_logger", "dep:parse-size"]
 default = ["binary", "parallel", "zopfli", "filetime"]
 parallel = ["dep:rayon", "indexmap/rayon", "dep:crossbeam-channel"]
 freestanding = ["libdeflater/freestanding"]

src/cli.rs

@@ -1,6 +1,7 @@
 use std::{num::NonZeroU64, path::PathBuf};
 
 use clap::{Arg, ArgAction, Command, builder::ArgPredicate, value_parser};
+use parse_size::parse_size;
 
 include!("display_chunks.rs");
 
@@ -365,6 +366,21 @@ high compression levels.")
                 .value_name("secs")
                 .value_parser(value_parser!(u64)),
         )
+        .arg(
+            Arg::new("max-size")
+                .help("Skip image if the decompressed size exceeds this limit")
+                .long_help("\
+Maximum size to allow for the input image. If the raw, decompressed image data (or the \
+file size) of the image exceeds this size, it will be skipped. This is useful for limiting \
+memory usage or avoiding long processing times on large images. The value may be specified \
+with a unit suffix such as k, KB, m, MB, etc.
+
+The decompressed size of an image is roughly equal to width * height * bit-depth / 8. E.g. \
+a 1920x1080 image with 24-bit color depth would be roughly 6MB.")
+                .long("max-raw-size")
+                .value_name("bytes")
+                .value_parser(|s: &str| parse_size(s)),
+        )
         .arg(
             Arg::new("threads")
                 .help("Number of threads to use [default: num logical CPUs]")

src/error.rs

@@ -43,7 +43,10 @@ impl fmt::Display for PngError {
                 f,
                 "Data length {l1} does not match the expected length {l2}"
             ),
-            PngError::InflatedDataTooLong(_) => f.write_str("Inflated data too long"),
+            PngError::InflatedDataTooLong(max) => write!(
+                f,
+                "Inflated data would exceed the maximum size ({max} bytes)"
+            ),
             PngError::InvalidData => f.write_str("Invalid data found; unable to read PNG file"),
             PngError::InvalidDepthForType(d, ref c) => {
                 write!(f, "Invalid bit depth {d} for color type {c}")

src/headers.rs

@@ -248,7 +248,7 @@ pub fn read_be_u32(bytes: &[u8]) -> u32 {
 }
 
 /// Extract and decompress the ICC profile from an iCCP chunk
-pub fn extract_icc(iccp: &Chunk) -> Option<Vec<u8>> {
+pub fn extract_icc(iccp: &Chunk, max_size: Option<usize>) -> Option<Vec<u8>> {
     // Skip (useless) profile name
     let mut data = iccp.data.as_slice();
     loop {
@@ -264,8 +264,11 @@ pub fn extract_icc(iccp: &Chunk) -> Option<Vec<u8>> {
         return None; // The profile is supposed to be compressed (method 0)
     }
     // The decompressed size is unknown so we have to guess the required buffer size
-    let max_size = compressed_data.len() * 2 + 1000;
+    let mut out_size = compressed_data.len() * 2 + 1000;
-    match inflate(compressed_data, max_size) {
+    if let Some(max) = max_size {
+        out_size = out_size.min(max);
+    }
+    match inflate(compressed_data, out_size) {
         Ok(icc) => Some(icc),
         Err(e) => {
             // Log the error so we can know if the buffer size needs to be adjusted
@@ -331,7 +334,7 @@ pub fn preprocess_chunks(aux_chunks: &mut Vec<Chunk>, opts: &mut Options) {
             trace!("Removing iCCP chunk due to conflict with sRGB chunk");
             aux_chunks.remove(iccp_idx);
             allow_grayscale = true;
-        } else if let Some(icc) = extract_icc(&aux_chunks[iccp_idx]) {
+        } else if let Some(icc) = extract_icc(&aux_chunks[iccp_idx], opts.max_decompressed_size) {
             let intent = if may_replace_iccp {
                 srgb_rendering_intent(&icc)
             } else {

src/main.rs

@@ -259,6 +259,8 @@ fn parse_opts_into_struct(
 
     opts.fix_errors = matches.get_flag("fix");
 
+    opts.max_decompressed_size = matches.get_one::<u64>("max-size").map(|&x| x as usize);
+
     opts.bit_depth_reduction = !matches.get_flag("no-bit-reduction");
 
     opts.color_type_reduction = !matches.get_flag("no-color-reduction");
@@ -425,6 +427,13 @@ fn parse_numeric_range_opts(
 }
 
 fn process_file(input: &InFile, output: &OutFile, opts: &Options) -> OptimizationResult {
+    if let (Some(max_size), InFile::Path(path)) = (opts.max_decompressed_size, input) {
+        if path.metadata().is_ok_and(|m| m.len() > max_size as u64) {
+            warn!("{input}: Skipped: File exceeds the maximum size ({max_size} bytes)");
+            return OptimizationResult::Skipped;
+        }
+    }
+
     match oxipng::optimize(input, output, opts) {
         // For optimizing single files, this will return the correct exit code always.
         // For recursive optimization, the correct choice is a bit subjective.
@@ -434,8 +443,8 @@ fn process_file(input: &InFile, output: &OutFile, opts: &Options) -> Optimizatio
         // PNG files, and return an error for them.
         // We don't really want to return an error code for those files.
         Ok(_) => OptimizationResult::Ok,
-        Err(e @ PngError::C2PAMetadataPreventsChanges) => {
+        Err(e @ PngError::C2PAMetadataPreventsChanges | e @ PngError::InflatedDataTooLong(_)) => {
-            warn!("{input}: {e}");
+            warn!("{input}: Skipped: {e}");
             OptimizationResult::Skipped
         }
         Err(e) => {

src/options.rs

@@ -154,6 +154,11 @@ pub struct Options {
     ///
     /// Default: `None`
     pub timeout: Option<Duration>,
+    /// Maximum decompressed size of the input IDAT.
+    /// If decompression would exceed this size, it will be rejected.
+    ///
+    /// Default: `None`
+    pub max_decompressed_size: Option<usize>,
 }
 
 impl Options {
@@ -263,6 +268,7 @@ impl Default for Options {
             deflater: Deflater::Libdeflater { compression: 11 },
             fast_evaluation: true,
             timeout: None,
+            max_decompressed_size: None,
         }
     }
 }

src/png/mod.rs

@@ -147,6 +147,11 @@ impl PngData {
             key_chunks.remove(b"PLTE"),
             key_chunks.remove(b"tRNS"),
         )?;
+        if let Some(max) = opts.max_decompressed_size {
+            if ihdr.raw_data_size() > max {
+                return Err(PngError::InflatedDataTooLong(max));
+            }
+        }
 
         let raw = PngImage::new(ihdr, &idat_data)?;
 

tests/flags.rs

@@ -634,6 +634,19 @@ fn scale_16() {
     );
 }
 
+#[test]
+fn max_size() {
+    let result = oxipng::optimize(
+        &"tests/files/max_size.png".into(),
+        &OutFile::None,
+        &Options {
+            max_decompressed_size: Some(150_000),
+            ..Options::default()
+        },
+    );
+    assert!(matches!(result, Err(PngError::InflatedDataTooLong(_))));
+}
+
 #[test]
 #[cfg(feature = "zopfli")]
 fn zopfli_mode() {

xtask/Cargo.lock

@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "anstream"
@@ -100,6 +100,12 @@ version = "1.70.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf"
 
+[[package]]
+name = "parse-size"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "487f2ccd1e17ce8c1bfab3a65c89525af41cfad4c8659021a1e9a2aacd73b89b"
+
 [[package]]
 name = "roff"
 version = "0.2.2"
@@ -197,4 +203,5 @@ version = "0.1.0"
 dependencies = [
  "clap",
  "clap_mangen",
+ "parse-size",
 ]

xtask/Cargo.toml

@@ -8,3 +8,4 @@ publish = false
 [dependencies]
 clap = "4.5.21"
 clap_mangen = "0.2.24"
+parse-size = "1.1.0"