From dece0383469198e71a4a610baf71b9db2db05d1f Mon Sep 17 00:00:00 2001
From: Jonathan Giddy <jongiddy@gmail.com>
Date: Sat, 10 Jun 2023 07:23:03 +0100
Subject: [PATCH] Fix infinite loop after invalid data

Closes #98
---
 src/write.rs | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/write.rs b/src/write.rs
index 5d80000e..51519579 100644
--- a/src/write.rs
+++ b/src/write.rs
@@ -254,11 +254,17 @@ impl<W: Write> Write for BzDecoder<W> {
             let res = self.data.decompress_vec(data, &mut self.buf);
             let written = (self.total_in() - before) as usize;
 
-            let res = res.map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
-
-            if res == Status::StreamEnd {
-                self.done = true;
+            match res {
+                Err(e) => {
+                    self.done = true;
+                    return Err(io::Error::new(io::ErrorKind::InvalidInput, e));
+                }
+                Ok(Status::StreamEnd) => {
+                    self.done = true;
+                }
+                Ok(_) => {}
             }
+
             if written > 0 || data.is_empty() || self.done {
                 return Ok(written);
             }
@@ -309,6 +315,14 @@ mod tests {
         assert_eq!(&data[..], b"");
     }
 
+    // https://github.com/alexcrichton/bzip2-rs/issues/98
+    #[test]
+    fn write_invalid() {
+        let mut d = BzDecoder::new(Vec::new());
+        let e = d.write(b"BZh\xfb").unwrap_err();
+        assert_eq!(e.kind(), std::io::ErrorKind::InvalidInput);
+    }
+
     #[test]
     fn qc() {
         ::quickcheck::quickcheck(test as fn(_) -> _);