From 10da026ff68f5b0eeff54fc0d1f900368a301034 Mon Sep 17 00:00:00 2001
From: zhanghar <zhanghar@iu.edu>
Date: Wed, 23 Dec 2020 14:37:09 -0500
Subject: [PATCH] feat: support append to an existing archive

---
 src/read.rs         |  5 +++--
 src/write.rs        | 42 +++++++++++++++++++++++++++++++++++++++++-
 tests/end_to_end.rs | 19 +++++++++++++++++++
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/src/read.rs b/src/read.rs
index 3aac00f6..0ce37c0a 100644
--- a/src/read.rs
+++ b/src/read.rs
@@ -209,7 +209,7 @@ fn make_reader<'a>(
 impl<R: Read + io::Seek> ZipArchive<R> {
     /// Get the directory start offset and number of files. This is done in a
     /// separate function to ease the control flow design.
-    fn get_directory_counts(
+    pub fn get_directory_counts(
         reader: &mut R,
         footer: &spec::CentralDirectoryEnd,
         cde_start_pos: u64,
@@ -514,7 +514,8 @@ fn unsupported_zip_error<T>(detail: &'static str) -> ZipResult<T> {
     Err(ZipError::UnsupportedArchive(detail))
 }
 
-fn central_header_to_zip_file<R: Read + io::Seek>(
+/// Parse a central directory entry to collect the information for the file.
+pub fn central_header_to_zip_file<R: Read + io::Seek>(
     reader: &mut R,
     archive_offset: u64,
 ) -> ZipResult<ZipFileData> {
diff --git a/src/write.rs b/src/write.rs
index bc688172..2994b40a 100644
--- a/src/write.rs
+++ b/src/write.rs
@@ -1,7 +1,7 @@
 //! Types for creating ZIP archives
 
 use crate::compression::CompressionMethod;
-use crate::read::ZipFile;
+use crate::read::{central_header_to_zip_file, ZipArchive, ZipFile};
 use crate::result::{ZipError, ZipResult};
 use crate::spec;
 use crate::types::{DateTime, System, ZipFileData, DEFAULT_VERSION};
@@ -194,6 +194,46 @@ impl ZipWriterStats {
     }
 }
 
+impl<A: Read + Write + io::Seek> ZipWriter<A> {
+    /// Initializes the archive from an existing ZIP archive, making it ready for append.
+    pub fn new_append(mut readwriter: A) -> ZipResult<ZipWriter<A>> {
+        let (footer, cde_start_pos) = spec::CentralDirectoryEnd::find_and_parse(&mut readwriter)?;
+
+        if footer.disk_number != footer.disk_with_central_directory {
+            return Err(ZipError::UnsupportedArchive(
+                "Support for multi-disk files is not implemented",
+            ));
+        }
+
+        let (archive_offset, directory_start, number_of_files) =
+            ZipArchive::get_directory_counts(&mut readwriter, &footer, cde_start_pos)?;
+
+        let mut files = Vec::new();
+
+        if let Err(_) = readwriter.seek(io::SeekFrom::Start(directory_start)) {
+            return Err(ZipError::InvalidArchive(
+                "Could not seek to start of central directory",
+            ));
+        }
+
+        for _ in 0..number_of_files {
+            let file = central_header_to_zip_file(&mut readwriter, archive_offset)?;
+            files.push(file);
+        }
+
+        let _ = readwriter.seek(io::SeekFrom::Start(directory_start)); // seek directory_start to overwrite it
+
+        Ok(ZipWriter {
+            inner: GenericZipWriter::Storer(readwriter),
+            files,
+            stats: Default::default(),
+            writing_to_file: false,
+            comment: String::new(),
+            writing_raw: true, // avoid recomputing the last file's header
+        })
+    }
+}
+
 impl<W: Write + io::Seek> ZipWriter<W> {
     /// Initializes the archive.
     ///
diff --git a/tests/end_to_end.rs b/tests/end_to_end.rs
index b826f548..c658a6f5 100644
--- a/tests/end_to_end.rs
+++ b/tests/end_to_end.rs
@@ -46,6 +46,25 @@ fn copy() {
     check_zip_file_contents(&mut tgt_archive, COPY_ENTRY_NAME);
 }
 
+// This test asserts that after appending to a `ZipWriter`, then reading its contents back out,
+// both the prior data and the appended data will be exactly the same as their originals.
+#[test]
+fn append() {
+    let mut file = &mut Cursor::new(Vec::new());
+    write_to_zip(file).expect("file written");
+
+    {
+        let mut zip = zip::ZipWriter::new_append(&mut file).unwrap();
+        zip.start_file(COPY_ENTRY_NAME, Default::default()).unwrap();
+        zip.write_all(LOREM_IPSUM).unwrap();
+        zip.finish().unwrap();
+    }
+
+    let mut zip = zip::ZipArchive::new(&mut file).unwrap();
+    check_zip_file_contents(&mut zip, ENTRY_NAME);
+    check_zip_file_contents(&mut zip, COPY_ENTRY_NAME);
+}
+
 fn write_to_zip(file: &mut Cursor<Vec<u8>>) -> zip::result::ZipResult<()> {
     let mut zip = zip::ZipWriter::new(file);