From 1322416ba4dc387bce9873104ac0deeac11eb14d Mon Sep 17 00:00:00 2001
From: Kornel <kornel@geekhood.net>
Date: Sun, 5 Aug 2018 04:02:56 +0100
Subject: [PATCH] Cleanup on error (#132)

Closes #126
---
 src/deflate/mod.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/deflate/mod.rs b/src/deflate/mod.rs
index ae51b847..d14f8395 100644
--- a/src/deflate/mod.rs
+++ b/src/deflate/mod.rs
@@ -81,8 +81,14 @@ pub fn cfzlib_deflate(
         stream.avail_out = out.capacity() as uInt;
         match deflate(&mut stream, Z_FINISH) {
             Z_STREAM_END => {}
-            Z_OK => return Err(PngError::DeflatedDataTooLong(stream.total_out as usize)),
-            _ => return Err(PngError::new("deflate")),
+            Z_OK | Z_BUF_ERROR => {
+                deflateEnd(&mut stream);
+                return Err(PngError::DeflatedDataTooLong(stream.total_out as usize));
+            },
+            _ => {
+                deflateEnd(&mut stream);
+                return Err(PngError::new("deflate"));
+            },
         }
         if Z_OK != deflateEnd(&mut stream) {
             return Err(PngError::new("deflateEnd"));