Add abort_file, and call it when validations fail

fuzz/fuzz_targets/fuzz_write.rs

@@ -69,10 +69,13 @@ fn do_operation<T>(writer: &mut zip_next::ZipWriter<T>,
     Ok(())
 }
 
-fuzz_target!(|data: Vec<FileOperation>| {
+fuzz_target!(|data: Vec<(FileOperation, bool)>| {
     let mut writer = zip_next::ZipWriter::new(Cursor::new(Vec::new()));
-    for operation in data {
+    for (operation, close_and_reopen) 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());
 });

src/write.rs

@@ -226,9 +226,7 @@ impl<W: Write + Seek> Write for ZipWriter<W> {
                         if self.stats.bytes_written > spec::ZIP64_BYTES_THR
                             && !self.files.last_mut().unwrap().large_file
                         {
-                            self.finish_file()?;
+                            let _ = self.abort_file();
-                            self.files_by_name
-                                .remove(&*self.files.pop().unwrap().file_name);
                             return Err(io::Error::new(
                                 io::ErrorKind::Other,
                                 "Large file option has not been set",
@@ -471,7 +469,8 @@ impl<W: Write + Seek> ZipWriter<W> {
             // Implicitly calling [`ZipWriter::end_extra_data`] for empty files.
             self.end_extra_data()?;
         }
-        self.inner.switch_to(CompressionMethod::Stored, None)?;
+        let make_plain_writer = self.inner.prepare_switch_to(CompressionMethod::Stored, None)?;
+        self.inner.switch_to(make_plain_writer)?;
         let writer = self.inner.get_plain();
 
         if !self.writing_raw {
@@ -495,6 +494,18 @@ impl<W: Write + Seek> ZipWriter<W> {
         Ok(())
     }
 
+    /// Removes the file currently being written from the archive.
+    pub fn abort_file(&mut self) -> ZipResult<()> {
+        self.files_by_name
+            .remove(&*self.files.pop().unwrap().file_name);
+        let make_plain_writer
+            = self.inner.prepare_switch_to(CompressionMethod::Stored, None)?;
+        self.inner.switch_to(make_plain_writer)?;
+        self.writing_to_file = false;
+        self.writing_raw = false;
+        Ok(())
+    }
+
     /// Create a file in the archive and start writing its' contents. The file must not have the
     /// same name as a file already in the archive.
     ///
@@ -504,9 +515,10 @@ impl<W: Write + Seek> ZipWriter<W> {
         S: Into<String>,
     {
         Self::normalize_options(&mut options);
+        let make_new_self = self.inner
+            .prepare_switch_to(options.compression_method, options.compression_level)?;
         self.start_entry(name, options, None)?;
-        self.inner
+        self.inner.switch_to(make_new_self)?;
-            .switch_to(options.compression_method, options.compression_level)?;
         self.writing_to_file = true;
         Ok(())
     }
@@ -668,7 +680,19 @@ impl<W: Write + Seek> ZipWriter<W> {
         }
         let file = self.files.last_mut().unwrap();
 
-        validate_extra_data(file)?;
+        if let Err(e) = validate_extra_data(file) {
+            let _ = self.abort_file();
+            return Err(e);
+        }
+
+        let make_compressing_writer = match self.inner
+            .prepare_switch_to(file.compression_method, file.compression_level) {
+            Ok(writer) => writer,
+            Err(e) => {
+                let _ = self.abort_file();
+                return Err(e);
+            }
+        };
 
         let mut data_start_result = file.data_start.load();
 
@@ -692,11 +716,8 @@ impl<W: Write + Seek> ZipWriter<W> {
             writer.seek(SeekFrom::Start(file.header_start + 28))?;
             writer.write_u16::<LittleEndian>(extra_field_length)?;
             writer.seek(SeekFrom::Start(header_end))?;
-
-            self.inner
-                .switch_to(file.compression_method, file.compression_level)?;
         }
-
+        self.inner.switch_to(make_compressing_writer)?;
         self.writing_to_extra_field = false;
         self.writing_to_central_extra_field_only = false;
         Ok(data_start_result)
@@ -957,33 +978,28 @@ impl<W: Write + Seek> Drop for ZipWriter<W> {
 }
 
 impl<W: Write + Seek> GenericZipWriter<W> {
-    fn switch_to(
+    fn prepare_switch_to(
         &mut self,
         compression: CompressionMethod,
         compression_level: Option<i32>,
-    ) -> ZipResult<()> {
+    ) -> ZipResult<Box<dyn FnOnce(W) -> GenericZipWriter<W>>> {
-        match self.current_compression() {
-            Some(method) if method == compression => return Ok(()),
-            _ => {}
-        }
-
         if let Closed = self {
             return Err(
                 io::Error::new(io::ErrorKind::BrokenPipe, "ZipWriter was already closed").into(),
             );
         }
 
-        let make_new_self: Box<dyn FnOnce(W) -> GenericZipWriter<W>> = {
+        {
             #[allow(deprecated)]
             match compression {
                 CompressionMethod::Stored => {
                     if compression_level.is_some() {
-                        return Err(ZipError::UnsupportedArchive(
+                        Err(ZipError::UnsupportedArchive(
                             "Unsupported compression level",
-                        ));
+                        ))
+                    } else {
+                        Ok(Box::new(|bare| Storer(bare)))
                     }
-
-                    Box::new(|bare| Storer(bare))
                 }
                 #[cfg(any(
                 feature = "deflate",
@@ -998,12 +1014,12 @@ impl<W: Write + Seek> GenericZipWriter<W> {
                         .ok_or(ZipError::UnsupportedArchive(
                             "Unsupported compression level",
                         ))? as u32;
-                    Box::new(move |bare| {
+                    Ok(Box::new(move |bare| {
                         GenericZipWriter::Deflater(DeflateEncoder::new(
                             bare,
                             flate2::Compression::new(level),
                         ))
-                    })
+                    }))
                 }
                 #[cfg(feature = "bzip2")]
                 CompressionMethod::Bzip2 => {
@@ -1014,15 +1030,15 @@ impl<W: Write + Seek> GenericZipWriter<W> {
                         .ok_or(ZipError::UnsupportedArchive(
                             "Unsupported compression level",
                         ))? as u32;
-                    Box::new(move |bare| {
+                    Ok(Box::new(move |bare| {
                         GenericZipWriter::Bzip2(BzEncoder::new(
                             bare,
                             bzip2::Compression::new(level),
                         ))
-                    })
+                    }))
                 }
                 CompressionMethod::AES => {
-                    return Err(ZipError::UnsupportedArchive(
+                    Err(ZipError::UnsupportedArchive(
                         "AES compression is not supported for writing",
                     ))
                 }
@@ -1035,16 +1051,19 @@ impl<W: Write + Seek> GenericZipWriter<W> {
                         .ok_or(ZipError::UnsupportedArchive(
                             "Unsupported compression level",
                         ))?;
-                    Box::new(move |bare| {
+                    Ok(Box::new(move |bare| {
                         GenericZipWriter::Zstd(ZstdEncoder::new(bare, level).unwrap())
-                    })
+                    }))
                 }
                 CompressionMethod::Unsupported(..) => {
-                    return Err(ZipError::UnsupportedArchive("Unsupported compression"))
+                    Err(ZipError::UnsupportedArchive("Unsupported compression"))
+                }
+            }
         }
     }
-        };
 
+    fn switch_to(&mut self, make_new_self: Box<dyn FnOnce(W) -> GenericZipWriter<W>>)
+            -> ZipResult<()>{
         let bare = match mem::replace(self, Closed) {
             Storer(w) => w,
             #[cfg(any(
@@ -1061,8 +1080,7 @@ impl<W: Write + Seek> GenericZipWriter<W> {
                 return Err(io::Error::new(
                     io::ErrorKind::BrokenPipe,
                     "ZipWriter was already closed",
-                )
+                ).into());
-                .into())
             }
         };
         *self = (make_new_self)(bare);
@@ -1097,23 +1115,6 @@ impl<W: Write + Seek> GenericZipWriter<W> {
         }
     }
 
-    fn current_compression(&self) -> Option<CompressionMethod> {
-        match *self {
-            Storer(..) => Some(CompressionMethod::Stored),
-            #[cfg(any(
-                feature = "deflate",
-                feature = "deflate-miniz",
-                feature = "deflate-zlib"
-            ))]
-            GenericZipWriter::Deflater(..) => Some(CompressionMethod::Deflated),
-            #[cfg(feature = "bzip2")]
-            GenericZipWriter::Bzip2(..) => Some(CompressionMethod::Bzip2),
-            #[cfg(feature = "zstd")]
-            GenericZipWriter::Zstd(..) => Some(CompressionMethod::Zstd),
-            Closed => None,
-        }
-    }
-
     fn unwrap(self) -> W {
         match self {
             Storer(w) => w,