Zlib-rs round 2

Cargo.toml

@@ -26,7 +26,7 @@ required-features = ["zlib-rs"]
 libz-sys = { version = "1.1.20", optional = true, default-features = false }
 libz-ng-sys = { version = "1.1.16", optional = true }
 # this matches the default features, but we don't want to depend on the default features staying the same
-zlib-rs = { version = "0.5.3", optional = true, default-features = false, features = ["std", "rust-allocator"] }
+zlib-rs = { version = "0.5.4", optional = true, default-features = false, features = ["std", "rust-allocator"] }
 cloudflare-zlib-sys = { version = "0.3.6", optional = true }
 miniz_oxide = { version = "0.8.5", optional = true, default-features = false, features = ["with-alloc", "simd"] }
 crc32fast = "1.2.0"

src/ffi/zlib_rs.rs

@@ -30,6 +30,7 @@ pub const MZ_FINISH: isize = DeflateFlush::Finish as isize;
 pub const MZ_DEFAULT_WINDOW_BITS: core::ffi::c_int = 15;
 
 use super::*;
+use crate::mem::{compress_failed, decompress_failed};
 
 impl From<::zlib_rs::Status> for crate::mem::Status {
     fn from(value: ::zlib_rs::Status) -> Self {
@@ -52,15 +53,18 @@ impl ErrorMessage {
 
 pub struct Inflate {
     pub(crate) inner: ::zlib_rs::Inflate,
+    // NOTE: these counts do not count the dictionary.
+    total_in: u64,
+    total_out: u64,
 }
 
 impl fmt::Debug for Inflate {
     fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
         write!(
             f,
             "zlib_rs inflate internal state. total_in: {}, total_out: {}",
-            self.inner.total_in(),
-            self.inner.total_out(),
+            self.total_in(),
+            self.total_out(),
         )
     }
 }
@@ -79,6 +83,8 @@ impl InflateBackend for Inflate {
     fn make(zlib_header: bool, window_bits: u8) -> Self {
         Inflate {
             inner: ::zlib_rs::Inflate::new(zlib_header, window_bits),
+            total_in: 0,
+            total_out: 0,
         }
     }
 
@@ -94,41 +100,67 @@ impl InflateBackend for Inflate {
             FlushDecompress::Finish => InflateFlush::Finish,
         };
 
-        match self.inner.decompress(input, output, flush) {
+        let total_in_start = self.inner.total_in();
+        let total_out_start = self.inner.total_out();
+
+        let result = self.inner.decompress(input, output, flush);
+
+        self.total_in += self.inner.total_in() - total_in_start;
+        self.total_out += self.inner.total_out() - total_out_start;
+
+        match result {
             Ok(status) => Ok(status.into()),
             Err(InflateError::NeedDict { dict_id }) => crate::mem::decompress_need_dict(dict_id),
-            Err(e) => crate::mem::decompress_failed(ErrorMessage(Some(e.as_str()))),
+            Err(_) => self.decompress_error(),
         }
     }
 
     fn reset(&mut self, zlib_header: bool) {
+        self.total_in = 0;
+        self.total_out = 0;
         self.inner.reset(zlib_header);
     }
 }
 
 impl Backend for Inflate {
     #[inline]
     fn total_in(&self) -> u64 {
-        self.inner.total_in()
+        self.total_in
     }
 
     #[inline]
     fn total_out(&self) -> u64 {
-        self.inner.total_out()
+        self.total_out
+    }
+}
+
+impl Inflate {
+    fn decompress_error<T>(&self) -> Result<T, DecompressError> {
+        decompress_failed(ErrorMessage(self.inner.error_message()))
+    }
+
+    pub fn set_dictionary(&mut self, dictionary: &[u8]) -> Result<u32, DecompressError> {
+        match self.inner.set_dictionary(dictionary) {
+            Ok(v) => Ok(v),
+            Err(_) => self.decompress_error(),
+        }
     }
 }
 
 pub struct Deflate {
     pub(crate) inner: ::zlib_rs::Deflate,
+    // NOTE: these counts do not count the dictionary.
+    total_in: u64,
+    total_out: u64,
 }
 
 impl fmt::Debug for Deflate {
     fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
         write!(
             f,
             "zlib_rs deflate internal state. total_in: {}, total_out: {}",
-            self.inner.total_in(),
-            self.inner.total_out(),
+            self.total_in(),
+            self.total_out(),
         )
     }
 }
@@ -140,6 +172,8 @@ impl DeflateBackend for Deflate {
 
         Deflate {
             inner: ::zlib_rs::Deflate::new(level.level() as i32, zlib_header, window_bits),
+            total_in: 0,
+            total_out: 0,
         }
     }
 
@@ -157,25 +191,64 @@ impl DeflateBackend for Deflate {
             FlushCompress::Finish => DeflateFlush::Finish,
         };
 
-        match self.inner.compress(input, output, flush) {
+        let total_in_start = self.inner.total_in();
+        let total_out_start = self.inner.total_out();
+
+        let result = self.inner.compress(input, output, flush);
+
+        self.total_in += self.inner.total_in() - total_in_start;
+        self.total_out += self.inner.total_out() - total_out_start;
+
+        match result {
             Ok(status) => Ok(status.into()),
-            Err(e) => crate::mem::compress_failed(ErrorMessage(Some(e.as_str()))),
+            Err(_) => self.compress_error(),
         }
     }
 
     fn reset(&mut self) {
+        self.total_in = 0;
+        self.total_out = 0;
         self.inner.reset();
     }
 }
 
 impl Backend for Deflate {
     #[inline]
     fn total_in(&self) -> u64 {
-        self.inner.total_in()
+        self.total_in
     }
 
     #[inline]
     fn total_out(&self) -> u64 {
-        self.inner.total_out()
+        self.total_out
+    }
+}
+
+impl Deflate {
+    fn compress_error<T>(&self) -> Result<T, CompressError> {
+        compress_failed(ErrorMessage(self.inner.error_message()))
+    }
+
+    pub fn set_dictionary(&mut self, dictionary: &[u8]) -> Result<u32, CompressError> {
+        match self.inner.set_dictionary(dictionary) {
+            Ok(v) => Ok(v),
+            Err(_) => self.compress_error(),
+        }
+    }
+
+    pub fn set_level(&mut self, level: Compression) -> Result<(), CompressError> {
+        use ::zlib_rs::Status;
+
+        match self.inner.set_level(level.level() as i32) {
+            Ok(status) => match status {
+                Status::Ok => Ok(()),
+
+                Status::BufError => compress_failed(ErrorMessage(Some("insufficient space"))),
+
+                Status::StreamEnd => unreachable!(),
+            },
+
+            Err(_) => self.compress_error(),
+        }
     }
 }

