Skip to main content

s_zip/
writer.rs

1//! Streaming ZIP writer that compresses data on-the-fly without temp files
2//!
3//! This eliminates:
4//! - Temp file disk I/O
5//! - File read buffers
6//! - Intermediate storage
7//!
8//! Expected RAM savings: 5-8 MB per file
9//!
10//! Now supports arbitrary writers (File, `Vec<u8>`, network streams, etc.)
11
12use crate::error::{Result, SZipError};
13use crc32fast::Hasher as Crc32;
14use flate2::write::DeflateEncoder;
15use flate2::Compression;
16use std::fs::File;
17use std::io::{Seek, Write};
18use std::path::Path;
19
20#[cfg(feature = "encryption")]
21use crate::encryption::{AesEncryptor, AesStrength};
22
23/// Compression method to use for ZIP entries
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum CompressionMethod {
26    /// No compression (stored)
27    Stored,
28    /// DEFLATE compression (most common)
29    Deflate,
30    /// Zstd compression (requires zstd-support feature)
31    #[cfg(feature = "zstd-support")]
32    Zstd,
33}
34
35impl CompressionMethod {
36    pub(crate) fn to_zip_method(self) -> u16 {
37        match self {
38            CompressionMethod::Stored => 0,
39            CompressionMethod::Deflate => 8,
40            #[cfg(feature = "zstd-support")]
41            CompressionMethod::Zstd => 93,
42        }
43    }
44}
45
46/// Entry being written to ZIP
47struct ZipEntry {
48    name: String,
49    local_header_offset: u64,
50    crc32: u32,
51    compressed_size: u64,
52    uncompressed_size: u64,
53    compression_method: u16,
54    #[cfg(feature = "encryption")]
55    #[allow(dead_code)] // Will be used for central directory in future versions
56    encryption_strength: Option<u16>,
57}
58
59/// Streaming ZIP writer that compresses data on-the-fly
60pub struct StreamingZipWriter<W: Write + Seek> {
61    output: W,
62    entries: Vec<ZipEntry>,
63    current_entry: Option<CurrentEntry>,
64    compression_level: u32,
65    compression_method: CompressionMethod,
66    #[cfg(feature = "encryption")]
67    password: Option<String>,
68    #[cfg(feature = "encryption")]
69    encryption_strength: AesStrength,
70}
71
72struct CurrentEntry {
73    name: String,
74    local_header_offset: u64,
75    encoder: Box<dyn CompressorWrite>,
76    counter: CrcCounter,
77    compression_method: u16,
78    #[cfg(feature = "encryption")]
79    encryptor: Option<AesEncryptor>,
80}
81
82trait CompressorWrite: Write {
83    fn finish_compression(self: Box<Self>) -> Result<CompressedBuffer>;
84    fn get_buffer_mut(&mut self) -> &mut CompressedBuffer;
85}
86
87struct DeflateCompressor {
88    encoder: DeflateEncoder<CompressedBuffer>,
89}
90
91impl Write for DeflateCompressor {
92    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
93        self.encoder.write(buf)
94    }
95
96    fn flush(&mut self) -> std::io::Result<()> {
97        self.encoder.flush()
98    }
99}
100
101impl CompressorWrite for DeflateCompressor {
102    fn finish_compression(self: Box<Self>) -> Result<CompressedBuffer> {
103        Ok(self.encoder.finish()?)
104    }
105
106    fn get_buffer_mut(&mut self) -> &mut CompressedBuffer {
107        self.encoder.get_mut()
108    }
109}
110
111/// Stored (no compression) pass-through compressor
112struct StoredCompressor {
113    buffer: CompressedBuffer,
114}
115
116impl Write for StoredCompressor {
117    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
118        self.buffer.write(buf)
119    }
120
121    fn flush(&mut self) -> std::io::Result<()> {
122        self.buffer.flush()
123    }
124}
125
126impl CompressorWrite for StoredCompressor {
127    fn finish_compression(self: Box<Self>) -> Result<CompressedBuffer> {
128        Ok(self.buffer)
129    }
130
131    fn get_buffer_mut(&mut self) -> &mut CompressedBuffer {
132        &mut self.buffer
133    }
134}
135
136#[cfg(feature = "zstd-support")]
137struct ZstdCompressor {
138    encoder: zstd::Encoder<'static, CompressedBuffer>,
139}
140
141#[cfg(feature = "zstd-support")]
142impl Write for ZstdCompressor {
143    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
144        self.encoder.write(buf)
145    }
146
147    fn flush(&mut self) -> std::io::Result<()> {
148        self.encoder.flush()
149    }
150}
151
152#[cfg(feature = "zstd-support")]
153impl CompressorWrite for ZstdCompressor {
154    fn finish_compression(self: Box<Self>) -> Result<CompressedBuffer> {
155        Ok(self.encoder.finish()?)
156    }
157
158    fn get_buffer_mut(&mut self) -> &mut CompressedBuffer {
159        self.encoder.get_mut()
160    }
161}
162
163/// Metadata tracker for CRC and byte counts
164struct CrcCounter {
165    crc: Crc32,
166    uncompressed_count: u64,
167    compressed_count: u64,
168}
169
170impl CrcCounter {
171    fn new() -> Self {
172        Self {
173            crc: Crc32::new(),
174            uncompressed_count: 0,
175            compressed_count: 0,
176        }
177    }
178
179    fn update_uncompressed(&mut self, data: &[u8]) {
180        self.crc.update(data);
181        self.uncompressed_count += data.len() as u64;
182    }
183
184    fn add_compressed(&mut self, count: u64) {
185        self.compressed_count += count;
186    }
187
188    fn finalize(&self) -> u32 {
189        self.crc.clone().finalize()
190    }
191}
192
193/// Buffered writer for compressed data with adaptive sizing
194///
195/// Automatically adjusts buffer capacity and flush threshold based on data size hints
196/// to optimize memory usage and performance for different file sizes.
197struct CompressedBuffer {
198    buffer: Vec<u8>,
199    flush_threshold: usize,
200}
201
202impl CompressedBuffer {
203    /// Create buffer with default capacity (for backward compatibility)
204    #[allow(dead_code)]
205    fn new() -> Self {
206        Self::with_size_hint(None)
207    }
208
209    /// Create buffer with adaptive sizing based on expected data size
210    ///
211    /// Optimizes initial capacity and flush threshold:
212    /// - Tiny files (<10KB): 8KB initial, 256KB threshold
213    /// - Small files (<100KB): 32KB initial, 512KB threshold  
214    /// - Medium files (<1MB): 128KB initial, 2MB threshold
215    /// - Large files (≥1MB): 256KB initial, 4MB threshold
216    fn with_size_hint(size_hint: Option<u64>) -> Self {
217        let (initial_capacity, flush_threshold) = match size_hint {
218            Some(size) if size < 10_000 => (8 * 1024, 256 * 1024), // Tiny: 8KB, 256KB
219            Some(size) if size < 100_000 => (32 * 1024, 512 * 1024), // Small: 32KB, 512KB
220            Some(size) if size < 1_000_000 => (128 * 1024, 2 * 1024 * 1024), // Medium: 128KB, 2MB
221            Some(size) if size < 10_000_000 => (256 * 1024, 4 * 1024 * 1024), // Large: 256KB, 4MB
222            _ => (512 * 1024, 8 * 1024 * 1024),                    // Very large: 512KB, 8MB
223        };
224
225        Self {
226            buffer: Vec::with_capacity(initial_capacity),
227            flush_threshold,
228        }
229    }
230
231    fn take(&mut self) -> Vec<u8> {
232        std::mem::take(&mut self.buffer)
233    }
234
235    fn should_flush(&self) -> bool {
236        self.buffer.len() >= self.flush_threshold
237    }
238}
239
240impl Write for CompressedBuffer {
241    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
242        self.buffer.extend_from_slice(buf);
243        Ok(buf.len())
244    }
245
246    fn flush(&mut self) -> std::io::Result<()> {
247        Ok(())
248    }
249}
250
251impl StreamingZipWriter<File> {
252    /// Create a new ZIP writer with default compression level (6) using DEFLATE
253    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
254        Self::with_compression(path, 6)
255    }
256
257    /// Create a new ZIP writer with custom compression level (0-9) using DEFLATE
258    pub fn with_compression<P: AsRef<Path>>(path: P, compression_level: u32) -> Result<Self> {
259        Self::with_method(path, CompressionMethod::Deflate, compression_level)
260    }
261
262    /// Create a new ZIP writer with specified compression method and level
263    ///
264    /// # Arguments
265    /// * `path` - Path to the output ZIP file
266    /// * `method` - Compression method to use (Deflate, Zstd, or Stored)
267    /// * `compression_level` - Compression level (0-9 for DEFLATE, 1-21 for Zstd)
268    pub fn with_method<P: AsRef<Path>>(
269        path: P,
270        method: CompressionMethod,
271        compression_level: u32,
272    ) -> Result<Self> {
273        let output = File::create(path)?;
274        Ok(Self {
275            output,
276            entries: Vec::new(),
277            current_entry: None,
278            compression_level,
279            compression_method: method,
280            #[cfg(feature = "encryption")]
281            password: None,
282            #[cfg(feature = "encryption")]
283            encryption_strength: AesStrength::Aes256,
284        })
285    }
286
287    /// Create a new ZIP writer with Zstd compression (requires zstd-support feature)
288    #[cfg(feature = "zstd-support")]
289    pub fn with_zstd<P: AsRef<Path>>(path: P, compression_level: i32) -> Result<Self> {
290        let output = File::create(path)?;
291        Ok(Self {
292            output,
293            entries: Vec::new(),
294            current_entry: None,
295            compression_level: compression_level as u32,
296            compression_method: CompressionMethod::Zstd,
297            #[cfg(feature = "encryption")]
298            password: None,
299            #[cfg(feature = "encryption")]
300            encryption_strength: AesStrength::Aes256,
301        })
302    }
303}
304
305impl<W: Write + Seek> StreamingZipWriter<W> {
306    /// Create a new ZIP writer from an arbitrary writer with default compression level (6) using DEFLATE
307    pub fn from_writer(writer: W) -> Result<Self> {
308        Self::from_writer_with_compression(writer, 6)
309    }
310
311    /// Create a new ZIP writer from an arbitrary writer with custom compression level
312    pub fn from_writer_with_compression(writer: W, compression_level: u32) -> Result<Self> {
313        Self::from_writer_with_method(writer, CompressionMethod::Deflate, compression_level)
314    }
315
316    /// Create a new ZIP writer from an arbitrary writer with specified compression method and level
317    ///
318    /// # Arguments
319    /// * `writer` - Any writer implementing Write + Seek
320    /// * `method` - Compression method to use (Deflate, Zstd, or Stored)
321    /// * `compression_level` - Compression level (0-9 for DEFLATE, 1-21 for Zstd)
322    pub fn from_writer_with_method(
323        writer: W,
324        method: CompressionMethod,
325        compression_level: u32,
326    ) -> Result<Self> {
327        Ok(Self {
328            output: writer,
329            entries: Vec::new(),
330            current_entry: None,
331            compression_level,
332            compression_method: method,
333            #[cfg(feature = "encryption")]
334            password: None,
335            #[cfg(feature = "encryption")]
336            encryption_strength: AesStrength::Aes256,
337        })
338    }
339
340    /// Set password for AES encryption (requires encryption feature)
341    ///
342    /// All subsequent entries will be encrypted with AES-256 using the provided password.
343    /// Call this method before `start_entry()` to encrypt files.
344    ///
345    /// # Arguments
346    /// * `password` - Password for encryption (minimum 8 characters recommended)
347    ///
348    /// # Example
349    /// ```no_run
350    /// use s_zip::StreamingZipWriter;
351    ///
352    /// let mut writer = StreamingZipWriter::new("encrypted.zip")?;
353    /// writer.set_password("my_secure_password");
354    ///
355    /// writer.start_entry("secret.txt")?;
356    /// writer.write_data(b"Confidential data")?;
357    /// writer.finish()?;
358    /// # Ok::<(), s_zip::SZipError>(())
359    /// ```
360    #[cfg(feature = "encryption")]
361    pub fn set_password(&mut self, password: impl Into<String>) -> &mut Self {
362        self.password = Some(password.into());
363        self
364    }
365
366    /// Set AES encryption strength (default: AES-256)
367    ///
368    /// # Arguments
369    /// * `strength` - AES encryption strength (Aes128, Aes192, or Aes256)
370    #[cfg(feature = "encryption")]
371    pub fn set_encryption_strength(&mut self, strength: AesStrength) -> &mut Self {
372        self.encryption_strength = strength;
373        self
374    }
375
376    /// Clear password (disable encryption for subsequent entries)
377    #[cfg(feature = "encryption")]
378    pub fn clear_password(&mut self) -> &mut Self {
379        self.password = None;
380        self
381    }
382
383    /// Start a new entry (file) in the ZIP
384    pub fn start_entry(&mut self, name: &str) -> Result<()> {
385        self.start_entry_with_hint(name, None)
386    }
387
388    /// Start a new entry with file metadata (modification time and Unix permissions).
389    ///
390    /// This is the recommended method when writing files that should preserve their
391    /// original timestamps and permissions. Entries created with `start_entry()` have
392    /// zero timestamps and no permission bits.
393    ///
394    /// # Example
395    /// ```no_run
396    /// # use s_zip::{StreamingZipWriter, EntryOptions};
397    /// # fn example() -> Result<(), Box<dyn std::error::Error>> {
398    /// let mut writer = StreamingZipWriter::new("output.zip")?;
399    /// let opts = EntryOptions {
400    ///     mtime: Some(std::time::SystemTime::now()),
401    ///     unix_mode: Some(0o644),
402    /// };
403    /// writer.start_entry_with_options("readme.txt", opts)?;
404    /// writer.write_data(b"Hello")?;
405    /// writer.finish()?;
406    /// # Ok(())
407    /// # }
408    /// ```
409    pub fn start_entry_with_options(
410        &mut self,
411        name: &str,
412        options: crate::EntryOptions,
413    ) -> Result<()> {
414        self.start_entry_with_options_and_hint(name, options, None)
415    }
416
417    /// Start a new entry with size hint for optimized buffering
418    ///
419    /// Providing an accurate size hint can improve performance by 15-25% for large files.
420    /// The hint is used to optimize buffer allocation and flush thresholds.
421    ///
422    /// # Arguments
423    /// * `name` - The name/path of the entry in the ZIP
424    /// * `size_hint` - Optional uncompressed size hint in bytes
425    ///
426    /// # Example
427    /// ```no_run
428    /// # use s_zip::StreamingZipWriter;
429    /// # fn example() -> Result<(), Box<dyn std::error::Error>> {
430    /// let mut writer = StreamingZipWriter::new("output.zip")?;
431    ///
432    /// // For large files, provide size hint for better performance
433    /// writer.start_entry_with_hint("large_file.bin", Some(10_000_000))?;
434    /// # Ok(())
435    /// # }
436    /// ```
437    pub fn start_entry_with_hint(&mut self, name: &str, size_hint: Option<u64>) -> Result<()> {
438        self.start_entry_with_options_and_hint(name, crate::EntryOptions::default(), size_hint)
439    }
440
441    fn start_entry_with_options_and_hint(
442        &mut self,
443        name: &str,
444        options: crate::EntryOptions,
445        size_hint: Option<u64>,
446    ) -> Result<()> {
447        // Finish previous entry if any
448        self.finish_current_entry()?;
449
450        let local_header_offset = self.output.stream_position()?;
451        let compression_method = self.compression_method.to_zip_method();
452
453        // Check if encryption is enabled
454        #[cfg(feature = "encryption")]
455        let (encryptor, encryption_flag) = if let Some(ref password) = self.password {
456            let enc = AesEncryptor::new(password, self.encryption_strength)?;
457            (Some(enc), 0x01) // bit 0 set for encryption
458        } else {
459            (None, 0x00)
460        };
461
462        #[cfg(not(feature = "encryption"))]
463        let encryption_flag = 0x00;
464
465        // Write local file header with data descriptor flag (bit 3) + encryption flag (bit 0)
466        self.output.write_all(&[0x50, 0x4b, 0x03, 0x04])?; // signature
467        self.output.write_all(&[51, 0])?; // version needed (5.1 for AES)
468        self.output.write_all(&[8 | encryption_flag, 0])?; // general purpose bit flag
469        self.output.write_all(&compression_method.to_le_bytes())?; // compression method
470
471        // MS-DOS timestamp (time, date)
472        let (dos_time, dos_date) = options.msdos_datetime();
473        self.output.write_all(&dos_time.to_le_bytes())?;
474        self.output.write_all(&dos_date.to_le_bytes())?;
475
476        self.output.write_all(&0u32.to_le_bytes())?; // crc32 placeholder
477        self.output.write_all(&0u32.to_le_bytes())?; // compressed size placeholder
478        self.output.write_all(&0u32.to_le_bytes())?; // uncompressed size placeholder
479        self.output.write_all(&(name.len() as u16).to_le_bytes())?;
480
481        // Extra field: AES (11 bytes) + Unix permissions (15 bytes) if set
482        let unix_extra = options.unix_extra_field();
483        #[cfg(feature = "encryption")]
484        let extra_len = if encryptor.is_some() { 11 } else { 0 } + unix_extra.len();
485        #[cfg(not(feature = "encryption"))]
486        let extra_len = unix_extra.len();
487
488        self.output.write_all(&(extra_len as u16).to_le_bytes())?; // extra len
489        self.output.write_all(name.as_bytes())?;
490
491        // Write AES extra field if encryption is enabled
492        #[cfg(feature = "encryption")]
493        if let Some(ref enc) = encryptor {
494            // AES extra field header (0x9901)
495            // Format per WinZip AE-2 spec:
496            //   ID(2) + Length(2) + Version(2) + Vendor(2) + Strength(1) + ActualCompression(2) = 7 bytes data
497            self.output.write_all(&[0x01, 0x99])?; // WinZip AES encryption marker
498            self.output.write_all(&[7, 0])?; // data size (7 bytes)
499            self.output.write_all(&[2, 0])?; // AE-2 format version
500            self.output.write_all(&[0x41, 0x45])?; // vendor ID "AE"
501            self.output
502                .write_all(&[enc.strength().to_winzip_code() as u8])?; // strength (1 byte!)
503            self.output.write_all(&compression_method.to_le_bytes())?; // actual compression (2 bytes)
504
505            // Write salt and password verification
506            self.output.write_all(enc.salt())?;
507            self.output.write_all(enc.password_verify())?;
508        }
509
510        // Create encoder for this entry based on compression method
511        // Use adaptive buffer if size hint is provided
512        let encoder: Box<dyn CompressorWrite> = match self.compression_method {
513            CompressionMethod::Deflate => Box::new(DeflateCompressor {
514                encoder: DeflateEncoder::new(
515                    CompressedBuffer::with_size_hint(size_hint),
516                    Compression::new(self.compression_level),
517                ),
518            }),
519            #[cfg(feature = "zstd-support")]
520            CompressionMethod::Zstd => {
521                let mut encoder = zstd::Encoder::new(
522                    CompressedBuffer::with_size_hint(size_hint),
523                    self.compression_level as i32,
524                )?;
525                encoder.include_checksum(false)?; // ZIP uses CRC32, not zstd checksum
526                Box::new(ZstdCompressor { encoder })
527            }
528            CompressionMethod::Stored => {
529                // Stored method: no compression, pass through data
530                Box::new(StoredCompressor {
531                    buffer: CompressedBuffer::new(),
532                })
533            }
534        };
535
536        #[cfg_attr(not(feature = "encryption"), allow(unused_mut))]
537        let mut counter = CrcCounter::new();
538
539        // Account for salt and password verify bytes in compressed size for encrypted entries
540        #[cfg(feature = "encryption")]
541        if let Some(ref enc) = encryptor {
542            let encryption_overhead = (enc.salt().len() + 2) as u64; // salt + password_verify
543            counter.add_compressed(encryption_overhead);
544        }
545
546        self.current_entry = Some(CurrentEntry {
547            name: name.to_string(),
548            local_header_offset,
549            encoder,
550            counter,
551            compression_method,
552            #[cfg(feature = "encryption")]
553            encryptor,
554        });
555
556        Ok(())
557    }
558
559    /// Write uncompressed data to current entry (will be compressed and/or encrypted on-the-fly)
560    pub fn write_data(&mut self, data: &[u8]) -> Result<()> {
561        let entry = self
562            .current_entry
563            .as_mut()
564            .ok_or_else(|| SZipError::InvalidFormat("No entry started".to_string()))?;
565
566        // Update CRC and size with uncompressed data
567        entry.counter.update_uncompressed(data);
568
569        // For AES encryption: Update HMAC with plaintext BEFORE compression
570        #[cfg(feature = "encryption")]
571        if let Some(ref mut encryptor) = entry.encryptor {
572            encryptor.update_hmac(data);
573        }
574
575        // Write to encoder (compresses data into buffer)
576        entry.encoder.write_all(data)?;
577
578        // Flush encoder to ensure all data is in buffer
579        entry.encoder.flush()?;
580
581        // Check if buffer should be flushed to output
582        let buffer = entry.encoder.get_buffer_mut();
583        if buffer.should_flush() {
584            // Flush buffer to output to keep memory usage low
585            let compressed_data = buffer.take();
586
587            // Encrypt compressed data if encryption is enabled and password is set
588            #[cfg(feature = "encryption")]
589            let data_to_write = if let Some(ref mut encryptor) = entry.encryptor {
590                let mut data_to_encrypt = compressed_data;
591                encryptor.encrypt(&mut data_to_encrypt)?;
592                data_to_encrypt
593            } else {
594                compressed_data
595            };
596
597            #[cfg(not(feature = "encryption"))]
598            let data_to_write = compressed_data;
599
600            self.output.write_all(&data_to_write)?;
601            entry.counter.add_compressed(data_to_write.len() as u64);
602        }
603
604        Ok(())
605    }
606
607    /// Finish current entry and write data descriptor
608    fn finish_current_entry(&mut self) -> Result<()> {
609        if let Some(mut entry) = self.current_entry.take() {
610            // Finish compression and get remaining buffered data
611            let mut buffer = entry.encoder.finish_compression()?;
612
613            // Flush any remaining data from buffer to output
614            let remaining_data = buffer.take();
615            if !remaining_data.is_empty() {
616                // Encrypt remaining compressed data if encryption is enabled and password is set
617                #[cfg(feature = "encryption")]
618                let data_to_write = if let Some(ref mut encryptor) = entry.encryptor {
619                    let mut data_to_encrypt = remaining_data;
620                    encryptor.encrypt(&mut data_to_encrypt)?;
621                    data_to_encrypt
622                } else {
623                    remaining_data
624                };
625
626                #[cfg(not(feature = "encryption"))]
627                let data_to_write = remaining_data;
628
629                self.output.write_all(&data_to_write)?;
630                entry.counter.add_compressed(data_to_write.len() as u64);
631            }
632
633            // Write authentication code for AES encryption
634            #[cfg(feature = "encryption")]
635            let (encryption_strength_code, auth_code_size) =
636                if let Some(encryptor) = entry.encryptor {
637                    let strength_code = encryptor.strength().to_winzip_code();
638                    let auth_code = encryptor.finalize();
639                    self.output.write_all(&auth_code)?;
640                    (Some(strength_code), auth_code.len() as u64)
641                } else {
642                    (None, 0)
643                };
644
645            #[cfg(not(feature = "encryption"))]
646            let auth_code_size = 0u64;
647
648            let crc = entry.counter.finalize();
649            let compressed_size = entry.counter.compressed_count + auth_code_size;
650            let uncompressed_size = entry.counter.uncompressed_count;
651
652            // Write data descriptor
653            // signature
654            self.output.write_all(&[0x50, 0x4b, 0x07, 0x08])?;
655            self.output.write_all(&crc.to_le_bytes())?;
656            // If sizes exceed 32-bit, write 64-bit sizes (ZIP64 data descriptor)
657            if compressed_size > u32::MAX as u64 || uncompressed_size > u32::MAX as u64 {
658                self.output.write_all(&compressed_size.to_le_bytes())?;
659                self.output.write_all(&uncompressed_size.to_le_bytes())?;
660            } else {
661                self.output
662                    .write_all(&(compressed_size as u32).to_le_bytes())?;
663                self.output
664                    .write_all(&(uncompressed_size as u32).to_le_bytes())?;
665            }
666
667            // Save entry info for central directory
668            self.entries.push(ZipEntry {
669                name: entry.name,
670                local_header_offset: entry.local_header_offset,
671                crc32: crc,
672                compressed_size,
673                uncompressed_size,
674                compression_method: entry.compression_method,
675                #[cfg(feature = "encryption")]
676                encryption_strength: encryption_strength_code,
677            });
678        }
679        Ok(())
680    }
681
682    /// Finish ZIP file (write central directory and return the writer)
683    pub fn finish(mut self) -> Result<W> {
684        // Finish last entry
685        self.finish_current_entry()?;
686
687        let central_dir_offset = self.output.stream_position()?;
688
689        // Write central directory
690        for entry in &self.entries {
691            self.output.write_all(&[0x50, 0x4b, 0x01, 0x02])?; // central dir sig
692            self.output.write_all(&[20, 0])?; // version made by
693            self.output.write_all(&[20, 0])?; // version needed
694
695            // Set encryption flag (bit 0) if entry was encrypted
696            #[cfg(feature = "encryption")]
697            let flags = if entry.encryption_strength.is_some() {
698                0x08 | 0x01 // bit 3 (data descriptor) + bit 0 (encryption)
699            } else {
700                0x08 // bit 3 only (data descriptor)
701            };
702            #[cfg(not(feature = "encryption"))]
703            let flags = 0x08;
704
705            self.output.write_all(&[flags, 0])?; // general purpose bit flag
706            self.output
707                .write_all(&entry.compression_method.to_le_bytes())?; // compression method
708            self.output.write_all(&[0, 0, 0, 0])?; // mod time/date
709            self.output.write_all(&entry.crc32.to_le_bytes())?;
710
711            // Write sizes (32-bit placeholders or actual values)
712            if entry.compressed_size > u32::MAX as u64 {
713                self.output.write_all(&0xFFFFFFFFu32.to_le_bytes())?;
714            } else {
715                self.output
716                    .write_all(&(entry.compressed_size as u32).to_le_bytes())?;
717            }
718
719            if entry.uncompressed_size > u32::MAX as u64 {
720                self.output.write_all(&0xFFFFFFFFu32.to_le_bytes())?;
721            } else {
722                self.output
723                    .write_all(&(entry.uncompressed_size as u32).to_le_bytes())?;
724            }
725
726            self.output
727                .write_all(&(entry.name.len() as u16).to_le_bytes())?;
728
729            // Prepare extra fields
730            let mut extra_field: Vec<u8> = Vec::new();
731
732            // Add AES extra field if entry was encrypted
733            #[cfg(feature = "encryption")]
734            if let Some(strength_code) = entry.encryption_strength {
735                // AES extra field header (0x9901)
736                extra_field.extend_from_slice(&[0x01, 0x99]); // WinZip AES encryption marker
737                extra_field.extend_from_slice(&[7, 0]); // data size
738                extra_field.extend_from_slice(&[2, 0]); // AE-2 format
739                extra_field.extend_from_slice(&[0x41, 0x45]); // vendor ID "AE"
740                extra_field.extend_from_slice(&strength_code.to_le_bytes()); // strength
741                extra_field.extend_from_slice(&entry.compression_method.to_le_bytes());
742                // actual compression
743            }
744
745            // Add ZIP64 extra field if needed
746            if entry.uncompressed_size > u32::MAX as u64
747                || entry.compressed_size > u32::MAX as u64
748                || entry.local_header_offset > u32::MAX as u64
749            {
750                // ZIP64 extra header ID 0x0001
751                extra_field.extend_from_slice(&0x0001u16.to_le_bytes());
752                // data size: we'll include uncompressed (8) if needed, compressed (8) if needed, and offset (8) if needed
753                let mut data: Vec<u8> = Vec::new();
754                if entry.uncompressed_size > u32::MAX as u64 {
755                    data.extend_from_slice(&entry.uncompressed_size.to_le_bytes());
756                }
757                if entry.compressed_size > u32::MAX as u64 {
758                    data.extend_from_slice(&entry.compressed_size.to_le_bytes());
759                }
760                if entry.local_header_offset > u32::MAX as u64 {
761                    data.extend_from_slice(&entry.local_header_offset.to_le_bytes());
762                }
763                extra_field.extend_from_slice(&(data.len() as u16).to_le_bytes());
764                extra_field.extend_from_slice(&data);
765            }
766
767            self.output
768                .write_all(&(extra_field.len() as u16).to_le_bytes())?; // extra len
769            self.output.write_all(&0u16.to_le_bytes())?; // file comment len
770            self.output.write_all(&0u16.to_le_bytes())?; // disk number start
771            self.output.write_all(&0u16.to_le_bytes())?; // internal attrs
772            self.output.write_all(&0u32.to_le_bytes())?; // external attrs
773
774            // local header offset (32-bit or 0xFFFFFFFF)
775            if entry.local_header_offset > u32::MAX as u64 {
776                self.output.write_all(&0xFFFFFFFFu32.to_le_bytes())?;
777            } else {
778                self.output
779                    .write_all(&(entry.local_header_offset as u32).to_le_bytes())?;
780            }
781
782            self.output.write_all(entry.name.as_bytes())?;
783            if !extra_field.is_empty() {
784                self.output.write_all(&extra_field)?;
785            }
786        }
787
788        let central_dir_size = self.output.stream_position()? - central_dir_offset;
789
790        // Determine if we need ZIP64 EOCD
791        let need_zip64 = self.entries.len() > u16::MAX as usize
792            || central_dir_size > u32::MAX as u64
793            || central_dir_offset > u32::MAX as u64;
794
795        if need_zip64 {
796            // Write ZIP64 End of Central Directory Record
797            // signature
798            self.output.write_all(&[0x50, 0x4b, 0x06, 0x06])?; // 0x06064b50
799                                                               // size of zip64 eocd record (size of remaining fields)
800                                                               // We'll write fixed-size fields: version made by(2)+version needed(2)+disk numbers(4+4)+entries on disk(8)+total entries(8)+cd size(8)+cd offset(8)
801            let zip64_eocd_size: u64 = 44;
802            self.output.write_all(&zip64_eocd_size.to_le_bytes())?;
803            // version made by, version needed
804            self.output.write_all(&[20, 0])?;
805            self.output.write_all(&[20, 0])?;
806            // disk number, disk where central dir starts
807            self.output.write_all(&0u32.to_le_bytes())?;
808            self.output.write_all(&0u32.to_le_bytes())?;
809            // entries on this disk (8)
810            self.output
811                .write_all(&(self.entries.len() as u64).to_le_bytes())?;
812            // total entries (8)
813            self.output
814                .write_all(&(self.entries.len() as u64).to_le_bytes())?;
815            // central directory size (8)
816            self.output.write_all(&central_dir_size.to_le_bytes())?;
817            // central directory offset (8)
818            self.output.write_all(&central_dir_offset.to_le_bytes())?;
819
820            // Write ZIP64 EOCD locator
821            // signature
822            self.output.write_all(&[0x50, 0x4b, 0x06, 0x07])?; // 0x07064b50
823                                                               // disk with ZIP64 EOCD (4)
824            self.output.write_all(&0u32.to_le_bytes())?;
825            // relative offset of ZIP64 EOCD (8)
826            let zip64_eocd_pos = central_dir_offset + central_dir_size; // directly after central dir
827            self.output.write_all(&zip64_eocd_pos.to_le_bytes())?;
828            // total number of disks
829            self.output.write_all(&0u32.to_le_bytes())?;
830        }
831
832        // Write end of central directory (classic)
833        self.output.write_all(&[0x50, 0x4b, 0x05, 0x06])?;
834        self.output.write_all(&0u16.to_le_bytes())?; // disk number
835        self.output.write_all(&0u16.to_le_bytes())?; // disk with central dir
836
837        // number of entries (16-bit or 0xFFFF if ZIP64 used)
838        if self.entries.len() > u16::MAX as usize {
839            self.output.write_all(&0xFFFFu16.to_le_bytes())?;
840            self.output.write_all(&0xFFFFu16.to_le_bytes())?;
841        } else {
842            self.output
843                .write_all(&(self.entries.len() as u16).to_le_bytes())?;
844            self.output
845                .write_all(&(self.entries.len() as u16).to_le_bytes())?;
846        }
847
848        // central dir size and offset (32-bit or 0xFFFFFFFF)
849        if central_dir_size > u32::MAX as u64 {
850            self.output.write_all(&0xFFFFFFFFu32.to_le_bytes())?;
851        } else {
852            self.output
853                .write_all(&(central_dir_size as u32).to_le_bytes())?;
854        }
855
856        if central_dir_offset > u32::MAX as u64 {
857            self.output.write_all(&0xFFFFFFFFu32.to_le_bytes())?;
858        } else {
859            self.output
860                .write_all(&(central_dir_offset as u32).to_le_bytes())?;
861        }
862
863        self.output.write_all(&0u16.to_le_bytes())?; // comment len
864
865        self.output.flush()?;
866        Ok(self.output)
867    }
868}