Improved limiting of max gzip size (#116)

src/deflate/miniz_stream.rs

@@ -1,12 +1,13 @@
+use atomicmin::AtomicMin;
 use error::PngError;
 use miniz_oxide::deflate::core::*;
 
-pub fn compress_to_vec_oxipng(input: &[u8], level: u8, window_bits: i32, strategy: i32, max_size: Option<usize>) -> Result<Vec<u8>, PngError> {
+pub fn compress_to_vec_oxipng(input: &[u8], level: u8, window_bits: i32, strategy: i32, max_size: &AtomicMin) -> Result<Vec<u8>, PngError> {
     // The comp flags function sets the zlib flag if the window_bits parameter is > 0.
     let flags = create_comp_flags_from_zip_params(level.into(), window_bits, strategy);
     let mut compressor = CompressorOxide::new(flags);
     // if max size is known, then expect that much data (but no more than input.len())
-    let mut output = Vec::with_capacity(max_size.unwrap_or(input.len() / 2).min(input.len()));
+    let mut output = Vec::with_capacity(max_size.get().unwrap_or(input.len() / 2).min(input.len()));
     // # Unsafe
     // We trust compress to not read the uninitialized bytes.
     unsafe {
@@ -32,7 +33,7 @@ pub fn compress_to_vec_oxipng(input: &[u8], level: u8, window_bits: i32, strateg
                 break;
             }
             TDEFLStatus::Okay => {
-                if let Some(max) = max_size {
+                if let Some(max) = max_size.get() {
                     if output.len() > max {
                         return Err(PngError::DeflatedDataTooLong(output.len()))
                     }

src/deflate/mod.rs

@@ -1,3 +1,4 @@
+use atomicmin::AtomicMin;
 use error::PngError;
 use miniz_oxide;
 use std::cmp::max;
@@ -13,7 +14,7 @@ pub fn inflate(data: &[u8]) -> Result<Vec<u8>, PngError> {
 }
 
 /// Compress a data stream using the DEFLATE algorithm
-pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8, max_size: Option<usize>) -> Result<Vec<u8>, PngError> {
+pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8, max_size: &AtomicMin) -> Result<Vec<u8>, PngError> {
     #[cfg(feature = "cfzlib")]
     {
         if is_cfzlib_supported() {
@@ -40,7 +41,7 @@ fn is_cfzlib_supported() -> bool {
 }
 
 #[cfg(feature = "cfzlib")]
-pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8, max_size: Option<usize>) -> Result<Vec<u8>, PngError> {
+pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8, max_size: &AtomicMin) -> Result<Vec<u8>, PngError> {
     use std::mem;
     use cloudflare_zlib_sys::*;
 
@@ -58,7 +59,7 @@ pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8, max
         }
 
         let upper_bound = deflateBound(&mut stream, data.len() as uLong) as usize;
-        let max_size = max_size.unwrap_or(upper_bound).min(upper_bound);
+        let max_size = max_size.get().unwrap_or(upper_bound).min(upper_bound);
         // it's important to have the capacity pre-allocated,
         // as unsafe set_len is called later
         let mut out = Vec::with_capacity(max_size);
@@ -70,7 +71,7 @@ pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8, max
         stream.avail_out = out.capacity() as uInt;
         match deflate(&mut stream, Z_FINISH) {
             Z_STREAM_END => {},
-            Z_OK => return Err(PngError::DeflatedDataTooLong(max_size)),
+            Z_OK => return Err(PngError::DeflatedDataTooLong(stream.total_out as usize)),
             _ => return Err(PngError::new("deflate")),
         }
         if Z_OK != deflateEnd(&mut stream) {

src/lib.rs

@@ -506,7 +506,7 @@ fn optimize_png(
             .filter_map(|trial| {
                 let filtered = &filters[&trial.filter];
                 let new_idat = if opts.deflate == Deflaters::Zlib {
-                    deflate::deflate(filtered, trial.compression, trial.strategy, opts.window, best_size.get())
+                    deflate::deflate(filtered, trial.compression, trial.strategy, opts.window, &best_size)
                 } else {
                     deflate::zopfli_deflate(filtered)
                 };

src/png/mod.rs

@@ -625,7 +625,7 @@ impl PngData {
                             STD_COMPRESSION,
                             STD_STRATEGY,
                             STD_WINDOW,
-                            best_size.get(),
+                            &best_size,
                         ).ok()
                         .as_ref().map(|l| {
                             best_size.set_min(l.len());