Fix inflate() flush modes and implement state.data_type

libz-rs-sys/src/tests/inflate.rs

@@ -1698,3 +1698,115 @@ fn op_len_edge_case() {
 
     assert_eq!(output_rs, output_ng);
 }
+
+// Fills the provided buffer with pseudorandom bytes based on the given seed
+// Duplicates bytes by `step` in a row
+fn prng_bytes(seed: u64, bytes: &mut [u8], step: usize) {
+    const M: u64 = 2u64.pow(32);
+    const A: u64 = 1664525;
+    const C: u64 = 1013904223;
+    let mut state = seed;
+    for chunk in bytes.chunks_mut(4 * step) {
+        state = (A * state + C) % M;
+        let rand_bytes = state.to_le_bytes();
+        for (i, byte) in chunk.iter_mut().enumerate() {
+            *byte = rand_bytes[i / step];
+        }
+    }
+}
+
+#[test]
+fn test_inflate_flush_block() {
+    let window_bits = -15; // Raw
+    const CHUNK: usize = 16384;
+
+    // Create a compressed vector of random data that's bigger then the zlib block size
+    let mut data = vec![0u8; 160000];
+    prng_bytes(314159, &mut data, 4);
+    let config = DeflateConfig {
+        window_bits,
+        ..DeflateConfig::default()
+    };
+    let mut output = vec![0u8; 80000];
+    // Compress the data
+    let (compressed_data, return_code) = compress_slice(&mut output, &data, config);
+    assert_eq!(return_code, ReturnCode::Ok);
+
+    // Log the stream positions and data_type output from libz
+    let mut zlib_log = Vec::new();
+    {
+        let mut stream = MaybeUninit::<libz_sys::z_stream>::zeroed();
+
+        let ret = unsafe {
+            libz_sys::inflateInit2_(
+                stream.as_mut_ptr(),
+                window_bits,
+                libz_sys::zlibVersion(),
+                core::mem::size_of::<libz_sys::z_stream>() as c_int,
+            )
+        };
+        assert_eq!(ReturnCode::from(ret), ReturnCode::Ok);
+
+        let mut output = vec![0u8; CHUNK * 2];
+        let stream = unsafe { stream.assume_init_mut() };
+        stream.next_in = compressed_data.as_ptr() as *mut u8;
+        stream.avail_in = compressed_data.len() as _;
+        loop {
+            stream.next_out = output.as_mut_ptr();
+            stream.avail_out = output.len() as _;
+
+            let ret = unsafe { libz_sys::inflate(stream, InflateFlush::Block as i32) };
+
+            let log = format!(
+                "In:{} Out:{} DT:{}",
+                stream.avail_in, stream.avail_out, stream.data_type
+            );
+            zlib_log.push(log);
+
+            assert_eq!(ReturnCode::from(ret), ReturnCode::Ok);
+
+            if stream.avail_in == 0 {
+                break;
+            }
+        }
+    }
+
+    // Log the stream positions and data_type output from libz_rs and compare
+    {
+        let mut stream = MaybeUninit::<z_stream>::zeroed();
+
+        let ret = unsafe {
+            inflateInit2_(
+                stream.as_mut_ptr(),
+                window_bits,
+                zlibVersion(),
+                core::mem::size_of::<z_stream>() as c_int,
+            )
+        };
+        assert_eq!(ReturnCode::from(ret), ReturnCode::Ok);
+
+        let mut output = vec![0u8; CHUNK * 2];
+        let stream = unsafe { stream.assume_init_mut() };
+        stream.next_in = compressed_data.as_ptr() as *mut u8;
+        stream.avail_in = compressed_data.len() as _;
+        loop {
+            stream.next_out = output.as_mut_ptr();
+            stream.avail_out = output.len() as _;
+
+            let ret = unsafe { inflate(stream, InflateFlush::Block as i32) };
+
+            let log = format!(
+                "In:{} Out:{} DT:{}",
+                stream.avail_in, stream.avail_out, stream.data_type
+            );
+            // Compare log entries
+            assert_eq!(zlib_log.remove(0), log);
+
+            assert_eq!(ReturnCode::from(ret), ReturnCode::Ok);
+
+            if stream.avail_in == 0 {
+                break;
+            }
+        }
+    }
+}

zlib-rs/src/inflate.rs

