async_deflate_zip/writer/
entry_writer.rs

1use super::helpers::CountWriter;
2use super::stored_entry::StoredEntry;
3use super::zip_writer::ZipWriter;
4
5use crate::deflate_encoder::DeflateEncoder;
6use crate::error::ZipError;
7use crate::header;
8
9use std::io;
10use std::pin::Pin;
11use std::task::{Context, Poll};
12use std::time::SystemTime;
13use tokio::io::{AsyncWrite, AsyncWriteExt};
14
15pin_project_lite::pin_project! {
16    /// A streaming writer for a single file entry in a ZIP archive.
17    ///
18    /// Obtained from [`ZipWriter::append_file`]. Data written through this
19    /// writer is compressed with DEFLATE and streamed to the underlying output.
20    ///
21    /// # Important
22    ///
23    /// The [`close`](EntryWriter::close) method **must** be called after all
24    /// data is written to finalize the deflate frame and write the Data
25    /// Descriptor. Dropping without closing will lose the entry.
26    pub struct EntryWriter<'a, W>
27    where
28        W: AsyncWrite,
29        W: Unpin,
30    {
31        pub(crate) zip: &'a mut ZipWriter<W>,
32        #[pin]
33        pub(crate) deflate_encoder: Option<DeflateEncoder<CountWriter<W>>>,
34        #[pin]
35        pub(crate) passthrough: Option<CountWriter<W>>,
36        pub(crate) is_stored: bool,
37        pub(crate) crc_hasher: crc32fast::Hasher,
38        pub(crate) uncompressed_size: u64,
39        pub(crate) local_header_offset: u64,
40        pub(crate) name: String,
41        pub(crate) mtime: SystemTime,
42        pub(crate) unix_permissions: Option<u32>,
43        pub(crate) uid_gid: Option<(u32, u32)>,
44        pub(crate) comment: Option<Vec<u8>>,
45    }
46
47    impl<W> PinnedDrop for EntryWriter<'_, W>
48    where
49        W: AsyncWrite,
50        W: Unpin,
51    {
52        fn drop(this: Pin<&mut Self>) {
53            let this = this.project();
54            if this.deflate_encoder.is_some() || this.passthrough.is_some() {
55                // close() was never called — mark the ZipWriter as poisoned
56                this.zip.poisoned = true;
57            }
58        }
59    }
60}
61
62impl<W: AsyncWrite + Unpin> EntryWriter<'_, W> {
63    /// Finalize the deflate frame, compute the CRC-32 checksum, and write
64    /// the Data Descriptor.
65    ///
66    /// This consumes the `EntryWriter`, flushes the deflate encoder, extracts
67    /// the compressed size, computes the CRC-32 of the uncompressed data, and
68    /// writes the trailing CRC-32 and sizes (Data Descriptor) after the compressed
69    /// data. The inner writer is returned to the parent [`ZipWriter`].
70    ///
71    /// # Errors
72    ///
73    /// Returns [`ZipError`] if `close` is called more than once, if the deflate
74    /// encoder fails to shut down (I/O error), or if writing the Data
75    /// Descriptor fails (I/O error).
76    pub async fn close(mut self) -> Result<(), ZipError> {
77        let (compressed_size, mut inner) = if self.is_stored {
78            let cw = self
79                .passthrough
80                .take()
81                .ok_or_else(|| ZipError::InvalidState("entry already closed".to_string()))?;
82            (cw.count, cw.inner)
83        } else {
84            let mut encoder = self
85                .deflate_encoder
86                .take()
87                .ok_or_else(|| ZipError::InvalidState("entry already closed".to_string()))?;
88            encoder.shutdown().await?;
89
90            // Extract the inner writer from the encoder stack
91            let count_writer: CountWriter<W> = encoder.into_inner();
92            let compressed_size = count_writer.count;
93            (compressed_size, count_writer.inner)
94        };
95
96        let crc32 = self.crc_hasher.clone().finalize();
97
98        let dd = header::DataDescriptor {
99            crc32,
100            compressed_size,
101            uncompressed_size: self.uncompressed_size,
102            // Use ZIP64 DD when any entry field exceeds 32 bits, consistent with
103            // CentralDirEntry::serialize() which also checks local_header_offset.
104            zip64: compressed_size > header::U32_MAX
105                || self.uncompressed_size > header::U32_MAX
106                || self.local_header_offset > header::U32_MAX,
107        };
108        let dd_bytes = dd.serialize();
109        inner.write_all(&dd_bytes).await.map_err(|e| {
110            self.zip.poisoned = true;
111            ZipError::Io(e)
112        })?;
113
114        // Update position tracker: compressed data + data descriptor
115        self.zip.pos += compressed_size + dd_bytes.len() as u64;
116
117        let (mtime_msdos, unix_mtime) = header::mtime_to_ms_dos_and_unix(self.mtime);
118
119        self.zip.entries.push(StoredEntry {
120            name: self.name.clone(),
121            crc32,
122            compressed_size,
123            uncompressed_size: self.uncompressed_size,
124            local_header_offset: self.local_header_offset,
125            is_directory: false,
126            is_symlink: false,
127            is_stored: self.is_stored,
128            mtime: mtime_msdos,
129            unix_mtime,
130            unix_permissions: self.unix_permissions,
131            uid_gid: self.uid_gid,
132            comment: self.comment.clone(),
133        });
134
135        // Return the inner writer to ZipWriter
136        self.zip.inner = Some(inner);
137        Ok(())
138    }
139}
140
141impl<W: AsyncWrite + Unpin> AsyncWrite for EntryWriter<'_, W> {
142    fn poll_write(
143        self: Pin<&mut Self>,
144        cx: &mut Context<'_>,
145        buf: &[u8],
146    ) -> Poll<io::Result<usize>> {
147        let this = self.project();
148        let result = if *this.is_stored {
149            match this.passthrough.as_pin_mut() {
150                Some(w) => w.poll_write(cx, buf),
151                None => {
152                    this.zip.poisoned = true;
153                    return Poll::Ready(Err(ZipError::Poisoned(
154                        "write after entry closed".to_string(),
155                    )
156                    .into()));
157                }
158            }
159        } else {
160            match this.deflate_encoder.as_pin_mut() {
161                Some(e) => e.poll_write(cx, buf),
162                None => {
163                    this.zip.poisoned = true;
164                    return Poll::Ready(Err(ZipError::Poisoned(
165                        "write after entry closed".to_string(),
166                    )
167                    .into()));
168                }
169            }
170        };
171        match result {
172            Poll::Ready(Ok(n)) => {
173                this.crc_hasher.update(&buf[..n]);
174                *this.uncompressed_size += n as u64;
175                Poll::Ready(Ok(n))
176            }
177            other => other,
178        }
179    }
180
181    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
182        let this = self.project();
183        if *this.is_stored {
184            match this.passthrough.as_pin_mut() {
185                Some(w) => w.poll_flush(cx),
186                None => {
187                    this.zip.poisoned = true;
188                    Poll::Ready(Err(ZipError::Poisoned(
189                        "flush after entry closed".to_string(),
190                    )
191                    .into()))
192                }
193            }
194        } else {
195            match this.deflate_encoder.as_pin_mut() {
196                Some(e) => e.poll_flush(cx),
197                None => {
198                    this.zip.poisoned = true;
199                    Poll::Ready(Err(ZipError::Poisoned(
200                        "flush after entry closed".to_string(),
201                    )
202                    .into()))
203                }
204            }
205        }
206    }
207
208    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
209        let this = self.project();
210        if *this.is_stored {
211            match this.passthrough.as_pin_mut() {
212                Some(w) => w.poll_shutdown(cx),
213                None => {
214                    this.zip.poisoned = true;
215                    Poll::Ready(Err(ZipError::Poisoned(
216                        "shutdown after entry closed".to_string(),
217                    )
218                    .into()))
219                }
220            }
221        } else {
222            match this.deflate_encoder.as_pin_mut() {
223                Some(e) => e.poll_shutdown(cx),
224                None => {
225                    this.zip.poisoned = true;
226                    Poll::Ready(Err(ZipError::Poisoned(
227                        "shutdown after entry closed".to_string(),
228                    )
229                    .into()))
230                }
231            }
232        }
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::super::*;
239    use std::time::SystemTime;
240    use tokio::io::AsyncWriteExt;
241
242    fn opts_mtime(t: SystemTime) -> EntryOptions {
243        EntryOptions {
244            mtime: t,
245            permissions: None,
246            uid_gid: None,
247            comment: None,
248        }
249    }
250
251    fn opts_perms(mode: u32) -> EntryOptions {
252        EntryOptions {
253            mtime: SystemTime::now(),
254            permissions: Some(mode),
255            uid_gid: None,
256            comment: None,
257        }
258    }
259
260    #[tokio::test]
261    async fn test_entry_mtime_epoch() {
262        let mut buf = Vec::new();
263        let mut zip = ZipWriter::new(&mut buf);
264        let mut entry = zip
265            .append_file("epoch.txt", opts_mtime(SystemTime::UNIX_EPOCH))
266            .await
267            .unwrap();
268        entry.write_all(b"test").await.unwrap();
269        entry.close().await.unwrap();
270        zip.finalize().await.unwrap();
271
272        let pos = buf.windows(4).position(|w| w == b"PK\x01\x02").unwrap();
273        let cd = &buf[pos..];
274
275        let time = u16::from_le_bytes(cd[12..14].try_into().unwrap());
276        let date = u16::from_le_bytes(cd[14..16].try_into().unwrap());
277        let local_offset = time::UtcOffset::current_local_offset().unwrap_or(time::UtcOffset::UTC);
278        let local_epoch =
279            time::OffsetDateTime::from(SystemTime::UNIX_EPOCH).to_offset(local_offset);
280        let expected_time = (local_epoch.hour() as u16) << 11
281            | (local_epoch.minute() as u16) << 5
282            | (local_epoch.second() as u16).div_ceil(2);
283        assert_eq!(time, expected_time, "expected local time for epoch");
284        assert_eq!(date, (1 << 5) | 1, "expected MS-DOS date for 1980-01-01");
285    }
286
287    #[tokio::test]
288    async fn test_entry_permissions() {
289        let mut buf = Vec::new();
290        let mut zip = ZipWriter::new(&mut buf);
291        let mut entry = zip
292            .append_file("perm_test.txt", opts_perms(0o644))
293            .await
294            .unwrap();
295        entry.write_all(b"test").await.unwrap();
296        entry.close().await.unwrap();
297        zip.finalize().await.unwrap();
298
299        let pos = buf.windows(4).position(|w| w == b"PK\x01\x02").unwrap();
300        let cd = &buf[pos..];
301        let efa = u32::from_le_bytes(cd[38..42].try_into().unwrap());
302        assert_eq!(efa, ((0o644 | 0o100000) as u32) << 16);
303        let vmb = u16::from_le_bytes(cd[4..6].try_into().unwrap());
304        assert!(vmb >> 8 == 3, "expected Unix host OS");
305    }
306
307    #[tokio::test]
308    async fn test_entry_setuid_permissions() {
309        let mut buf = Vec::new();
310        let mut zip = ZipWriter::new(&mut buf);
311        let mut entry = zip
312            .append_file("setuid_test.txt", opts_perms(0o4755))
313            .await
314            .unwrap();
315        entry.write_all(b"test").await.unwrap();
316        entry.close().await.unwrap();
317        zip.finalize().await.unwrap();
318
319        let pos = buf.windows(4).position(|w| w == b"PK\x01\x02").unwrap();
320        let cd = &buf[pos..];
321        let efa = u32::from_le_bytes(cd[38..42].try_into().unwrap());
322        assert_eq!(efa, ((0o4755 | 0o100000) as u32) << 16);
323    }
324
325    #[tokio::test]
326    async fn test_entry_mtime_appears_in_cd_extra() {
327        let mut buf = Vec::new();
328        let mut zip = ZipWriter::new(&mut buf);
329        let mut entry = zip
330            .append_file("mtime_test.txt", opts_mtime(SystemTime::now()))
331            .await
332            .unwrap();
333        entry.write_all(b"hello").await.unwrap();
334        entry.close().await.unwrap();
335        zip.finalize().await.unwrap();
336
337        let pos = buf.windows(4).position(|w| w == b"PK\x01\x02").unwrap();
338        let cd = &buf[pos..];
339        let name_len = u16::from_le_bytes(cd[28..30].try_into().unwrap()) as usize;
340        let extra_len = u16::from_le_bytes(cd[30..32].try_into().unwrap()) as usize;
341
342        let extra_start = 46 + name_len;
343        let extra = &cd[extra_start..extra_start + extra_len];
344        let has_ts_extra = extra.windows(2).any(|w| w == b"UT");
345        assert!(
346            has_ts_extra,
347            "CD entry extra should contain extended timestamp (0x5455/UT) when mtime is set"
348        );
349        assert!(
350            extra_len >= 4,
351            "extra_len should be >= 4 when mtime is set, got {extra_len}"
352        );
353        let vmb = u16::from_le_bytes(cd[4..6].try_into().unwrap());
354        assert_eq!(vmb >> 8, 3, "expected Unix host OS when mtime is set");
355    }
356
357    #[tokio::test]
358    async fn test_entry_drop_poisons_zip_writer() {
359        let mut buf = Vec::new();
360        let mut zip = ZipWriter::new(&mut buf);
361
362        drop(
363            zip.append_file("lost.txt", EntryOptions::file())
364                .await
365                .unwrap(),
366        );
367
368        let result = zip.append_file("another.txt", EntryOptions::file()).await;
369        assert!(result.is_err(), "expected Err, got Ok");
370        let err = result.err().unwrap();
371        assert!(
372            err.to_string().contains("archive corrupted"),
373            "expected 'archive corrupted', got: {err}"
374        );
375    }
376
377    #[tokio::test]
378    async fn test_entry_drop_poison_affects_finalize() {
379        let mut buf = Vec::new();
380        let mut zip = ZipWriter::new(&mut buf);
381
382        drop(
383            zip.append_file("lost.txt", EntryOptions::file())
384                .await
385                .unwrap(),
386        );
387
388        let err = zip.finalize().await.unwrap_err();
389        assert!(
390            err.to_string().contains("archive corrupted"),
391            "expected 'archive corrupted', got: {err}"
392        );
393    }
394}
async_deflate_zip/writer/entry_writer.rs

async_deflate_zip/writer/
entry_writer.rs
