async_zip/base/write/
entry_stream.rs

1// Copyright (c) 2021 Harry [Majored] [hello@majored.pw]
2// MIT License (https://github.com/Majored/rs-async-zip/blob/main/LICENSE)
3
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;
9use crate::base::write::ZipFileWriter;
10use crate::entry::ZipEntry;
11use crate::error::{Result, Zip64ErrorCase, ZipError};
12use crate::spec::data_descriptor::{DataDescriptor, Zip64DataDescriptor};
13use crate::spec::extra_field::ExtraFieldAsBytes;
14use crate::spec::header::InfoZipUnicodeCommentExtraField;
15use crate::spec::header::InfoZipUnicodePathExtraField;
16use crate::spec::header::{
17    CentralDirectoryRecord, ExtraField, GeneralPurposeFlag, LocalFileHeader, Zip64ExtendedInformationExtraField,
18};
19use crate::string::StringEncoding;
20
21use std::io::Error;
22use std::pin::Pin;
23use std::task::{Context, Poll};
24
25use crate::base::read::get_zip64_extra_field_mut;
26use crate::spec::consts::{NON_ZIP64_MAX_NUM_FILES, NON_ZIP64_MAX_SIZE};
27use crc32fast::Hasher;
28use futures_lite::io::{AsyncWrite, AsyncWriteExt};
29
30/// An entry writer which supports the streaming of data (ie. the writing of unknown size or data at runtime).
31///
32/// # Note
33/// - This writer cannot be manually constructed; instead, use [`ZipFileWriter::write_entry_stream()`].
34/// - [`EntryStreamWriter::close()`] must be called before a stream writer goes out of scope.
35/// - Utilities for working with [`AsyncWrite`] values are provided by [`AsyncWriteExt`].
36pub struct EntryStreamWriter<'b, W: AsyncWrite + Unpin> {
37    writer: AsyncOffsetWriter<CompressedAsyncWriter<'b, W>>,
38    cd_entries: &'b mut Vec<CentralDirectoryEntry>,
39    entry: ZipEntry,
40    hasher: Hasher,
41    lfh: LocalFileHeader,
42    lfh_offset: u64,
43    data_offset: u64,
44    force_no_zip64: bool,
45    /// To write back to the original writer if zip64 is required.
46    is_zip64: &'b mut bool,
47}
48
49impl<'b, W: AsyncWrite + Unpin> EntryStreamWriter<'b, W> {
50    pub(crate) async fn from_raw(
51        writer: &'b mut ZipFileWriter<W>,
52        mut entry: ZipEntry,
53    ) -> Result<EntryStreamWriter<'b, W>> {
54        if writer.force_no_zip64 && writer.cd_entries.len() >= NON_ZIP64_MAX_NUM_FILES as usize {
55            return Err(ZipError::Zip64Needed(Zip64ErrorCase::TooManyFiles));
56        }
57
58        #[cfg(feature = "deflate64")]
59        if matches!(entry.compression(), crate::Compression::Deflate64) {
60            return Err(ZipError::FeatureNotSupported("Deflate64 writing"));
61        }
62
63        let lfh_offset = writer.writer.offset();
64        let lfh = EntryStreamWriter::write_lfh(writer, &mut entry).await?;
65        let data_offset = writer.writer.offset();
66        let force_no_zip64 = writer.force_no_zip64;
67
68        let cd_entries = &mut writer.cd_entries;
69        let is_zip64 = &mut writer.is_zip64;
70        let writer = AsyncOffsetWriter::new(CompressedAsyncWriter::from_raw(&mut writer.writer, entry.compression())?);
71
72        Ok(EntryStreamWriter {
73            writer,
74            cd_entries,
75            entry,
76            lfh,
77            lfh_offset,
78            data_offset,
79            hasher: Hasher::new(),
80            force_no_zip64,
81            is_zip64,
82        })
83    }
84
85    async fn write_lfh(writer: &'b mut ZipFileWriter<W>, entry: &mut ZipEntry) -> Result<LocalFileHeader> {
86        // Always emit a zip64 extended field, even if we don't need it, because we *might* need it.
87        // If we are forcing no zip, we will have to error later if the file is too large.
88        let (lfh_compressed, lfh_uncompressed) = if !writer.force_no_zip64 {
89            if !writer.is_zip64 {
90                writer.is_zip64 = true;
91            }
92            entry.extra_fields.push(ExtraField::Zip64ExtendedInformation(Zip64ExtendedInformationExtraField {
93                uncompressed_size: Some(entry.uncompressed_size),
94                compressed_size: Some(entry.compressed_size),
95                relative_header_offset: None,
96                disk_start_number: None,
97            }));
98
99            (NON_ZIP64_MAX_SIZE, NON_ZIP64_MAX_SIZE)
100        } else {
101            if entry.compressed_size > NON_ZIP64_MAX_SIZE as u64 || entry.uncompressed_size > NON_ZIP64_MAX_SIZE as u64
102            {
103                return Err(ZipError::Zip64Needed(Zip64ErrorCase::LargeFile));
104            }
105
106            (entry.compressed_size as u32, entry.uncompressed_size as u32)
107        };
108
109        let utf8_without_alternative =
110            entry.filename().is_utf8_without_alternative() && entry.comment().is_utf8_without_alternative();
111        if !utf8_without_alternative {
112            if matches!(entry.filename().encoding(), StringEncoding::Utf8) {
113                let u_file_name = entry.filename().as_bytes().to_vec();
114                if !u_file_name.is_empty() {
115                    let basic_crc32 =
116                        crc32fast::hash(entry.filename().alternative().unwrap_or_else(|| entry.filename().as_bytes()));
117                    let upath_field = get_or_put_info_zip_unicode_path_extra_field_mut(entry.extra_fields.as_mut());
118                    if let InfoZipUnicodePathExtraField::V1 { crc32, unicode } = upath_field {
119                        *crc32 = basic_crc32;
120                        *unicode = u_file_name;
121                    }
122                }
123            }
124            if matches!(entry.comment().encoding(), StringEncoding::Utf8) {
125                let u_comment = entry.comment().as_bytes().to_vec();
126                if !u_comment.is_empty() {
127                    let basic_crc32 =
128                        crc32fast::hash(entry.comment().alternative().unwrap_or_else(|| entry.comment().as_bytes()));
129                    let ucom_field = get_or_put_info_zip_unicode_comment_extra_field_mut(entry.extra_fields.as_mut());
130                    if let InfoZipUnicodeCommentExtraField::V1 { crc32, unicode } = ucom_field {
131                        *crc32 = basic_crc32;
132                        *unicode = u_comment;
133                    }
134                }
135            }
136        }
137
138        let filename_basic = entry.filename().alternative().unwrap_or_else(|| entry.filename().as_bytes());
139
140        let lfh = LocalFileHeader {
141            compressed_size: lfh_compressed,
142            uncompressed_size: lfh_uncompressed,
143            compression: entry.compression().into(),
144            crc: entry.crc32,
145            extra_field_length: entry
146                .extra_fields()
147                .count_bytes()
148                .try_into()
149                .map_err(|_| ZipError::ExtraFieldTooLarge)?,
150            file_name_length: filename_basic.len().try_into().map_err(|_| ZipError::FileNameTooLarge)?,
151            mod_time: entry.last_modification_date().time,
152            mod_date: entry.last_modification_date().date,
153            version: crate::spec::version::as_needed_to_extract(entry),
154            flags: GeneralPurposeFlag {
155                data_descriptor: true,
156                encrypted: false,
157                strong_encryption: false,
158                compressed_patched: false,
159                filename_unicode: utf8_without_alternative,
160            },
161        };
162
163        writer.writer.write_all(&crate::spec::consts::LFH_SIGNATURE.to_le_bytes()).await?;
164        writer.writer.write_all(&lfh.as_slice()).await?;
165        writer.writer.write_all(filename_basic).await?;
166        writer.writer.write_all(&entry.extra_fields().as_bytes()).await?;
167
168        Ok(lfh)
169    }
170
171    /// Consumes this entry writer and completes all closing tasks.
172    ///
173    /// This includes:
174    /// - Finalising the CRC32 hash value for the written data.
175    /// - Calculating the compressed and uncompressed byte sizes.
176    /// - Constructing a central directory header.
177    /// - Pushing that central directory header to the [`ZipFileWriter`]'s store.
178    ///
179    /// Failure to call this function before going out of scope would result in a corrupted ZIP file.
180    pub async fn close(mut self) -> Result<()> {
181        self.writer.close().await?;
182
183        let crc = self.hasher.finalize();
184        let uncompressed_size = self.writer.offset();
185        let inner_writer = self.writer.into_inner().into_inner();
186        let compressed_size = inner_writer.offset() - self.data_offset;
187
188        let (cdr_compressed_size, cdr_uncompressed_size, lh_offset) = if self.force_no_zip64 {
189            if uncompressed_size > NON_ZIP64_MAX_SIZE as u64
190                || compressed_size > NON_ZIP64_MAX_SIZE as u64
191                || self.lfh_offset > NON_ZIP64_MAX_SIZE as u64
192            {
193                return Err(ZipError::Zip64Needed(Zip64ErrorCase::LargeFile));
194            }
195            (uncompressed_size as u32, compressed_size as u32, self.lfh_offset as u32)
196        } else {
197            // When streaming an entry, we are always using a zip64 field.
198            match get_zip64_extra_field_mut(&mut self.entry.extra_fields) {
199                // This case shouldn't be necessary but is included for completeness.
200                None => {
201                    self.entry.extra_fields.push(ExtraField::Zip64ExtendedInformation(
202                        Zip64ExtendedInformationExtraField {
203                            uncompressed_size: Some(uncompressed_size),
204                            compressed_size: Some(compressed_size),
205                            relative_header_offset: Some(self.lfh_offset),
206                            disk_start_number: None,
207                        },
208                    ));
209                }
210                Some(zip64) => {
211                    zip64.uncompressed_size = Some(uncompressed_size);
212                    zip64.compressed_size = Some(compressed_size);
213                    zip64.relative_header_offset = Some(self.lfh_offset);
214                }
215            }
216            self.lfh.extra_field_length =
217                self.entry.extra_fields().count_bytes().try_into().map_err(|_| ZipError::ExtraFieldTooLarge)?;
218
219            (NON_ZIP64_MAX_SIZE, NON_ZIP64_MAX_SIZE, NON_ZIP64_MAX_SIZE)
220        };
221
222        inner_writer.write_all(&crate::spec::consts::DATA_DESCRIPTOR_SIGNATURE.to_le_bytes()).await?;
223        if self.force_no_zip64 {
224            let descriptor =
225                DataDescriptor { crc, compressed_size: cdr_compressed_size, uncompressed_size: cdr_uncompressed_size };
226            inner_writer.write_all(&descriptor.as_bytes()).await?;
227        } else {
228            let descriptor = Zip64DataDescriptor { crc, compressed_size, uncompressed_size };
229            inner_writer.write_all(&descriptor.as_bytes()).await?;
230        }
231
232        let comment_basic = self.entry.comment().alternative().unwrap_or_else(|| self.entry.comment().as_bytes());
233
234        let cdh = CentralDirectoryRecord {
235            compressed_size: cdr_compressed_size,
236            uncompressed_size: cdr_uncompressed_size,
237            crc,
238            v_made_by: crate::spec::version::as_made_by(),
239            v_needed: self.lfh.version,
240            compression: self.lfh.compression,
241            extra_field_length: self.lfh.extra_field_length,
242            file_name_length: self.lfh.file_name_length,
243            file_comment_length: comment_basic.len().try_into().map_err(|_| ZipError::CommentTooLarge)?,
244            mod_time: self.lfh.mod_time,
245            mod_date: self.lfh.mod_date,
246            flags: self.lfh.flags,
247            disk_start: 0,
248            inter_attr: self.entry.internal_file_attribute(),
249            exter_attr: self.entry.external_file_attribute(),
250            lh_offset,
251        };
252
253        self.cd_entries.push(CentralDirectoryEntry { header: cdh, entry: self.entry });
254        // Mark the archive as Zip64 once the central directory no longer fits in the legacy count field.
255        if self.cd_entries.len() > NON_ZIP64_MAX_NUM_FILES as usize && !*self.is_zip64 {
256            *self.is_zip64 = true;
257        }
258
259        Ok(())
260    }
261}
262
263impl<'a, W: AsyncWrite + Unpin> AsyncWrite for EntryStreamWriter<'a, W> {
264    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<std::result::Result<usize, Error>> {
265        let poll = Pin::new(&mut self.writer).poll_write(cx, buf);
266
267        if let Poll::Ready(Ok(written)) = poll {
268            self.hasher.update(&buf[0..written]);
269        }
270
271        poll
272    }
273
274    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<std::result::Result<(), Error>> {
275        Pin::new(&mut self.writer).poll_flush(cx)
276    }
277
278    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<std::result::Result<(), Error>> {
279        Pin::new(&mut self.writer).poll_close(cx)
280    }
281}
async_zip/base/write/entry_stream.rs

async_zip/base/write/
entry_stream.rs
