diff --git a/src/write.rs b/src/write.rs index 7e35c88fa..e453bbe3e 100644 --- a/src/write.rs +++ b/src/write.rs @@ -25,6 +25,7 @@ use core::mem::{self, offset_of, size_of}; use core::str::{Utf8Error, from_utf8}; use crc32fast::Hasher; use indexmap::IndexMap; +use std::io::Cursor; use std::io::ErrorKind; use std::io::{self, Read, Seek, Write}; use std::io::{BufReader, SeekFrom}; @@ -1768,6 +1769,66 @@ impl ZipWriter { self.add_symlink(path, target, options) } + /// Add a file whose contents were already compressed by a [`ZipFileBuilder`]. + /// + /// The compressed data is copied into the archive verbatim, so the compression work can be + /// done ahead of time, possibly on another thread. See [`ZipFileBuilder`] for an example of + /// parallelizing compression this way. The file must not have the same name as a file already + /// in the archive. + pub fn add_prepared_file(&mut self, file: PreparedZipFile) -> ZipResult<()> { + let PreparedZipFile { + file_name, + options, + crc32, + uncompressed_size, + compressed_size, + data, + } = file; + let raw_values = ZipRawValues { + crc32, + compressed_size, + uncompressed_size, + }; + + self.start_entry(&file_name, options, Some(raw_values))?; + self.writing_to_file = true; + self.writing_raw = true; + + // start_entry leaves the inner writer as a bare Storer (a compression encoder is only + // installed by start_file*), so the already-compressed bytes pass through unchanged. + // Write directly to the inner writer rather than through ZipWriter::write, whose stats + // bookkeeping (CRC-32 hashing and the large-file size check) is redundant here: the CRC + // and sizes are already known, and large_file was set from them in finish(). + let result = self + .inner + .try_inner_mut() + .and_then(|writer| writer.write_all(&data)); + self.ok_or_abort_file(result)?; + + // The local file header was written from the raw values, but for large files it contains + // a ZIP64 extra field with placeholder sizes; go back and patch it now. When the + // underlying writer doesn't support seeking, the entry uses a data descriptor instead, + // which finish_file() skips for raw entries, so write it here. + let auto_large_file = self.auto_large_file; + let result = (|| -> ZipResult<()> { + let writer = self.inner.try_inner_mut()?; + let Some((file_name_raw, file)) = self.files.last_mut() else { + return Err(ZipError::Io(io::Error::other("Cannot get last file"))); + }; + if file.is_using_data_descriptor() { + file.write_data_descriptor(writer, auto_large_file)?; + } else if file.large_file { + let file_end = writer.stream_position()?; + file.update_local_file_header(writer, file_name_raw)?; + writer.seek(SeekFrom::Start(file_end))?; + } + Ok(()) + })(); + self.ok_or_abort_file(result)?; + + self.finish_file() + } + fn finalize(&mut self) -> ZipResult { self.finish_file()?; @@ -2255,6 +2316,214 @@ impl GenericZipWriter { } } +/// Compresses the contents of a single file entry independently of any [`ZipWriter`], so that +/// the CPU-intensive part of adding files to an archive can be parallelized. +/// +/// The builder compresses the data written to it into memory as it arrives, while tracking the +/// metadata (CRC-32 and sizes) that the archive will need. Once [`ZipFileBuilder::finish`] has +/// been called, the resulting [`PreparedZipFile`] can be appended to an archive with +/// [`ZipWriter::add_prepared_file`], which copies the already-compressed bytes without +/// recompressing them. +/// +/// Builders are self-contained: they can be created, written to and finished on other threads +/// while a `ZipWriter` is in use, and the prepared files appended to the archive as they become +/// ready. +/// +/// Encryption is not supported; pass options without encryption or creation will fail. +/// +/// ``` +/// use std::io::{Cursor, Write}; +/// use zip::ZipWriter; +/// use zip::write::{SimpleFileOptions, ZipFileBuilder}; +/// +/// # fn main() -> zip::result::ZipResult<()> { +/// let options = SimpleFileOptions::default(); +/// +/// // Compress each file on its own thread... +/// let handles = ["first.txt", "second.txt"].map(|name| { +/// std::thread::spawn(move || -> zip::result::ZipResult<_> { +/// let mut builder = ZipFileBuilder::new(name, options)?; +/// builder.write_all(b"Contents compressed off-thread")?; +/// builder.finish() +/// }) +/// }); +/// +/// // ...then append them to the archive serially. +/// let mut zip = ZipWriter::new(Cursor::new(Vec::new())); +/// for handle in handles { +/// zip.add_prepared_file(handle.join().unwrap()?)?; +/// } +/// let zip_bytes = zip.finish()?.into_inner(); +/// # let mut archive = zip::ZipArchive::new(Cursor::new(zip_bytes))?; +/// # assert_eq!(archive.len(), 2); +/// # Ok(()) +/// # } +/// ``` +pub struct ZipFileBuilder { + compressor: GenericZipWriter>>, + file_name: Box<[u8]>, + options: FullFileOptions<'static, 'static>, + hasher: Hasher, + uncompressed_size: u64, +} + +impl Debug for ZipFileBuilder { + fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { + f.write_fmt(format_args!( + "ZipFileBuilder {{file_name: {:?}, uncompressed_size: {}}}", + self.file_name.escape_ascii().to_string(), + self.uncompressed_size + )) + } +} + +impl ZipFileBuilder { + /// Creates a builder that compresses a file's contents using the compression settings from + /// `options`. + /// + /// Fails if `options` requests encryption or an unsupported compression method or level. + pub fn new( + name: N, + mut options: FileOptions<'_, '_, T>, + ) -> ZipResult { + if options.encrypt_with.is_some() { + return Err(UnsupportedArchive( + "Encrypted files cannot be compressed outside of a ZipWriter", + )); + } + options.normalize(); + let owned_options = FullFileOptions { + compression_method: options.compression_method, + compression_level: options.compression_level, + external_attributes: options.external_attributes, + last_modified_time: options.last_modified_time, + permissions: options.permissions, + large_file: options.large_file, + encrypt_with: None, + extended_options: ExtendedFileOptions { + extra_fields: options + .extended_options + .extra_fields() + .cloned() + .unwrap_or_default(), + file_comment: options.extended_options.take_file_comment(), + }, + alignment: options.alignment, + #[cfg(feature = "deflate-zopfli")] + zopfli_buffer_size: options.zopfli_buffer_size, + system: options.system, + name: None, + }; + let mut compressor = + GenericZipWriter::Storer(MaybeEncrypted::Unencrypted(Cursor::new(Vec::new()))); + let make_compressor = compressor.prepare_next_writer( + owned_options.compression_method, + owned_options.compression_level, + #[cfg(feature = "deflate-zopfli")] + owned_options.zopfli_buffer_size, + )?; + compressor.switch_to(make_compressor)?; + Ok(Self { + compressor, + file_name: name.to_string().into_bytes().into_boxed_slice(), + options: owned_options, + hasher: Hasher::new(), + uncompressed_size: 0, + }) + } + + /// Like [`ZipFileBuilder::new`], but takes a [`Path`] as the file name. + /// + /// This function ensures that the '/' path separator is used and normalizes `.` and `..`. It + /// ignores any `..` or Windows drive letter that would produce a path outside the ZIP file's + /// root. + pub fn new_from_path>( + path: P, + options: FileOptions<'_, '_, T>, + ) -> ZipResult { + Self::new(path_to_string(path)?, options) + } + + /// Finishes compressing the file, and returns a [`PreparedZipFile`] that can be appended to + /// an archive with [`ZipWriter::add_prepared_file`]. + pub fn finish(mut self) -> ZipResult { + self.compressor + .switch_to(Box::new(|bare| Ok(GenericZipWriter::Storer(bare))))?; + let GenericZipWriter::Storer(MaybeEncrypted::Unencrypted(cursor)) = self.compressor else { + return Err(ZipError::Io(io::Error::other( + "Compressor was in an unexpected state", + ))); + }; + let data = cursor.into_inner().into_boxed_slice(); + let compressed_size = data.len() as u64; + let mut options = self.options; + if compressed_size >= spec::ZIP64_BYTES_THR + || self.uncompressed_size >= spec::ZIP64_BYTES_THR + { + options.large_file = true; + } + Ok(PreparedZipFile { + file_name: self.file_name, + options, + crc32: self.hasher.finalize(), + uncompressed_size: self.uncompressed_size, + compressed_size, + data, + }) + } +} + +impl Write for ZipFileBuilder { + fn write(&mut self, buf: &[u8]) -> io::Result { + if buf.is_empty() { + return Ok(0); + } + let Some(w) = self.compressor.ref_mut() else { + return Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "write(): ZipFileBuilder was already closed", + )); + }; + let count = w.write(buf)?; + self.hasher.update(&buf[..count]); + self.uncompressed_size += count as u64; + Ok(count) + } + + fn flush(&mut self) -> io::Result<()> { + match self.compressor.ref_mut() { + Some(w) => w.flush(), + None => Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "flush(): ZipFileBuilder was already closed", + )), + } + } +} + +/// A file entry whose contents have already been compressed by a [`ZipFileBuilder`], ready to be +/// appended to an archive with [`ZipWriter::add_prepared_file`]. +pub struct PreparedZipFile { + file_name: Box<[u8]>, + options: FullFileOptions<'static, 'static>, + crc32: u32, + uncompressed_size: u64, + compressed_size: u64, + data: Box<[u8]>, +} + +impl Debug for PreparedZipFile { + fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { + f.write_fmt(format_args!( + "PreparedZipFile {{file_name: {:?}, crc32: {}, uncompressed_size: {}, compressed_size: {}}}", + self.file_name.escape_ascii().to_string(), + self.crc32, + self.uncompressed_size, + self.compressed_size + )) + } +} + #[cfg(feature = "_deflate-any")] fn deflate_compression_level_range() -> std::ops::RangeInclusive { #[cfg(not(any(feature = "deflate-zopfli", feature = "deflate-flate2")))] @@ -4694,4 +4963,155 @@ mod tests { assert_eq!(std::mem::size_of::(), 4); assert_eq!(std::mem::size_of::(), 26); } + + #[test] + fn test_prepared_file_roundtrip() { + use super::ZipFileBuilder; + use std::io::Read; + + let options = SimpleFileOptions::default(); + let mut builder = ZipFileBuilder::new( + "prepared.txt", + FullFileOptions::default().with_file_comment("file comment"), + ) + .unwrap(); + builder.write_all(b"Contents of the prepared file").unwrap(); + let prepared = builder.finish().unwrap(); + let empty = ZipFileBuilder::new("empty.txt", options) + .unwrap() + .finish() + .unwrap(); + + // Interleave the prepared files with normally-written files. + let mut writer = ZipWriter::new(Cursor::new(Vec::new())); + writer.start_file("before.txt", options).unwrap(); + writer.write_all(b"before").unwrap(); + writer.add_prepared_file(prepared).unwrap(); + writer.start_file("after.txt", options).unwrap(); + writer.write_all(b"after").unwrap(); + writer.add_prepared_file(empty).unwrap(); + + let mut archive = writer.finish_into_readable().unwrap(); + let mut contents = String::new(); + for (name, expected) in [ + ("before.txt", "before"), + ("prepared.txt", "Contents of the prepared file"), + ("after.txt", "after"), + ("empty.txt", ""), + ] { + contents.clear(); + archive + .by_name(name) + .unwrap() + .read_to_string(&mut contents) + .unwrap(); + assert_eq!(contents, expected); + } + assert_eq!( + archive.by_name("prepared.txt").unwrap().comment(), + "file comment" + ); + } + + #[test] + fn test_prepared_file_duplicate_name() { + use super::ZipFileBuilder; + + let options = SimpleFileOptions::default(); + let mut writer = ZipWriter::new(Cursor::new(Vec::new())); + + let mut builder = ZipFileBuilder::new("dup.txt", options).unwrap(); + builder.write_all(b"first").unwrap(); + writer.add_prepared_file(builder.finish().unwrap()).unwrap(); + + let mut builder = ZipFileBuilder::new("dup.txt", options).unwrap(); + builder.write_all(b"second").unwrap(); + assert!(writer.add_prepared_file(builder.finish().unwrap()).is_err()); + + // The writer must still be usable after the failed add. + let mut builder = ZipFileBuilder::new("other.txt", options).unwrap(); + builder.write_all(b"third").unwrap(); + writer.add_prepared_file(builder.finish().unwrap()).unwrap(); + let archive = writer.finish_into_readable().unwrap(); + assert_eq!(archive.len(), 2); + } + + #[test] + fn test_prepared_file_stream_mode() { + use super::ZipFileBuilder; + use std::io::Read; + + let options = SimpleFileOptions::default(); + let mut builder = ZipFileBuilder::new("streamed.txt", options).unwrap(); + builder.write_all(b"streamed contents").unwrap(); + let prepared = builder.finish().unwrap(); + + let mut writer = ZipWriter::new_stream(Vec::new()); + writer.add_prepared_file(prepared).unwrap(); + let bytes = writer.finish().unwrap().into_inner(); + + let mut archive = ZipArchive::new(Cursor::new(bytes)).unwrap(); + let mut contents = String::new(); + archive + .by_name("streamed.txt") + .unwrap() + .read_to_string(&mut contents) + .unwrap(); + assert_eq!(contents, "streamed contents"); + } + + #[test] + fn test_prepared_file_rejects_encryption() { + use super::ZipFileBuilder; + + let options = SimpleFileOptions::default() + .compression_method(Stored) + .with_deprecated_encryption(b"password"); + assert!(ZipFileBuilder::new("secret.txt", options).is_err()); + } + + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn test_prepared_files_on_threads() { + use super::ZipFileBuilder; + use std::io::Read; + + let entries: Vec<(String, Vec)> = (0..8) + .map(|i| (format!("file{i}.bin"), vec![i as u8; 10_000])) + .collect(); + + let prepared: Vec<_> = std::thread::scope(|scope| { + let handles: Vec<_> = entries + .iter() + .map(|(name, data)| { + scope.spawn(move || -> ZipResult<_> { + let mut builder = ZipFileBuilder::new(name, SimpleFileOptions::default())?; + builder.write_all(data)?; + builder.finish() + }) + }) + .collect(); + handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect::>>() + }) + .unwrap(); + + let mut writer = ZipWriter::new(Cursor::new(Vec::new())); + for file in prepared { + writer.add_prepared_file(file).unwrap(); + } + let mut archive = writer.finish_into_readable().unwrap(); + let mut contents = Vec::new(); + for (name, data) in &entries { + contents.clear(); + archive + .by_name(name) + .unwrap() + .read_to_end(&mut contents) + .unwrap(); + assert_eq!(&contents, data); + } + } }