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

async_zip/base/write/
entry_stream.rs
