track and expose the number of codes used

Author: folkertdev

libz-rs-sys/src/lib.rs

@@ -908,8 +908,11 @@ pub unsafe extern "C-unwind" fn inflateResetKeep(strm: *mut z_stream) -> c_int {
 /// - `buf` is `NULL`
 /// - `buf` and `len` satisfy the requirements of [`core::slice::from_raw_parts`]
 #[cfg_attr(feature = "export-symbols", export_name = prefix!(inflateCodesUsed))]
-pub unsafe extern "C-unwind" fn inflateCodesUsed(_strm: *mut z_stream) -> c_ulong {
-    todo!()
+pub unsafe extern "C-unwind" fn inflateCodesUsed(strm: *mut z_stream) -> c_ulong {
+    match InflateStream::from_stream_mut(strm) {
+        Some(stream) => zlib_rs::inflate::codes_used(stream) as c_ulong,
+        None => c_ulong::MAX,
+    }
 }
 
 /// Compresses as much data as possible, and stops when the input buffer becomes empty or the output buffer becomes full.

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

@@ -2209,3 +2209,73 @@ fn blow_up_the_stack_2() {
     let (_, err) = uncompress_slice(&mut output_rs, INPUT, config);
     assert_eq!(err, ReturnCode::DataError);
 }
+
+#[test]
+fn codes_used() {
+    // -1 is returned on NULL
+    assert_eq_rs_ng!({ inflateCodesUsed(core::ptr::null_mut()) });
+    assert_eq!(
+        unsafe { libz_rs_sys::inflateCodesUsed(core::ptr::null_mut()) },
+        c_ulong::MAX
+    );
+
+    let inputs = &[
+        include_bytes!("test-data/compression-corpus/The fastest WASM zlib.md.gzip-0.gz")
+            .as_slice(),
+        include_bytes!("test-data/compression-corpus/The fastest WASM zlib.md.gzip-9.gz")
+            .as_slice(),
+        include_bytes!("test-data/compression-corpus/The fastest WASM zlib.md.gzip-filtered-9.gz")
+            .as_slice(),
+        include_bytes!("test-data/compression-corpus/The fastest WASM zlib.md.gzip-fixed-9.gz")
+            .as_slice(),
+        include_bytes!("test-data/compression-corpus/The fastest WASM zlib.md.gzip-huffman-9.gz")
+            .as_slice(),
+        include_bytes!("test-data/compression-corpus/The fastest WASM zlib.md.gzip-rle-9.gz")
+            .as_slice(),
+    ];
+
+    for input in inputs {
+        assert_eq_rs_ng!({
+            let mut output: Vec<u8> = vec![0u8; 1 << 15];
+            let mut codes_used = Vec::new();
+
+            let mut stream = MaybeUninit::<z_stream>::zeroed();
+
+            let err = unsafe {
+                inflateInit2_(
+                    stream.as_mut_ptr(),
+                    16 + 15,
+                    zlibVersion(),
+                    core::mem::size_of::<z_stream>() as c_int,
+                )
+            };
+            assert_eq!(ReturnCode::from(err), ReturnCode::Ok);
+
+            let stream = unsafe { stream.assume_init_mut() };
+            codes_used.push(inflateCodesUsed(stream));
+
+            stream.next_out = output.as_mut_ptr();
+            stream.avail_out = output.len() as _;
+
+            for chunk in input.chunks(16) {
+                stream.next_in = chunk.as_ptr().cast_mut();
+                stream.avail_in = chunk.len() as _;
+
+                let err = unsafe { inflate(stream, InflateFlush::NoFlush as _) };
+                codes_used.push(inflateCodesUsed(stream));
+
+                if err == ReturnCode::StreamEnd as i32 {
+                    break;
+                }
+
+                assert_eq!(err, ReturnCode::Ok as i32);
+            }
+
+            let err = inflateEnd(stream);
+            assert_eq!(ReturnCode::from(err), ReturnCode::Ok);
+
+            output.truncate(stream.total_out as usize);
+            (output, codes_used)
+        });
+    }
+}

test-libz-rs-sys/src/lib.rs

@@ -1,14 +1,14 @@
 #[cfg(test)]
 mod deflate;
 pub mod end_to_end;
+#[cfg(test)]
+#[cfg(feature = "gz")]
+mod gz;
 mod helpers;
 #[cfg(test)]
 mod inflate;
 #[cfg(test)]
 mod zlib_ng_cve;
-#[cfg(test)]
-#[cfg(feature = "gz")]
-mod gz;
 
 #[cfg(test)]
 #[macro_export]
@@ -21,6 +21,12 @@ macro_rules! assert_eq_rs_ng {
         let _ng = unsafe {
             use libz_sys::*;
 
+            // this function is not exposed by `libz_sys`
+            extern "C" {
+                #[allow(unused)]
+                fn inflateCodesUsed(strm: *mut z_stream) -> core::ffi::c_ulong;
+            }
+
             $tt
         };
 

test-libz-rs-sys/src/test-data/compression-corpus/The fastest WASM zlib.md.gzip-0.gz

@@ -1587,7 +1587,7 @@ impl State<'_> {
                             self.have += 1;
                         }
 
-                        let InflateTable::Success(root) = inflate_table(
+                        let InflateTable::Success { root, used } = inflate_table(
                             CodeType::Codes,
                             &self.lens,
                             19,
@@ -1599,6 +1599,7 @@ impl State<'_> {
                             break 'label self.bad("invalid code lengths set\0");
                         };
 
+                        self.next = used;
                         self.len_table.codes = Codes::Codes;
                         self.len_table.bits = root;
 
@@ -1685,7 +1686,7 @@ impl State<'_> {
 
                         // build code tables
 
-                        let InflateTable::Success(root) = inflate_table(
+                        let InflateTable::Success { root, used } = inflate_table(
                             CodeType::Lens,
                             &self.lens,
                             self.nlen,
@@ -1699,8 +1700,9 @@ impl State<'_> {
 
                         self.len_table.codes = Codes::Len;
                         self.len_table.bits = root;
+                        self.next = used;
 
-                        let InflateTable::Success(root) = inflate_table(
+                        let InflateTable::Success { root, used } = inflate_table(
                             CodeType::Dists,
                             &self.lens[self.nlen..],
                             self.ndist,
@@ -1714,6 +1716,7 @@ impl State<'_> {
 
                         self.dist_table.bits = root;
                         self.dist_table.codes = Codes::Dist;
+                        self.next += used;
 
                         mode = Mode::Len_;
 
@@ -2215,6 +2218,10 @@ pub fn reset_keep(stream: &mut InflateStream) -> ReturnCode {
     ReturnCode::Ok
 }
 
+pub fn codes_used(stream: &InflateStream) -> usize {
+    stream.state.next
+}
+
 pub unsafe fn inflate(stream: &mut InflateStream, flush: InflateFlush) -> ReturnCode {
     if stream.next_out.is_null() || (stream.next_in.is_null() && stream.avail_in != 0) {
         return ReturnCode::StreamError as _;