From 9a98a31a4609bcc010a3c2ebfbf390662f042f8f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kornel=20Lesin=CC=81ski?= <kornel@geekhood.net>
Date: Fri, 1 Jun 2018 15:06:16 +0100
Subject: [PATCH] deflate() error handling

---
 src/deflate/mod.rs | 22 ++++++++++++++--------
 src/lib.rs         |  5 ++++-
 src/png/mod.rs     | 19 ++++++++++---------
 3 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/src/deflate/mod.rs b/src/deflate/mod.rs
index f63f75e0..4d85090b 100644
--- a/src/deflate/mod.rs
+++ b/src/deflate/mod.rs
@@ -13,7 +13,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) -> Vec<u8> {
+pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8) -> Result<Vec<u8>, PngError> {
     #[cfg(feature = "cfzlib")]
     {
         if is_cfzlib_supported() {
@@ -21,7 +21,7 @@ pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8) -> Vec<u8> {
         }
     }
 
-    miniz_stream::compress_to_vec_oxipng(data, zc, zw.into(), zs.into())
+    Ok(miniz_stream::compress_to_vec_oxipng(data, zc, zw.into(), zs.into()))
 }
 
 #[cfg(feature = "cfzlib")]
@@ -40,20 +40,22 @@ fn is_cfzlib_supported() -> bool {
 }
 
 #[cfg(feature = "cfzlib")]
-pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8) -> Vec<u8> {
+pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8) -> Result<Vec<u8>, PngError> {
     use std::mem;
     use cloudflare_zlib_sys::*;
 
     assert!(data.len() < u32::max_value() as usize);
     unsafe {
         let mut stream = mem::zeroed();
-        assert_eq!(Z_OK, deflateInit2(
+        if Z_OK != deflateInit2(
                     &mut stream,
                     level.into(),
                     Z_DEFLATED,
                     window_bits.into(),
                     MAX_MEM_LEVEL,
-                    strategy.into()));
+                    strategy.into()) {
+            return Err(PngError::new("deflateInit2"));
+        }
 
         let max_size = deflateBound(&mut stream, data.len() as uLong) as usize;
         // it's important to have the capacity pre-allocated,
@@ -65,11 +67,15 @@ pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8) ->
         stream.avail_in = data.len() as uInt;
         stream.next_out = out.as_mut_ptr();
         stream.avail_out = out.capacity() as uInt;
-        assert_eq!(Z_STREAM_END, deflate(&mut stream, Z_FINISH));
-        assert_eq!(Z_OK, deflateEnd(&mut stream));
+        if Z_STREAM_END != deflate(&mut stream, Z_FINISH) {
+            return Err(PngError::new("deflate"));
+        }
+        if Z_OK != deflateEnd(&mut stream) {
+            return Err(PngError::new("deflateEnd"));
+        }
         debug_assert!(stream.total_out as usize <= out.capacity());
         out.set_len(stream.total_out as usize);
-        return out;
+        return Ok(out);
     }
 }
 
diff --git a/src/lib.rs b/src/lib.rs
index 5dd134e5..a49eaf2a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -501,7 +501,10 @@ fn optimize_png(
                 let new_idat = if opts.deflate == Deflaters::Zlib {
                     deflate::deflate(filtered, trial.compression, trial.strategy, opts.window)
                 } else {
-                    deflate::zopfli_deflate(filtered).unwrap()
+                    deflate::zopfli_deflate(filtered)
+                };
+                let new_idat = if let Ok(n) = new_idat {n} else {
+                    return None;
                 };
 
                 if opts.verbosity == Some(1) {
diff --git a/src/png/mod.rs b/src/png/mod.rs
index 0cf6d7af..12fdb3e2 100644
--- a/src/png/mod.rs
+++ b/src/png/mod.rs
@@ -603,27 +603,28 @@ impl PngData {
         assert!(!alphas.is_empty());
         let best = alphas
             .iter()
-            .map(|alpha| {
+            .filter_map(|alpha| {
                 let mut image = self.clone();
                 image.reduce_alpha_channel(*alpha);
-                let size = STD_FILTERS
+                STD_FILTERS
                     .iter()
-                    .map(|f| {
+                    .filter_map(|f| {
                         deflate::deflate(
                             &image.filter_image(*f),
                             STD_COMPRESSION,
                             STD_STRATEGY,
                             STD_WINDOW,
-                        ).len()
+                        ).ok()
+                        .as_ref().map(|l| l.len())
                     })
                     .min()
-                    .unwrap();
-                (size, image)
+                    .map(|size| (size, image))
             })
-            .min_by_key(|&(size, _)| size)
-            .unwrap();
+            .min_by_key(|&(size, _)| size);
 
-        self.raw_data = best.1.raw_data;
+        if let Some(best) = best {
+            self.raw_data = best.1.raw_data;
+        }
     }
 
     pub fn reduce_alpha_channel(&mut self, optim: AlphaOptim) -> bool {