From 4f46bacfdaa0cda95be636bcb16209cb9f4e497c Mon Sep 17 00:00:00 2001
From: Jonathan Behrens <fintelia@gmail.com>
Date: Sat, 16 Apr 2022 16:10:09 -0400
Subject: [PATCH] Fixes to PackBitsReader

---
 src/decoder/mod.rs    | 5 +----
 src/decoder/stream.rs | 8 ++++----
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs
index 750fff3..5923590 100644
--- a/src/decoder/mod.rs
+++ b/src/decoder/mod.rs
@@ -1141,10 +1141,7 @@ impl<R: Read + Seek> Decoder<R> {
             CompressionMethod::LZW => {
                 Box::new(LZWReader::new(reader, usize::try_from(compressed_length)?))
             }
-            CompressionMethod::PackBits => Box::new(PackBitsReader::new(
-                reader,
-                usize::try_from(compressed_length)?,
-            )),
+            CompressionMethod::PackBits => Box::new(PackBitsReader::new(reader, compressed_length)),
             CompressionMethod::Deflate | CompressionMethod::OldDeflate => {
                 Box::new(DeflateReader::new(reader))
             }
diff --git a/src/decoder/stream.rs b/src/decoder/stream.rs
index 05997cc..d5af64a 100644
--- a/src/decoder/stream.rs
+++ b/src/decoder/stream.rs
@@ -270,9 +270,9 @@ pub struct PackBitsReader<R: Read> {
 
 impl<R: Read> PackBitsReader<R> {
     /// Wraps a reader
-    pub fn new(reader: R, length: usize) -> Self {
+    pub fn new(reader: R, length: u64) -> Self {
         Self {
-            reader: reader.take(length as u64),
+            reader: reader.take(length),
             state: PackBitsReaderState::Header,
             count: 0,
         }
@@ -315,7 +315,7 @@ impl<R: Read> Read for PackBitsReader<R> {
         }
 
         self.count -= length;
-        if self.count <= buf.len() {
+        if self.count == 0 {
             self.state = PackBitsReaderState::Header;
         }
         return Ok(length);
@@ -391,7 +391,7 @@ mod test {
         let encoded_len = encoded.len();
 
         let buff = io::Cursor::new(encoded);
-        let mut decoder = PackBitsReader::new(buff, encoded_len);
+        let mut decoder = PackBitsReader::new(buff, encoded_len as u64);
 
         let mut decoded = Vec::new();
         decoder.read_to_end(&mut decoded).unwrap();