From 82e45fa8901824344cf636a37c03b2595478a4d1 Mon Sep 17 00:00:00 2001
From: Lukasz Anforowicz <lukasza@chromium.org>
Date: Tue, 29 Aug 2023 20:19:24 +0000
Subject: [PATCH 1/2] Refactoring: Dedupe code into
 `write_to_spare_capacity_of_vec` helper.

---
 src/mem.rs | 64 +++++++++++++++++++++++++++++++++---------------------
 1 file changed, 39 insertions(+), 25 deletions(-)

diff --git a/src/mem.rs b/src/mem.rs
index 6313c220..6d311a6f 100644
--- a/src/mem.rs
+++ b/src/mem.rs
@@ -1,7 +1,6 @@
 use std::error::Error;
 use std::fmt;
 use std::io;
-use std::slice;
 
 use crate::ffi::{self, Backend, Deflate, DeflateBackend, ErrorMessage, Inflate, InflateBackend};
 use crate::Compression;
@@ -342,19 +341,12 @@ impl Compress {
         output: &mut Vec<u8>,
         flush: FlushCompress,
     ) -> Result<Status, CompressError> {
-        let cap = output.capacity();
-        let len = output.len();
-
-        unsafe {
+        write_to_spare_capacity_of_vec(output, |out| {
             let before = self.total_out();
-            let ret = {
-                let ptr = output.as_mut_ptr().add(len);
-                let out = slice::from_raw_parts_mut(ptr, cap - len);
-                self.compress(input, out, flush)
-            };
-            output.set_len((self.total_out() - before) as usize + len);
-            ret
-        }
+            let ret = self.compress(input, out, flush);
+            let bytes_written = self.total_out() - before;
+            (bytes_written as usize, ret)
+        })
     }
 }
 
@@ -473,19 +465,12 @@ impl Decompress {
         output: &mut Vec<u8>,
         flush: FlushDecompress,
     ) -> Result<Status, DecompressError> {
-        let cap = output.capacity();
-        let len = output.len();
-
-        unsafe {
+        write_to_spare_capacity_of_vec(output, |out| {
             let before = self.total_out();
-            let ret = {
-                let ptr = output.as_mut_ptr().add(len);
-                let out = slice::from_raw_parts_mut(ptr, cap - len);
-                self.decompress(input, out, flush)
-            };
-            output.set_len((self.total_out() - before) as usize + len);
-            ret
-        }
+            let ret = self.decompress(input, out, flush);
+            let bytes_written = self.total_out() - before;
+            (bytes_written as usize, ret)
+        })
     }
 
     /// Specifies the decompression dictionary to use.
@@ -574,6 +559,35 @@ impl fmt::Display for CompressError {
     }
 }
 
+/// Allows `writer` to write data into the spare capacity of the `output` vector.
+/// This will not reallocate the vector provided or attempt to grow it, so space
+/// for the `output` must be reserved by the caller before calling this
+/// function.
+///
+/// `writer` needs to return the number of bytes written (and can also return
+/// another arbitrary return value).
+fn write_to_spare_capacity_of_vec<T>(
+    output: &mut Vec<u8>,
+    writer: impl FnOnce(&mut [u8]) -> (usize, T),
+) -> T {
+    let cap = output.capacity();
+    let len = output.len();
+
+    // FIXME: This is unsound - see https://github.com/rust-lang/flate2-rs/issues/220
+    // (The code below reimplements `Vec::spare_capacity_mut`, but returns `&mut [u8]`
+    // instead of `&mut [MaybeUninit<u8>]`.)
+    unsafe {
+        let (bytes_written, ret) = {
+            let ptr = output.as_mut_ptr().add(len);
+            let out = slice::from_raw_parts_mut(ptr, cap - len);
+            writer(out)
+        };
+        let new_len = core::cmp::min(len + bytes_written, cap);  // Sanitizes `bytes_written`.
+        output.set_len(new_len);
+        ret
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::io::Write;

From 69972b8262fbe03c28450c4c18961b41ad92af6a Mon Sep 17 00:00:00 2001
From: Lukasz Anforowicz <lukasza@chromium.org>
Date: Tue, 29 Aug 2023 20:52:45 +0000
Subject: [PATCH 2/2] Fix soundness of `write_to_spare_capacity_of_vec`.

Fixes https://github.com/rust-lang/flate2-rs/issues/220
---
 src/mem.rs | 20 +++++++-------------
 1 file changed, 7 insertions(+), 13 deletions(-)

diff --git a/src/mem.rs b/src/mem.rs
index 6d311a6f..d4a50917 100644
--- a/src/mem.rs
+++ b/src/mem.rs
@@ -573,19 +573,13 @@ fn write_to_spare_capacity_of_vec<T>(
     let cap = output.capacity();
     let len = output.len();
 
-    // FIXME: This is unsound - see https://github.com/rust-lang/flate2-rs/issues/220
-    // (The code below reimplements `Vec::spare_capacity_mut`, but returns `&mut [u8]`
-    // instead of `&mut [MaybeUninit<u8>]`.)
-    unsafe {
-        let (bytes_written, ret) = {
-            let ptr = output.as_mut_ptr().add(len);
-            let out = slice::from_raw_parts_mut(ptr, cap - len);
-            writer(out)
-        };
-        let new_len = core::cmp::min(len + bytes_written, cap);  // Sanitizes `bytes_written`.
-        output.set_len(new_len);
-        ret
-    }
+    output.resize(output.capacity(), 0);
+    let (bytes_written, ret) = writer(&mut output[len..]);
+
+    let new_len = core::cmp::min(len + bytes_written, cap); // Sanitizes `bytes_written`.
+    output.resize(new_len, 0 /* unused */);
+
+    ret
 }
 
 #[cfg(test)]