Skip to main content

async_zip/base/write/
entry_seekable.rs

1// MIT License (https://github.com/Majored/rs-async-zip/blob/main/LICENSE)
2
3use crate::base::read::get_zip64_extra_field_mut;
4use crate::base::write::compressed_writer::CompressedAsyncWriter;
5use crate::base::write::get_or_put_info_zip_unicode_comment_extra_field_mut;
6use crate::base::write::get_or_put_info_zip_unicode_path_extra_field_mut;
7use crate::base::write::io::offset::AsyncOffsetWriter;
8use crate::base::write::{CentralDirectoryEntry, ZipFileWriter};
9use crate::entry::ZipEntry;
10use crate::error::{Result, Zip64ErrorCase, ZipError};
11use crate::spec::consts::{NON_ZIP64_MAX_NUM_FILES, NON_ZIP64_MAX_SIZE};
12use crate::spec::extra_field::ExtraFieldAsBytes;
13use crate::spec::header::{
14    CentralDirectoryRecord, ExtraField, GeneralPurposeFlag, InfoZipUnicodeCommentExtraField,
15    InfoZipUnicodePathExtraField, LocalFileHeader, Zip64ExtendedInformationExtraField,
16};
17use crate::StringEncoding;
18
19use crc32fast::Hasher;
20use futures_lite::io::{AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, SeekFrom};
21use std::io::Error;
22use std::pin::Pin;
23use std::task::{Context, Poll};
24
25const ZIP64_VERSION_NEEDED: u16 = 45;
26
27/// An entry writer which streams data to a seekable ZIP output.
28///
29/// Unlike [`EntryStreamWriter`](crate::base::write::EntryStreamWriter), this writer doesn't use
30/// data descriptors. Instead, it writes a placeholder local file header, streams the entry data,
31/// then seeks back and patches the header with the final CRC and sizes.
32///
33/// If the final compressed or uncompressed size requires Zip64 but no Zip64 size fields were
34/// reserved up front, closing this writer will fail. Use [`ZipEntryBuilder::size`] to reserve those
35/// fields when the size is known to exceed the non-Zip64 limit, or use
36/// [`ZipFileWriter::write_entry_stream`] for fully unknown Zip64-sized entries.
37///
38/// [`ZipEntryBuilder::size`]: crate::ZipEntryBuilder::size
39pub struct EntrySeekableWriter<'b, W: AsyncWrite + AsyncSeek + Unpin> {
40    writer: AsyncOffsetWriter<CompressedAsyncWriter<'b, W>>,
41    cd_entries: &'b mut Vec<CentralDirectoryEntry>,
42    entry: ZipEntry,
43    hasher: Hasher,
44    lfh: LocalFileHeader,
45    lfh_offset: u64,
46    data_offset: u64,
47    local_header_has_zip64_sizes: bool,
48    force_no_zip64: bool,
49    /// To write back to the original writer if Zip64 is required.
50    is_zip64: &'b mut bool,
51}
52
53impl<'b, W: AsyncWrite + AsyncSeek + Unpin> EntrySeekableWriter<'b, W> {
54    pub(crate) async fn from_raw(
55        writer: &'b mut ZipFileWriter<W>,
56        mut entry: ZipEntry,
57    ) -> Result<EntrySeekableWriter<'b, W>> {
58        if writer.force_no_zip64 && writer.cd_entries.len() >= NON_ZIP64_MAX_NUM_FILES as usize {
59            return Err(ZipError::Zip64Needed(Zip64ErrorCase::TooManyFiles));
60        }
61
62        #[cfg(feature = "deflate64")]
63        if matches!(entry.compression(), crate::Compression::Deflate64) {
64            return Err(ZipError::FeatureNotSupported("Deflate64 writing"));
65        }
66
67        let lfh_offset = writer.writer.offset();
68        let (lfh, local_header_has_zip64_sizes) = EntrySeekableWriter::write_lfh(writer, &mut entry).await?;
69        let data_offset = writer.writer.offset();
70        let force_no_zip64 = writer.force_no_zip64;
71
72        let cd_entries = &mut writer.cd_entries;
73        let is_zip64 = &mut writer.is_zip64;
74        let writer = AsyncOffsetWriter::new(CompressedAsyncWriter::from_raw(&mut writer.writer, entry.compression())?);
75
76        Ok(EntrySeekableWriter {
77            writer,
78            cd_entries,
79            entry,
80            lfh,
81            lfh_offset,
82            data_offset,
83            local_header_has_zip64_sizes,
84            hasher: Hasher::new(),
85            force_no_zip64,
86            is_zip64,
87        })
88    }
89
90    async fn write_lfh(writer: &'b mut ZipFileWriter<W>, entry: &mut ZipEntry) -> Result<(LocalFileHeader, bool)> {
91        let local_header_has_zip64_sizes =
92            entry.uncompressed_size >= NON_ZIP64_MAX_SIZE as u64 || entry.compressed_size >= NON_ZIP64_MAX_SIZE as u64;
93        if local_header_has_zip64_sizes {
94            if writer.force_no_zip64 {
95                return Err(ZipError::Zip64Needed(Zip64ErrorCase::LargeFile));
96            }
97            if !writer.is_zip64 {
98                writer.is_zip64 = true;
99            }
100            // Reserve Zip64 size slots up front so the later header patch stays the same width.
101            match get_zip64_extra_field_mut(&mut entry.extra_fields) {
102                Some(zip64) => {
103                    zip64.uncompressed_size = Some(entry.uncompressed_size);
104                    zip64.compressed_size = Some(entry.compressed_size);
105                }
106                None => {
107                    entry.extra_fields.push(ExtraField::Zip64ExtendedInformation(Zip64ExtendedInformationExtraField {
108                        uncompressed_size: Some(entry.uncompressed_size),
109                        compressed_size: Some(entry.compressed_size),
110                        relative_header_offset: None,
111                        disk_start_number: None,
112                    }));
113                }
114            }
115        }
116
117        let utf8_without_alternative =
118            entry.filename().is_utf8_without_alternative() && entry.comment().is_utf8_without_alternative();
119        if !utf8_without_alternative {
120            if matches!(entry.filename().encoding(), StringEncoding::Utf8) {
121                let u_file_name = entry.filename().as_bytes().to_vec();
122                if !u_file_name.is_empty() {
123                    let basic_crc32 =
124                        crc32fast::hash(entry.filename().alternative().unwrap_or_else(|| entry.filename().as_bytes()));
125                    let upath_field = get_or_put_info_zip_unicode_path_extra_field_mut(entry.extra_fields.as_mut());
126                    if let InfoZipUnicodePathExtraField::V1 { crc32, unicode } = upath_field {
127                        *crc32 = basic_crc32;
128                        *unicode = u_file_name;
129                    }
130                }
131            }
132            if matches!(entry.comment().encoding(), StringEncoding::Utf8) {
133                let u_comment = entry.comment().as_bytes().to_vec();
134                if !u_comment.is_empty() {
135                    let basic_crc32 =
136                        crc32fast::hash(entry.comment().alternative().unwrap_or_else(|| entry.comment().as_bytes()));
137                    let ucom_field = get_or_put_info_zip_unicode_comment_extra_field_mut(entry.extra_fields.as_mut());
138                    if let InfoZipUnicodeCommentExtraField::V1 { crc32, unicode } = ucom_field {
139                        *crc32 = basic_crc32;
140                        *unicode = u_comment;
141                    }
142                }
143            }
144        }
145
146        let filename_basic = entry.filename().alternative().unwrap_or_else(|| entry.filename().as_bytes());
147
148        let lfh = LocalFileHeader {
149            compressed_size: if local_header_has_zip64_sizes { NON_ZIP64_MAX_SIZE } else { 0 },
150            uncompressed_size: if local_header_has_zip64_sizes { NON_ZIP64_MAX_SIZE } else { 0 },
151            compression: entry.compression().into(),
152            crc: entry.crc32,
153            extra_field_length: entry
154                .extra_fields()
155                .count_bytes()
156                .try_into()
157                .map_err(|_| ZipError::ExtraFieldTooLarge)?,
158            file_name_length: filename_basic.len().try_into().map_err(|_| ZipError::FileNameTooLarge)?,
159            mod_time: entry.last_modification_date().time,
160            mod_date: entry.last_modification_date().date,
161            version: crate::spec::version::as_needed_to_extract(entry),
162            flags: GeneralPurposeFlag {
163                data_descriptor: false,
164                encrypted: false,
165                strong_encryption: false,
166                compressed_patched: false,
167                filename_unicode: utf8_without_alternative,
168            },
169        };
170
171        writer.writer.write_all(&crate::spec::consts::LFH_SIGNATURE.to_le_bytes()).await?;
172        writer.writer.write_all(&lfh.as_slice()).await?;
173        writer.writer.write_all(filename_basic).await?;
174        writer.writer.write_all(&entry.extra_fields().as_bytes()).await?;
175
176        Ok((lfh, local_header_has_zip64_sizes))
177    }
178
179    /// Consumes this entry writer and completes all closing tasks.
180    ///
181    /// This includes:
182    /// - Finalising the CRC32 hash value for the written data.
183    /// - Calculating the compressed and uncompressed byte sizes.
184    /// - Seeking back to patch the local file header.
185    /// - Constructing a central directory header.
186    /// - Pushing that central directory header to the [`ZipFileWriter`]'s store.
187    ///
188    /// Failure to call this function before going out of scope would result in a corrupted ZIP file.
189    pub async fn close(mut self) -> Result<()> {
190        self.writer.close().await?;
191
192        let crc = self.hasher.finalize();
193        let uncompressed_size = self.writer.offset();
194        let inner_writer = self.writer.into_inner().into_inner();
195        let compressed_size = inner_writer.offset() - self.data_offset;
196        let end_offset = inner_writer.offset();
197
198        let requires_zip64_sizes =
199            uncompressed_size >= NON_ZIP64_MAX_SIZE as u64 || compressed_size >= NON_ZIP64_MAX_SIZE as u64;
200        let requires_zip64_offset = self.lfh_offset >= NON_ZIP64_MAX_SIZE as u64;
201
202        if self.force_no_zip64 && (requires_zip64_sizes || requires_zip64_offset) {
203            return Err(ZipError::Zip64Needed(Zip64ErrorCase::LargeFile));
204        }
205        if requires_zip64_sizes && !self.local_header_has_zip64_sizes {
206            return Err(ZipError::Zip64Needed(Zip64ErrorCase::LargeFile));
207        }
208
209        let uses_zip64_sizes = requires_zip64_sizes || self.local_header_has_zip64_sizes;
210        let uses_zip64_metadata = uses_zip64_sizes || requires_zip64_offset;
211        if uses_zip64_sizes {
212            if !*self.is_zip64 {
213                *self.is_zip64 = true;
214            }
215            self.lfh.compressed_size = NON_ZIP64_MAX_SIZE;
216            self.lfh.uncompressed_size = NON_ZIP64_MAX_SIZE;
217            match get_zip64_extra_field_mut(&mut self.entry.extra_fields) {
218                Some(zip64) => {
219                    zip64.uncompressed_size = Some(uncompressed_size);
220                    zip64.compressed_size = Some(compressed_size);
221                }
222                None => {
223                    self.entry.extra_fields.push(ExtraField::Zip64ExtendedInformation(
224                        Zip64ExtendedInformationExtraField {
225                            uncompressed_size: Some(uncompressed_size),
226                            compressed_size: Some(compressed_size),
227                            relative_header_offset: None,
228                            disk_start_number: None,
229                        },
230                    ));
231                }
232            }
233        } else {
234            self.lfh.compressed_size = compressed_size as u32;
235            self.lfh.uncompressed_size = uncompressed_size as u32;
236        }
237        if uses_zip64_metadata {
238            self.lfh.version = self.lfh.version.max(ZIP64_VERSION_NEEDED);
239        }
240        self.lfh.crc = crc;
241        self.lfh.extra_field_length =
242            self.entry.extra_fields().count_bytes().try_into().map_err(|_| ZipError::ExtraFieldTooLarge)?;
243
244        let filename_basic = self.entry.filename().alternative().unwrap_or_else(|| self.entry.filename().as_bytes());
245        let local_extra_fields = self.entry.extra_fields().as_bytes();
246
247        inner_writer.seek(SeekFrom::Start(self.lfh_offset + crate::spec::consts::SIGNATURE_LENGTH as u64)).await?;
248        inner_writer.write_all(&self.lfh.as_slice()).await?;
249        inner_writer.write_all(filename_basic).await?;
250        inner_writer.write_all(&local_extra_fields).await?;
251        inner_writer.seek(SeekFrom::Start(end_offset)).await?;
252
253        let lh_offset = if requires_zip64_offset {
254            if !*self.is_zip64 {
255                *self.is_zip64 = true;
256            }
257            match get_zip64_extra_field_mut(&mut self.entry.extra_fields) {
258                Some(zip64) => {
259                    zip64.relative_header_offset = Some(self.lfh_offset);
260                }
261                None => {
262                    self.entry.extra_fields.push(ExtraField::Zip64ExtendedInformation(
263                        Zip64ExtendedInformationExtraField {
264                            uncompressed_size: None,
265                            compressed_size: None,
266                            relative_header_offset: Some(self.lfh_offset),
267                            disk_start_number: None,
268                        },
269                    ));
270                }
271            }
272            NON_ZIP64_MAX_SIZE
273        } else {
274            self.lfh_offset as u32
275        };
276
277        let comment_basic = self.entry.comment().alternative().unwrap_or_else(|| self.entry.comment().as_bytes());
278
279        let cdh = CentralDirectoryRecord {
280            compressed_size: self.lfh.compressed_size,
281            uncompressed_size: self.lfh.uncompressed_size,
282            crc,
283            v_made_by: crate::spec::version::as_made_by(),
284            v_needed: self.lfh.version,
285            compression: self.lfh.compression,
286            extra_field_length: self
287                .entry
288                .extra_fields()
289                .count_bytes()
290                .try_into()
291                .map_err(|_| ZipError::ExtraFieldTooLarge)?,
292            file_name_length: self.lfh.file_name_length,
293            file_comment_length: comment_basic.len().try_into().map_err(|_| ZipError::CommentTooLarge)?,
294            mod_time: self.lfh.mod_time,
295            mod_date: self.lfh.mod_date,
296            flags: self.lfh.flags,
297            disk_start: 0,
298            inter_attr: self.entry.internal_file_attribute(),
299            exter_attr: self.entry.external_file_attribute(),
300            lh_offset,
301        };
302
303        self.cd_entries.push(CentralDirectoryEntry { header: cdh, entry: self.entry });
304        // Mark the archive as Zip64 once the central directory no longer fits in the legacy count field.
305        if self.cd_entries.len() > NON_ZIP64_MAX_NUM_FILES as usize && !*self.is_zip64 {
306            *self.is_zip64 = true;
307        }
308
309        Ok(())
310    }
311}
312
313impl<'a, W: AsyncWrite + AsyncSeek + Unpin> AsyncWrite for EntrySeekableWriter<'a, W> {
314    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<std::result::Result<usize, Error>> {
315        let poll = Pin::new(&mut self.writer).poll_write(cx, buf);
316
317        if let Poll::Ready(Ok(written)) = poll {
318            self.hasher.update(&buf[0..written]);
319        }
320
321        poll
322    }
323
324    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<std::result::Result<(), Error>> {
325        Pin::new(&mut self.writer).poll_flush(cx)
326    }
327
328    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<std::result::Result<(), Error>> {
329        Pin::new(&mut self.writer).poll_close(cx)
330    }
331}