src/mem.rs

@@ -286,6 +286,14 @@ impl Compress {
         }
     }
 
+    /// Specifies the compression dictionary to use.
+    ///
+    /// Returns the Adler-32 checksum of the dictionary.
+    #[cfg(all(not(feature = "any_zlib"), feature = "zlib-rs"))]
+    pub fn set_dictionary(&mut self, dictionary: &[u8]) -> Result<u32, CompressError> {
+        self.inner.set_dictionary(dictionary)
+    }
+
     /// Quickly resets this compressor without having to reallocate anything.
     ///
     /// This is equivalent to dropping this object and then creating a new one.
@@ -303,22 +311,31 @@ impl Compress {
     /// the compression of the available input data before changing the
     /// compression level. Flushing the stream before calling this method
     /// ensures that the function will succeed on the first call.
-    #[cfg(feature = "any_zlib")]
+    #[cfg(any(feature = "any_zlib", feature = "zlib-rs"))]
     pub fn set_level(&mut self, level: Compression) -> Result<(), CompressError> {
-        use std::os::raw::c_int;
-        // SAFETY: The field `inner` must always be accessed as a raw pointer,
-        // since it points to a cyclic structure. No copies of `inner` can be
-        // retained for longer than the lifetime of `self.inner.inner.stream_wrapper`.
-        let stream = self.inner.inner.stream_wrapper.inner;
-        unsafe {
-            (*stream).msg = std::ptr::null_mut();
+        #[cfg(all(not(feature = "any_zlib"), feature = "zlib-rs"))]
+        {
+            self.inner.set_level(level)
         }
-        let rc = unsafe { ffi::deflateParams(stream, level.0 as c_int, ffi::MZ_DEFAULT_STRATEGY) };
 
-        match rc {
-            ffi::MZ_OK => Ok(()),
-            ffi::MZ_BUF_ERROR => compress_failed(self.inner.inner.msg()),
-            c => panic!("unknown return code: {}", c),
+        #[cfg(feature = "any_zlib")]
+        {
+            use std::os::raw::c_int;
+            // SAFETY: The field `inner` must always be accessed as a raw pointer,
+            // since it points to a cyclic structure. No copies of `inner` can be
+            // retained for longer than the lifetime of `self.inner.inner.stream_wrapper`.
+            let stream = self.inner.inner.stream_wrapper.inner;
+            unsafe {
+                (*stream).msg = std::ptr::null_mut();
+            }
+            let rc =
+                unsafe { ffi::deflateParams(stream, level.0 as c_int, ffi::MZ_DEFAULT_STRATEGY) };
+
+            match rc {
+                ffi::MZ_OK => Ok(()),
+                ffi::MZ_BUF_ERROR => compress_failed(self.inner.inner.msg()),
+                c => panic!("unknown return code: {}", c),
+            }
         }
     }
 
@@ -540,6 +557,12 @@ impl Decompress {
         }
     }
 
+    /// Specifies the decompression dictionary to use.
+    #[cfg(all(not(feature = "any_zlib"), feature = "zlib-rs"))]
+    pub fn set_dictionary(&mut self, dictionary: &[u8]) -> Result<u32, DecompressError> {
+        self.inner.set_dictionary(dictionary)
+    }
+
     /// Performs the equivalent of replacing this decompression state with a
     /// freshly allocated copy.
     ///
@@ -806,7 +829,7 @@ mod tests {
         assert_eq!(&decoded[..decoder.total_out() as usize], string);
     }
 
-    #[cfg(feature = "any_zlib")]
+    #[cfg(any(feature = "any_zlib", feature = "zlib-rs"))]
     #[test]
     fn test_error_message() {
         let mut decoder = Decompress::new(false);