@@ -252,6 +252,7 @@ pub enum Mode {
     Stored,
     CopyBlock,
     Check,
+    Len_,
     Len,
     Lit,
     LenExt,
@@ -505,6 +506,7 @@ impl<'a> State<'a> {
             Mode::CopyBlock => self.copy_block(),
             Mode::Check => self.check(),
             Mode::Len => self.len(),
+            Mode::Len_ => self.len_(),
             Mode::LenExt => self.len_ext(),
             Mode::Lit => self.lit(),
             Mode::Dist => self.dist(),
@@ -965,13 +967,14 @@ impl<'a> State<'a> {
                     bits: 5,
                 };
 
+                self.mode = Mode::Len_;
+
                 self.bit_reader.drop_bits(2);
 
                 if let InflateFlush::Trees = self.flush {
                     self.inflate_leave(ReturnCode::Ok)
                 } else {
-                    self.mode = Mode::Len;
-                    self.len()
+                    self.len_()
                 }
             }
             2 => {
@@ -1046,6 +1049,11 @@ impl<'a> State<'a> {
         self.type_()
     }
 
+    fn len_(&mut self) -> ReturnCode {
+        self.mode = Mode::Len;
+        self.len()
+    }
+
     fn len(&mut self) -> ReturnCode {
         let avail_in = self.bit_reader.bytes_remaining();
         let avail_out = self.writer.remaining();
@@ -1463,13 +1471,13 @@ impl<'a> State<'a> {
         self.dist_table.bits = root;
         self.dist_table.codes = Codes::Dist;
 
-        self.mode = Mode::Len;
+        self.mode = Mode::Len_;
 
         if matches!(self.flush, InflateFlush::Trees) {
             return self.inflate_leave(ReturnCode::Ok);
         }
 
-        self.len()
+        self.len_()
     }
 
     fn dict_id(&mut self) -> ReturnCode {
@@ -1499,7 +1507,7 @@ impl<'a> State<'a> {
     }
 
     fn bad(&mut self, msg: &'static str) -> ReturnCode {
-        #[cfg(feature = "std")]
+        #[cfg(all(feature = "std", test))]
         dbg!(msg);
         self.error_message = Some(msg);
         self.inflate_leave(ReturnCode::DataError)
@@ -1511,6 +1519,22 @@ impl<'a> State<'a> {
         // actual logic is in `inflate` itself
         return_code
     }
+
+    /// Stored in the `z_stream.data_type` field
+    fn decoding_state(&self) -> i32 {
+        let bit_reader_bits = self.bit_reader.bits_in_buffer() as i32;
+        debug_assert!(bit_reader_bits < 64);
+
+        let last = if self.last { 64 } else { 0 };
+
+        let mode = match self.mode {
+            Mode::Type => 128,
+            Mode::Len_ | Mode::CopyBlock => 256,
+            _ => 0,
+        };
+
+        bit_reader_bits | last | mode
+    }
 }
 
 fn inflate_fast_help(state: &mut State, _start: usize) -> ReturnCode {
@@ -1881,6 +1905,13 @@ pub unsafe fn inflate(stream: &mut InflateStream, flush: InflateFlush) -> Return
 
     let state = &mut stream.state;
 
+    // skip check
+    if let Mode::Type = state.mode {
+        state.mode = Mode::TypeDo;
+    }
+
+    state.flush = flush;
+
     state.bit_reader.update_slice(source_slice);
     state.writer = ReadBuf::new(dest_slice);
 
@@ -1922,32 +1953,34 @@ pub unsafe fn inflate(stream: &mut InflateStream, flush: InflateFlush) -> Return
     if must_update_window {
         'blk: {
             // initialize the window if needed
-            if stream.state.window.size() == 0 {
-                match Window::new_in(&stream.alloc, stream.state.wbits) {
-                    Some(window) => stream.state.window = window,
+            if state.window.size() == 0 {
+                match Window::new_in(&stream.alloc, state.wbits) {
+                    Some(window) => state.window = window,
                     None => {
-                        stream.state.mode = Mode::Mem;
+                        state.mode = Mode::Mem;
                         err = ReturnCode::MemError;
                         break 'blk;
                     }
                 }
             }
 
-            stream.state.window.extend(
-                &stream.state.writer.filled()[..out_written],
-                stream.state.flags,
+            state.window.extend(
+                &state.writer.filled()[..out_written],
+                state.flags,
                 update_checksum,
-                &mut stream.state.checksum,
-                &mut stream.state.crc_fold,
+                &mut state.checksum,
+                &mut state.crc_fold,
             );
         }
     }
 
-    if let Some(msg) = stream.state.error_message {
+    if let Some(msg) = state.error_message {
         assert!(msg.ends_with(|c| c == '\0'));
         stream.msg = msg.as_ptr() as *mut u8 as *mut core::ffi::c_char;
     }
 
+    stream.data_type = state.decoding_state();
+
     if ((in_read == 0 && out_written == 0) || flush == InflateFlush::Finish as _)
         && err == (ReturnCode::Ok as _)
     {