diff --git a/Cargo.lock b/Cargo.lock index 8673a4fd..19110c15 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 02761bae..2024aeba 100644 --- a/Cargo.toml +++ b/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"] diff --git a/src/cli.rs b/src/cli.rs index 8977cf99..108c89bd 100644 --- a/src/cli.rs +++ b/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]") diff --git a/src/error.rs b/src/error.rs index 0d022812..59320cb2 100644 --- a/src/error.rs +++ b/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}") diff --git a/src/headers.rs b/src/headers.rs index 3e3f6a99..06e09974 100644 --- a/src/headers.rs +++ b/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> { +pub fn extract_icc(iccp: &Chunk, max_size: Option) -> Option> { // Skip (useless) profile name let mut data = iccp.data.as_slice(); loop { @@ -264,8 +264,11 @@ pub fn extract_icc(iccp: &Chunk) -> Option> { 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; - match inflate(compressed_data, max_size) { + let mut out_size = compressed_data.len() * 2 + 1000; + 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, 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 { diff --git a/src/main.rs b/src/main.rs index f2ee9e3b..2823d232 100644 --- a/src/main.rs +++ b/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::("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) => { - warn!("{input}: {e}"); + Err(e @ PngError::C2PAMetadataPreventsChanges | e @ PngError::InflatedDataTooLong(_)) => { + warn!("{input}: Skipped: {e}"); OptimizationResult::Skipped } Err(e) => { diff --git a/src/options.rs b/src/options.rs index a712c5dd..476673cf 100644 --- a/src/options.rs +++ b/src/options.rs @@ -154,6 +154,11 @@ pub struct Options { /// /// Default: `None` pub timeout: Option, + /// Maximum decompressed size of the input IDAT. + /// If decompression would exceed this size, it will be rejected. + /// + /// Default: `None` + pub max_decompressed_size: Option, } impl Options { @@ -263,6 +268,7 @@ impl Default for Options { deflater: Deflater::Libdeflater { compression: 11 }, fast_evaluation: true, timeout: None, + max_decompressed_size: None, } } } diff --git a/src/png/mod.rs b/src/png/mod.rs index e0f0acae..e5129b91 100644 --- a/src/png/mod.rs +++ b/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)?; diff --git a/tests/files/max_size.png b/tests/files/max_size.png new file mode 100644 index 00000000..7c7f12ce Binary files /dev/null and b/tests/files/max_size.png differ diff --git a/tests/flags.rs b/tests/flags.rs index de6afbde..ca2f9766 100644 --- a/tests/flags.rs +++ b/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() { diff --git a/xtask/Cargo.lock b/xtask/Cargo.lock index 3ea263df..398fc7ad 100644 --- a/xtask/Cargo.lock +++ b/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", ] diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index 4d0d24a6..e501424f 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -8,3 +8,4 @@ publish = false [dependencies] clap = "4.5.21" clap_mangen = "0.2.24" +parse-size = "1.1.0"