unsafe reworkings

Author: folkertdev

libz-rs-sys/src/lib.rs

@@ -90,7 +90,10 @@ pub unsafe extern "C" fn inflate(strm: *mut z_stream, flush: i32) -> i32 {
 
 pub unsafe extern "C" fn inflateEnd(strm: *mut z_stream) -> i32 {
     match InflateStream::from_stream_mut(strm) {
-        Some(stream) => zlib_rs::inflate::end(stream) as _,
+        Some(stream) => {
+            zlib_rs::inflate::end(stream);
+            ReturnCode::Ok as _
+        }
         None => ReturnCode::StreamError as _,
     }
 }
@@ -345,7 +348,10 @@ pub extern "C" fn compressBound(sourceLen: c_ulong) -> c_ulong {
 
 pub unsafe extern "C" fn deflateEnd(strm: *mut z_stream) -> i32 {
     match DeflateStream::from_stream_mut(strm) {
-        Some(stream) => zlib_rs::deflate::end(stream) as _,
+        Some(stream) => match zlib_rs::deflate::end(stream) {
+            Ok(_) => ReturnCode::Ok as _,
+            Err(_) => ReturnCode::DataError as _,
+        },
         None => ReturnCode::StreamError as _,
     }
 }

zlib-rs/src/allocate.rs

@@ -69,6 +69,7 @@ unsafe fn zng_free(ptr: *mut c_void) {
     unsafe { libc::free(ptr) };
 }
 
+#[derive(Clone, Copy)]
 #[repr(C)]
 pub(crate) struct Allocator<'a> {
     pub(crate) zalloc: crate::c_api::alloc_func,

zlib-rs/src/deflate.rs

@@ -35,9 +35,7 @@ pub struct DeflateStream<'a> {
     pub(crate) total_out: crate::c_api::z_size,
     pub(crate) msg: *const libc::c_char,
     pub(crate) state: &'a mut State<'a>,
-    pub(crate) zalloc: crate::c_api::alloc_func,
-    pub(crate) zfree: crate::c_api::free_func,
-    pub(crate) opaque: crate::c_api::voidpf,
+    pub(crate) alloc: Allocator<'a>,
     pub(crate) data_type: libc::c_int,
     pub(crate) adler: crate::c_api::z_checksum,
     pub(crate) reserved: crate::c_api::uLong,
@@ -74,6 +72,11 @@ impl<'a> DeflateStream<'a> {
 
         Some(stream)
     }
+
+    fn as_z_stream_mut(&mut self) -> &mut z_stream {
+        // safety: a valid &mut DeflateStream is also a valid &mut z_stream
+        unsafe { &mut *(self as *mut DeflateStream as *mut z_stream) }
+    }
 }
 
 /// number of elements in hash table
@@ -258,11 +261,11 @@ pub fn init(stream: &mut z_stream, config: DeflateConfig) -> ReturnCode {
                 if let Some(mut sym_buf) = sym_buf {
                     alloc.deallocate(sym_buf.as_mut_ptr(), sym_buf.capacity())
                 }
-                if let Some(mut pending) = pending {
+                if let Some(pending) = pending {
                     pending.drop_in(&alloc);
                 }
                 if let Some(head) = head {
-                    alloc.deallocate(head.as_mut_ptr(), HASH_SIZE)
+                    alloc.deallocate(head.as_mut_ptr(), 1)
                 }
                 if let Some(prev) = prev {
                     alloc.deallocate(prev.as_mut_ptr(), prev.len())
@@ -434,12 +437,7 @@ pub fn copy<'a>(
         core::ptr::copy_nonoverlapping(source, dest.as_mut_ptr(), 1);
     }
 
-    let alloc = Allocator {
-        zalloc: source.zalloc,
-        zfree: source.zfree,
-        opaque: source.opaque,
-        _marker: PhantomData,
-    };
+    let alloc = &source.alloc;
 
     // allocated here to have the same order as zlib
     let Some(state_allocation) = alloc.allocate::<State>() else {
@@ -448,13 +446,13 @@ pub fn copy<'a>(
 
     let source_state = &source.state;
 
-    let window = source_state.window.clone_in(&alloc);
+    let window = source_state.window.clone_in(alloc);
 
     let prev = alloc.allocate_slice::<u16>(source_state.w_size);
     let head = alloc.allocate::<[u16; HASH_SIZE]>();
 
-    let pending = source_state.pending.clone_in(&alloc);
-    let sym_buf = source_state.sym_buf.clone_in(&alloc);
+    let pending = source_state.pending.clone_in(alloc);
+    let sym_buf = source_state.sym_buf.clone_in(alloc);
 
     // if any allocation failed, clean up allocations that did succeed
     let (window, prev, head, pending, sym_buf) = match (window, prev, head, pending, sym_buf) {
@@ -471,8 +469,8 @@ pub fn copy<'a>(
                 if let Some(mut sym_buf) = sym_buf {
                     alloc.deallocate(sym_buf.as_mut_ptr(), sym_buf.capacity())
                 }
-                if let Some(mut pending) = pending {
-                    pending.drop_in(&alloc);
+                if let Some(pending) = pending {
+                    pending.drop_in(alloc);
                 }
                 if let Some(head) = head {
                     alloc.deallocate(head.as_mut_ptr(), HASH_SIZE)
@@ -481,7 +479,7 @@ pub fn copy<'a>(
                     alloc.deallocate(prev.as_mut_ptr(), prev.len())
                 }
                 if let Some(mut window) = window {
-                    window.drop_in(&alloc);
+                    window.drop_in(alloc);
                 }
 
                 alloc.deallocate(state_allocation.as_mut_ptr(), 1);
@@ -558,37 +556,37 @@ pub fn copy<'a>(
     ReturnCode::Ok
 }
 
-pub fn end(stream: &mut DeflateStream) -> ReturnCode {
+/// # Returns
+///
+/// - Err when deflate is not done. A common cause is insufficient output space
+/// - Ok otherwise
+pub fn end<'a>(stream: &'a mut DeflateStream) -> Result<&'a mut z_stream, &'a mut z_stream> {
     let status = stream.state.status;
 
-    let sym_buf = stream.state.sym_buf.as_mut_ptr();
-    let pending = stream.state.pending.as_mut_ptr();
-    let head = stream.state.head.as_mut_ptr();
-    let prev = stream.state.prev.as_mut_ptr();
-    let window = stream.state.window.as_mut_ptr();
-    let state = stream.state as *mut State;
+    let alloc = stream.alloc;
 
-    let opaque = stream.opaque;
-    let free = stream.zfree;
+    // deallocate in reverse order of allocations
+    unsafe {
+        // safety: we make sure that these fields are not used (by invalidating the state pointer)
+        stream.state.sym_buf.drop_in(&alloc);
+        stream.state.pending.drop_in(&alloc);
+        alloc.deallocate(stream.state.head, 1);
+        alloc.deallocate(stream.state.prev.as_mut_ptr(), stream.state.prev.len());
+        stream.state.window.drop_in(&alloc);
+    }
 
-    // safety: a valid &mut DeflateStream is also a valid &mut z_stream
-    let stream = unsafe { &mut *(stream as *mut DeflateStream as *mut z_stream) };
+    let state = stream.state as *mut State;
+    let stream = stream.as_z_stream_mut();
     stream.state = std::ptr::null_mut();
 
-    // deallocate in reverse order of allocations
+    // safety: `state` is not used later
     unsafe {
-        free(opaque, sym_buf.cast());
-        free(opaque, pending.cast());
-        free(opaque, head.cast());
-        free(opaque, prev.cast());
-        free(opaque, window.cast());
-
-        free(opaque, state.cast());
+        alloc.deallocate(state, 1);
     }
 
     match status {
-        Status::Busy => ReturnCode::DataError,
-        _ => ReturnCode::Ok,
+        Status::Busy => Err(stream),
+        _ => Ok(stream),
     }
 }
 
@@ -2561,7 +2559,8 @@ pub fn compress<'a>(
     };
 
     if let Some(stream) = unsafe { DeflateStream::from_stream_mut(&mut stream) } {
-        end(stream);
+        // error gets ignored here
+        let _ = end(stream);
     }
 
     (output_slice, ReturnCode::Ok)
@@ -3041,7 +3040,7 @@ mod test {
         assert_eq!(deflate(stream, Flush::NoFlush), ReturnCode::Ok);
 
         // but end is not
-        assert_eq!(end(stream), ReturnCode::DataError);
+        assert!(end(stream).is_err());
     }
 
     #[test]
@@ -3512,7 +3511,7 @@ mod test {
 
         let n = stream.total_out as usize;
 
-        assert_eq!(end(stream), ReturnCode::Ok);
+        assert!(end(stream).is_ok());
 
         let output_rs = &mut output[..n];
 
@@ -3567,7 +3566,7 @@ mod test {
 
         let n = stream.total_out as usize;
 
-        assert_eq!(end(stream), ReturnCode::Ok);
+        assert!(end(stream).is_ok());
 
         let output_rs = &mut output[..n];
 

zlib-rs/src/deflate/pending.rs

@@ -11,10 +11,6 @@ pub struct Pending<'a> {
 }
 
 impl<'a> Pending<'a> {
-    pub fn as_mut_ptr(&mut self) -> *mut u8 {
-        self.buf
-    }
-
     pub fn reset_keep(&mut self) {
         // keep the buffer as it is
         self.pending = 0;
@@ -92,7 +88,7 @@ impl<'a> Pending<'a> {
         Some(clone)
     }
 
-    pub(crate) unsafe fn drop_in(&mut self, alloc: &Allocator) {
+    pub(crate) unsafe fn drop_in(&self, alloc: &Allocator) {
         let len = self.end as usize - self.buf as usize;
         alloc.deallocate(self.buf, len);
     }

zlib-rs/src/deflate/window.rs

@@ -67,10 +67,6 @@ impl<'a> Window<'a> {
         unsafe { slice_assume_init_mut(slice) }
     }
 
-    pub fn as_mut_ptr(&mut self) -> *mut u8 {
-        self.buf.as_mut_ptr() as *mut u8
-    }
-
     /// # Safety
     ///
     /// `src` must point to `range.end - range.start` valid (initialized!) bytes

zlib-rs/src/inflate.rs

@@ -115,6 +115,11 @@ impl<'a> InflateStream<'a> {
 
         Some(stream)
     }
+
+    fn as_z_stream_mut(&mut self) -> &mut z_stream {
+        // safety: a valid &mut InflateStream is also a valid &mut z_stream
+        unsafe { &mut *(self as *mut _ as *mut z_stream) }
+    }
 }
 
 const MAX_BITS: u8 = 15; // maximum number of bits in a code
@@ -212,7 +217,7 @@ pub fn uncompress<'a>(
 
     let avail_out = stream.avail_out;
 
-    unsafe { end(stream) };
+    end(stream);
 
     let ret = match err {
         ReturnCode::StreamEnd => ReturnCode::Ok,
@@ -2119,33 +2124,26 @@ pub fn set_dictionary(stream: &mut InflateStream, dictionary: &[u8]) -> ReturnCo
     ReturnCode::Ok
 }
 
-/// # Safety
-///
-/// The `strm` must be either NULL or a valid mutable reference to a z_stream value where the state
-/// has been initialized with `inflateInit_` or `inflateInit2_`.
-pub unsafe extern "C" fn end(stream: &mut InflateStream) -> i32 {
+pub fn end<'a>(stream: &'a mut InflateStream) -> &'a mut z_stream {
     let mut state = State::new(&[], ReadBuf::new(&mut []));
     std::mem::swap(&mut state, stream.state);
 
     let mut window = Window::empty();
     std::mem::swap(&mut window, &mut state.window);
 
-    window.drop_in(&stream.alloc);
+    let alloc = stream.alloc;
 
-    let alloc = Allocator {
-        zalloc: stream.alloc.zalloc,
-        zfree: stream.alloc.zfree,
-        opaque: stream.alloc.opaque,
-        _marker: PhantomData,
-    };
+    // safety: window is not used again
+    unsafe { window.drop_in(&stream.alloc) };
 
-    // safety: a valid &mut InflateStream is also a valid &mut z_stream
-    let stream = unsafe { &mut *(stream as *mut _ as *mut z_stream) };
+    let stream = stream.as_z_stream_mut();
 
     let state_ptr = std::mem::replace(&mut stream.state, std::ptr::null_mut());
-    alloc.deallocate(state_ptr, 1);
 
-    ReturnCode::Ok as _
+    // safety: state_ptr is not used again
+    unsafe { alloc.deallocate(state_ptr as *mut State, 1) };
+
+    stream
 }
 
 pub fn get_header<'a>(

zlib-rs/src/inflate/window.rs

@@ -168,7 +168,7 @@ impl<'a> Window<'a> {
         })
     }
 
-    pub unsafe fn drop_in(&mut self, alloc: &Allocator) {
+    pub unsafe fn drop_in(mut self, alloc: &Allocator) {
         if !self.buf.is_empty() {
             let buf = core::mem::take(&mut self.buf);
             alloc.deallocate(buf.as_mut_ptr(), self.buf.len());