From 3af7f187d737b8add5fa57773d28a9172b07b35b Mon Sep 17 00:00:00 2001
From: Chris Hennick <hennickc@amazon.com>
Date: Wed, 10 May 2023 14:54:21 -0700
Subject: [PATCH] Strengthen fuzz_write: can now close and reopen before
 copying

---
 fuzz/fuzz_targets/fuzz_write.rs | 49 +++++++++++++++++++++------------
 1 file changed, 31 insertions(+), 18 deletions(-)

diff --git a/fuzz/fuzz_targets/fuzz_write.rs b/fuzz/fuzz_targets/fuzz_write.rs
index eb700cf9..ff495ffa 100644
--- a/fuzz/fuzz_targets/fuzz_write.rs
+++ b/fuzz/fuzz_targets/fuzz_write.rs
@@ -1,5 +1,6 @@
 #![no_main]
 
+use std::cell::RefCell;
 use libfuzzer_sys::fuzz_target;
 use arbitrary::Arbitrary;
 use std::io::{Cursor, Read, Seek, Write};
@@ -14,15 +15,18 @@ pub struct File {
 pub enum FileOperation {
     Write {
         file: File,
-        options: zip_next::write::FileOptions
+        options: zip_next::write::FileOptions,
+        reopen: bool,
     },
     ShallowCopy {
         base: Box<FileOperation>,
-        new_name: String
+        new_name: String,
+        reopen: bool,
     },
     DeepCopy {
         base: Box<FileOperation>,
-        new_name: String
+        new_name: String,
+        reopen: bool,
     }
 }
 
@@ -34,40 +38,49 @@ impl FileOperation {
             FileOperation::DeepCopy {new_name, ..} => new_name
         }.to_owned()
     }
+
+    pub fn should_reopen(&self) -> bool {
+        match self {
+            FileOperation::Write {reopen, ..} => *reopen,
+            FileOperation::ShallowCopy {reopen, ..} => *reopen,
+            FileOperation::DeepCopy {reopen, ..} => *reopen
+        }
+    }
 }
 
-fn do_operation<T>(writer: &mut zip_next::ZipWriter<T>,
+fn do_operation<T>(writer: &mut RefCell<zip_next::ZipWriter<T>>,
                    operation: &FileOperation) -> Result<(), Box<dyn std::error::Error>>
                    where T: Read + Write + Seek {
     match operation {
-        FileOperation::Write {file, mut options} => {
+        FileOperation::Write {file, mut options, ..} => {
             if file.contents.iter().map(Vec::len).sum::<usize>() >= u32::MAX as usize {
                 options = options.large_file(true);
             }
-            writer.start_file(file.name.to_owned(), options)?;
+            writer.borrow_mut().start_file(file.name.to_owned(), options)?;
             for chunk in &file.contents {
-                writer.write_all(chunk.as_slice())?;
+                writer.borrow_mut().write_all(chunk.as_slice())?;
             }
         }
-        FileOperation::ShallowCopy {base, new_name} => {
+        FileOperation::ShallowCopy {base, new_name, .. } => {
             do_operation(writer, base)?;
-            writer.shallow_copy_file(&base.get_name(), new_name)?;
+            writer.borrow_mut().shallow_copy_file(&base.get_name(), new_name)?;
         }
-        FileOperation::DeepCopy {base, new_name} => {
+        FileOperation::DeepCopy {base, new_name, .. } => {
             do_operation(writer, base)?;
-            writer.deep_copy_file(&base.get_name(), new_name)?;
+            writer.borrow_mut().deep_copy_file(&base.get_name(), new_name)?;
         }
     }
+    if operation.should_reopen() {
+        let new_writer = zip_next::ZipWriter::new_append(writer.borrow_mut().finish().unwrap()).unwrap();
+        *writer = new_writer.into();
+    }
     Ok(())
 }
 
-fuzz_target!(|data: Vec<(FileOperation, bool)>| {
-    let mut writer = zip_next::ZipWriter::new(Cursor::new(Vec::new()));
-    for (operation, close_and_reopen) in data {
+fuzz_target!(|data: Vec<FileOperation>| {
+    let mut writer = RefCell::new(zip_next::ZipWriter::new(Cursor::new(Vec::new())));
+    for operation in data {
         let _ = do_operation(&mut writer, &operation);
-        if close_and_reopen {
-            writer = zip_next::ZipWriter::new_append(writer.finish().unwrap()).unwrap();
-        }
     }
-    let _ = zip_next::ZipArchive::new(writer.finish().unwrap());
+    let _ = zip_next::ZipArchive::new(writer.borrow_mut().finish().unwrap());
 });
\ No newline at end of file