From 99a1a75f37ee436549659845a2b9371dea80f1cb Mon Sep 17 00:00:00 2001
From: oyvindln <oyvindln@users.noreply.github.com>
Date: Sat, 4 Jul 2020 18:46:55 +0200
Subject: [PATCH] fix: Avoid infinitely looping on sync flush with short buffer
 writers

Closes #47
---
 src/compress.rs      | 14 ++++++++++++--
 src/deflate_state.rs |  8 ++++++++
 src/writer.rs        |  3 +++
 tests/test.rs        | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/src/compress.rs b/src/compress.rs
index 034a406..a771ee3 100644
--- a/src/compress.rs
+++ b/src/compress.rs
@@ -121,7 +121,7 @@ pub fn compress_data_dynamic_n<W: Write>(
         }
 
         if deflate_state.lz77_state.is_last_block() {
-            // The last block has already been written, so we don't ave anything to compress.
+            // The last block has already been written, so we don't have anything to compress.
             break;
         }
 
@@ -133,6 +133,11 @@ pub fn compress_data_dynamic_n<W: Write>(
             flush,
         );
 
+        if written > 0 {
+            // Reset sync status if we compress more data.
+            deflate_state.sync_was_output_last = false;
+        }
+
         // Bytes written in this call
         bytes_written += written;
         // Total bytes written since the compression process started
@@ -253,7 +258,12 @@ pub fn compress_data_dynamic_n<W: Write>(
         if status == LZ77Status::Finished {
             // This flush mode means that there should be an empty stored block at the end.
             if flush == Flush::Sync {
-                write_stored_block(&[], &mut deflate_state.encoder_state.writer, false);
+                // Only write this if we didn't already write it to the output buffer in a previous
+                // call.
+                if deflate_state.sync_was_output_last == false {
+                    write_stored_block(&[], &mut deflate_state.encoder_state.writer, false);
+                    deflate_state.sync_was_output_last = true;
+                }
             } else if !deflate_state.lz77_state.is_last_block() {
                 // Make sure a block with the last block header has been output.
                 // Not sure this can actually happen, but we make sure to finish properly
diff --git a/src/deflate_state.rs b/src/deflate_state.rs
index 4e25fb8..b126417 100644
--- a/src/deflate_state.rs
+++ b/src/deflate_state.rs
@@ -88,6 +88,13 @@ pub struct DeflateState<W: Write> {
     /// Number of bytes written as calculated by sum of block input lengths.
     /// Used to check that they are correct when `debug_assertions` are enabled.
     pub bytes_written_control: DebugCounter,
+    /// Whether the last data written to the output buffer was a sync flush.
+    /// Need to keep track of this in order to avoid infinitely looping on writers
+    /// that can't output the sync flush bytes in one go.
+    /// This is implemented in a somewhat clunky manner at the moment,
+    /// ideally it should be done in a more fail-safe way to avoid
+    /// further bugs.
+    pub sync_was_output_last: bool,
 }
 
 impl<W: Write> DeflateState<W> {
@@ -108,6 +115,7 @@ impl<W: Write> DeflateState<W> {
             output_buf_pos: 0,
             flush_mode: Flush::None,
             bytes_written_control: DebugCounter::default(),
+            sync_was_output_last: false,
         }
     }
 
diff --git a/src/writer.rs b/src/writer.rs
index f31d437..bc90fd9 100644
--- a/src/writer.rs
+++ b/src/writer.rs
@@ -51,6 +51,9 @@ pub fn compress_until_done<W: Write>(
         }
     }
 
+    // Reset sync status so it's possible to sync again later.
+    deflate_state.sync_was_output_last = false;
+
     debug_assert_eq!(
         deflate_state.bytes_written,
         deflate_state.bytes_written_control.get()
diff --git a/tests/test.rs b/tests/test.rs
index f2a8154..97788f5 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -159,3 +159,39 @@ fn afl_regressions_default_compression() {
         }
     }
 }
+
+
+mod issue_47 {
+    use std::io::{self, Write};
+
+    #[test]
+    fn issue_47() {
+        let _ = deflate::write::ZlibEncoder::new(SmallWriter::new(vec![], 2), deflate::Compression::Fast).flush();
+    }
+
+    struct SmallWriter<W: Write> {
+        writer: W,
+        small: usize,
+    }
+
+    impl<W: Write> SmallWriter<W> {
+        fn new(writer: W, buf_len: usize) -> SmallWriter<W> {
+            SmallWriter {
+                writer,
+                small: buf_len,
+            }
+        }
+    }
+
+    impl<W: Write> Write for SmallWriter<W> {
+        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+            // Never write more than `small` bytes at a time.
+            let small = buf.len().min(self.small);
+            self.writer.write(&buf[..small])
+        }
+
+        fn flush(&mut self) -> io::Result<()> {
+            Ok(())
+        }
+    }
+}