async_zip/base/write/
mod.rs

1// Copyright (c) 2021-2022 Harry [Majored] [hello@majored.pw]
2// MIT License (https://github.com/Majored/rs-async-zip/blob/main/LICENSE)
3
4//! A module which supports writing ZIP files.
5//!
6//! # Example
7//! ### Whole data (u8 slice)
8//! ```no_run
9//! # #[cfg(feature = "deflate")]
10//! # {
11//! # use async_zip::{Compression, ZipEntryBuilder, base::write::ZipFileWriter};
12//! # use async_zip::error::ZipError;
13//! #
14//! # async fn run() -> Result<(), ZipError> {
15//! let mut writer = ZipFileWriter::new(Vec::<u8>::new());
16//!
17//! let data = b"This is an example file.";
18//! let opts = ZipEntryBuilder::new(String::from("foo.txt").into(), Compression::Deflate);
19//!
20//! writer.write_entry_whole(opts, data).await?;
21//! writer.close().await?;
22//! #   Ok(())
23//! # }
24//! # }
25//! ```
26//! ### Stream data (unknown size & data)
27//! ```no_run
28//! # #[cfg(feature = "deflate")]
29//! # {
30//! # use async_zip::{Compression, ZipEntryBuilder, base::write::ZipFileWriter};
31//! # use std::io::Cursor;
32//! # use async_zip::error::ZipError;
33//! # use futures_lite::io::AsyncWriteExt;
34//! # use tokio_util::compat::TokioAsyncWriteCompatExt;
35//! #
36//! # async fn run() -> Result<(), ZipError> {
37//! let mut writer = ZipFileWriter::new(Vec::<u8>::new());
38//!
39//! let data = b"This is an example file.";
40//! let opts = ZipEntryBuilder::new(String::from("bar.txt").into(), Compression::Deflate);
41//!
42//! let mut entry_writer = writer.write_entry_stream(opts).await?;
43//! entry_writer.write_all(data).await.unwrap();
44//!
45//! entry_writer.close().await?;
46//! writer.close().await?;
47//! #   Ok(())
48//! # }
49//! # }
50//! ```
51
52pub(crate) mod compressed_writer;
53pub(crate) mod entry_seekable;
54pub(crate) mod entry_stream;
55pub(crate) mod entry_whole;
56pub(crate) mod io;
57
58pub use entry_seekable::EntrySeekableWriter;
59pub use entry_stream::EntryStreamWriter;
60
61#[cfg(feature = "tokio")]
62use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt};
63
64use crate::entry::ZipEntry;
65use crate::error::Result;
66use crate::spec::extra_field::ExtraFieldAsBytes;
67use crate::spec::header::{
68    CentralDirectoryRecord, EndOfCentralDirectoryHeader, ExtraField, InfoZipUnicodeCommentExtraField,
69    InfoZipUnicodePathExtraField, Zip64EndOfCentralDirectoryLocator, Zip64EndOfCentralDirectoryRecord,
70};
71
72#[cfg(feature = "tokio")]
73use crate::tokio::write::ZipFileWriter as TokioZipFileWriter;
74
75use entry_whole::EntryWholeWriter;
76use io::offset::AsyncOffsetWriter;
77
78use crate::spec::consts::{NON_ZIP64_MAX_NUM_FILES, NON_ZIP64_MAX_SIZE};
79use futures_lite::io::{AsyncSeek, AsyncWrite, AsyncWriteExt};
80
81pub(crate) struct CentralDirectoryEntry {
82    pub header: CentralDirectoryRecord,
83    pub entry: ZipEntry,
84}
85
86/// A ZIP file writer which acts over AsyncWrite implementers.
87///
88/// # Note
89/// - [`ZipFileWriter::close()`] must be called before a stream writer goes out of scope.
90pub struct ZipFileWriter<W> {
91    pub(crate) writer: AsyncOffsetWriter<W>,
92    pub(crate) cd_entries: Vec<CentralDirectoryEntry>,
93    /// If true, will error if a Zip64 struct must be written.
94    force_no_zip64: bool,
95    /// Whether to write Zip64 end of directory structs.
96    pub(crate) is_zip64: bool,
97    comment_opt: Option<String>,
98}
99
100impl<W: AsyncWrite + Unpin> ZipFileWriter<W> {
101    /// Construct a new ZIP file writer from a mutable reference to a writer.
102    pub fn new(writer: W) -> Self {
103        Self {
104            writer: AsyncOffsetWriter::new(writer),
105            cd_entries: Vec::new(),
106            comment_opt: None,
107            is_zip64: false,
108            force_no_zip64: false,
109        }
110    }
111
112    /// Force the ZIP writer to operate in non-ZIP64 mode.
113    /// If any files would need ZIP64, an error will be raised.
114    pub fn force_no_zip64(mut self) -> Self {
115        self.force_no_zip64 = true;
116        self
117    }
118
119    /// Force the ZIP writer to emit Zip64 structs at the end of the archive.
120    /// Zip64 extended fields will only be written if needed.
121    pub fn force_zip64(mut self) -> Self {
122        self.is_zip64 = true;
123        self
124    }
125
126    /// Write a new ZIP entry of known size and data.
127    pub async fn write_entry_whole<E: Into<ZipEntry>>(&mut self, entry: E, data: &[u8]) -> Result<()> {
128        EntryWholeWriter::from_raw(self, entry.into(), data).write().await
129    }
130
131    /// Write an entry of unknown size and data via streaming (ie. using a data descriptor).
132    /// The generated Local File Header will be invalid, with no compressed size, uncompressed size,
133    /// and a null CRC. This might cause problems with the destination reader.
134    pub async fn write_entry_stream<E: Into<ZipEntry>>(&mut self, entry: E) -> Result<EntryStreamWriter<'_, W>> {
135        EntryStreamWriter::from_raw(self, entry.into()).await
136    }
137
138    /// Write an entry of unknown size and data via streaming to a seekable output.
139    ///
140    /// This avoids data descriptors by seeking back to patch the local file header after
141    /// the entry is written.
142    pub async fn write_entry_seekable<E: Into<ZipEntry>>(&mut self, entry: E) -> Result<EntrySeekableWriter<'_, W>>
143    where
144        W: AsyncSeek,
145    {
146        EntrySeekableWriter::from_raw(self, entry.into()).await
147    }
148
149    /// Set the ZIP file comment.
150    pub fn comment(&mut self, comment: String) {
151        self.comment_opt = Some(comment);
152    }
153
154    /// Returns a mutable reference to the inner writer.
155    ///
156    /// Care should be taken when using this inner writer as doing so may invalidate internal state of this writer.
157    pub fn inner_mut(&mut self) -> &mut W {
158        self.writer.inner_mut()
159    }
160
161    /// Consumes this ZIP writer and completes all closing tasks.
162    ///
163    /// This includes:
164    /// - Writing all central directory headers.
165    /// - Writing the end of central directory header.
166    /// - Writing the file comment.
167    ///
168    /// Failure to call this function before going out of scope would result in a corrupted ZIP file.
169    pub async fn close(mut self) -> Result<W> {
170        let file_comment_length = self
171            .comment_opt
172            .as_ref()
173            .map(|comment| comment.len().try_into())
174            .transpose()
175            .map_err(|_| crate::error::ZipError::CommentTooLarge)?
176            .unwrap_or_default();
177        let cd_offset = self.writer.offset();
178
179        for entry in &self.cd_entries {
180            let filename_basic =
181                entry.entry.filename().alternative().unwrap_or_else(|| entry.entry.filename().as_bytes());
182            let comment_basic = entry.entry.comment().alternative().unwrap_or_else(|| entry.entry.comment().as_bytes());
183
184            self.writer.write_all(&crate::spec::consts::CDH_SIGNATURE.to_le_bytes()).await?;
185            self.writer.write_all(&entry.header.as_slice()).await?;
186            self.writer.write_all(filename_basic).await?;
187            self.writer.write_all(&entry.entry.extra_fields().as_bytes()).await?;
188            self.writer.write_all(comment_basic).await?;
189        }
190
191        let central_directory_size = self.writer.offset() - cd_offset;
192        let central_directory_size_u32 =
193            central_directory_size_field(central_directory_size, self.force_no_zip64, &mut self.is_zip64)?;
194        let num_entries_in_directory = self.cd_entries.len() as u64;
195        let num_entries_in_directory_u16 = if num_entries_in_directory > NON_ZIP64_MAX_NUM_FILES as u64 {
196            NON_ZIP64_MAX_NUM_FILES
197        } else {
198            num_entries_in_directory as u16
199        };
200        let cd_offset_u32 = if cd_offset > NON_ZIP64_MAX_SIZE as u64 {
201            if self.force_no_zip64 {
202                return Err(crate::error::ZipError::Zip64Needed(crate::error::Zip64ErrorCase::LargeFile));
203            } else {
204                self.is_zip64 = true;
205            }
206            NON_ZIP64_MAX_SIZE
207        } else {
208            cd_offset as u32
209        };
210
211        // Add the zip64 EOCDR and EOCDL if we are in zip64 mode.
212        if self.is_zip64 {
213            let eocdr_offset = self.writer.offset();
214
215            let eocdr = Zip64EndOfCentralDirectoryRecord {
216                size_of_zip64_end_of_cd_record: 44,
217                version_made_by: crate::spec::version::as_made_by(),
218                version_needed_to_extract: 46,
219                disk_number: 0,
220                disk_number_start_of_cd: 0,
221                num_entries_in_directory_on_disk: num_entries_in_directory,
222                num_entries_in_directory,
223                directory_size: central_directory_size,
224                offset_of_start_of_directory: cd_offset,
225            };
226            self.writer.write_all(&crate::spec::consts::ZIP64_EOCDR_SIGNATURE.to_le_bytes()).await?;
227            self.writer.write_all(&eocdr.as_bytes()).await?;
228
229            let eocdl = Zip64EndOfCentralDirectoryLocator {
230                number_of_disk_with_start_of_zip64_end_of_central_directory: 0,
231                relative_offset: eocdr_offset,
232                total_number_of_disks: 1,
233            };
234            self.writer.write_all(&crate::spec::consts::ZIP64_EOCDL_SIGNATURE.to_le_bytes()).await?;
235            self.writer.write_all(&eocdl.as_bytes()).await?;
236        }
237
238        let header = EndOfCentralDirectoryHeader {
239            disk_num: 0,
240            start_cent_dir_disk: 0,
241            num_of_entries_disk: num_entries_in_directory_u16,
242            num_of_entries: num_entries_in_directory_u16,
243            size_cent_dir: central_directory_size_u32,
244            cent_dir_offset: cd_offset_u32,
245            file_comm_length: file_comment_length,
246        };
247
248        self.writer.write_all(&crate::spec::consts::EOCDR_SIGNATURE.to_le_bytes()).await?;
249        self.writer.write_all(&header.as_slice()).await?;
250        if let Some(comment) = self.comment_opt {
251            self.writer.write_all(comment.as_bytes()).await?;
252        }
253
254        Ok(self.writer.into_inner())
255    }
256}
257
258pub(crate) fn central_directory_size_field(
259    central_directory_size: u64,
260    force_no_zip64: bool,
261    is_zip64: &mut bool,
262) -> Result<u32> {
263    if central_directory_size > NON_ZIP64_MAX_SIZE as u64 {
264        if force_no_zip64 {
265            return Err(crate::error::ZipError::Zip64Needed(crate::error::Zip64ErrorCase::LargeFile));
266        }
267        *is_zip64 = true;
268        Ok(NON_ZIP64_MAX_SIZE)
269    } else {
270        Ok(central_directory_size as u32)
271    }
272}
273
274#[cfg(feature = "tokio")]
275impl<W> ZipFileWriter<Compat<W>>
276where
277    W: tokio::io::AsyncWrite + Unpin,
278{
279    /// Construct a new ZIP file writer from a mutable reference to a writer.
280    pub fn with_tokio(writer: W) -> TokioZipFileWriter<W> {
281        Self {
282            writer: AsyncOffsetWriter::new(writer.compat_write()),
283            cd_entries: Vec::new(),
284            comment_opt: None,
285            is_zip64: false,
286            force_no_zip64: false,
287        }
288    }
289}
290
291pub(crate) fn get_or_put_info_zip_unicode_path_extra_field_mut(
292    extra_fields: &mut Vec<ExtraField>,
293) -> &mut InfoZipUnicodePathExtraField {
294    if !extra_fields.iter().any(|field| matches!(field, ExtraField::InfoZipUnicodePath(_))) {
295        extra_fields
296            .push(ExtraField::InfoZipUnicodePath(InfoZipUnicodePathExtraField::V1 { crc32: 0, unicode: vec![] }));
297    }
298
299    for field in extra_fields.iter_mut() {
300        if let ExtraField::InfoZipUnicodePath(extra_field) = field {
301            return extra_field;
302        }
303    }
304
305    panic!("InfoZipUnicodePathExtraField not found after insertion")
306}
307
308pub(crate) fn get_or_put_info_zip_unicode_comment_extra_field_mut(
309    extra_fields: &mut Vec<ExtraField>,
310) -> &mut InfoZipUnicodeCommentExtraField {
311    if !extra_fields.iter().any(|field| matches!(field, ExtraField::InfoZipUnicodeComment(_))) {
312        extra_fields
313            .push(ExtraField::InfoZipUnicodeComment(InfoZipUnicodeCommentExtraField::V1 { crc32: 0, unicode: vec![] }));
314    }
315
316    for field in extra_fields.iter_mut() {
317        if let ExtraField::InfoZipUnicodeComment(extra_field) = field {
318            return extra_field;
319        }
320    }
321
322    panic!("InfoZipUnicodeCommentExtraField not found after insertion")
323}
async_zip/base/write/mod.rs

async_zip/base/write/
mod.rs
