1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum CompressionMethod {
26 Stored,
28 Deflate,
30 #[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
46struct 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)] encryption_strength: Option<u16>,
57}
58
59pub 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
111struct 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
163struct 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
193struct CompressedBuffer {
198 buffer: Vec<u8>,
199 flush_threshold: usize,
200}
201
202impl CompressedBuffer {
203 #[allow(dead_code)]
205 fn new() -> Self {
206 Self::with_size_hint(None)
207 }
208
209 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), Some(size) if size < 100_000 => (32 * 1024, 512 * 1024), Some(size) if size < 1_000_000 => (128 * 1024, 2 * 1024 * 1024), Some(size) if size < 10_000_000 => (256 * 1024, 4 * 1024 * 1024), _ => (512 * 1024, 8 * 1024 * 1024), };
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 pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
254 Self::with_compression(path, 6)
255 }
256
257 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 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 #[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 pub fn from_writer(writer: W) -> Result<Self> {
308 Self::from_writer_with_compression(writer, 6)
309 }
310
311 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 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 #[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 #[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 #[cfg(feature = "encryption")]
378 pub fn clear_password(&mut self) -> &mut Self {
379 self.password = None;
380 self
381 }
382
383 pub fn start_entry(&mut self, name: &str) -> Result<()> {
385 self.start_entry_with_hint(name, None)
386 }
387
388 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 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 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 #[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) } else {
459 (None, 0x00)
460 };
461
462 #[cfg(not(feature = "encryption"))]
463 let encryption_flag = 0x00;
464
465 self.output.write_all(&[0x50, 0x4b, 0x03, 0x04])?; self.output.write_all(&[51, 0])?; self.output.write_all(&[8 | encryption_flag, 0])?; self.output.write_all(&compression_method.to_le_bytes())?; 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())?; self.output.write_all(&0u32.to_le_bytes())?; self.output.write_all(&0u32.to_le_bytes())?; self.output.write_all(&(name.len() as u16).to_le_bytes())?;
480
481 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())?; self.output.write_all(name.as_bytes())?;
490
491 #[cfg(feature = "encryption")]
493 if let Some(ref enc) = encryptor {
494 self.output.write_all(&[0x01, 0x99])?; self.output.write_all(&[7, 0])?; self.output.write_all(&[2, 0])?; self.output.write_all(&[0x41, 0x45])?; self.output
502 .write_all(&[enc.strength().to_winzip_code() as u8])?; self.output.write_all(&compression_method.to_le_bytes())?; self.output.write_all(enc.salt())?;
507 self.output.write_all(enc.password_verify())?;
508 }
509
510 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)?; Box::new(ZstdCompressor { encoder })
527 }
528 CompressionMethod::Stored => {
529 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 #[cfg(feature = "encryption")]
541 if let Some(ref enc) = encryptor {
542 let encryption_overhead = (enc.salt().len() + 2) as u64; 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 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 entry.counter.update_uncompressed(data);
568
569 #[cfg(feature = "encryption")]
571 if let Some(ref mut encryptor) = entry.encryptor {
572 encryptor.update_hmac(data);
573 }
574
575 entry.encoder.write_all(data)?;
577
578 entry.encoder.flush()?;
580
581 let buffer = entry.encoder.get_buffer_mut();
583 if buffer.should_flush() {
584 let compressed_data = buffer.take();
586
587 #[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 fn finish_current_entry(&mut self) -> Result<()> {
609 if let Some(mut entry) = self.current_entry.take() {
610 let mut buffer = entry.encoder.finish_compression()?;
612
613 let remaining_data = buffer.take();
615 if !remaining_data.is_empty() {
616 #[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 #[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 self.output.write_all(&[0x50, 0x4b, 0x07, 0x08])?;
655 self.output.write_all(&crc.to_le_bytes())?;
656 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 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 pub fn finish(mut self) -> Result<W> {
684 self.finish_current_entry()?;
686
687 let central_dir_offset = self.output.stream_position()?;
688
689 for entry in &self.entries {
691 self.output.write_all(&[0x50, 0x4b, 0x01, 0x02])?; self.output.write_all(&[20, 0])?; self.output.write_all(&[20, 0])?; #[cfg(feature = "encryption")]
697 let flags = if entry.encryption_strength.is_some() {
698 0x08 | 0x01 } else {
700 0x08 };
702 #[cfg(not(feature = "encryption"))]
703 let flags = 0x08;
704
705 self.output.write_all(&[flags, 0])?; self.output
707 .write_all(&entry.compression_method.to_le_bytes())?; self.output.write_all(&[0, 0, 0, 0])?; self.output.write_all(&entry.crc32.to_le_bytes())?;
710
711 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 let mut extra_field: Vec<u8> = Vec::new();
731
732 #[cfg(feature = "encryption")]
734 if let Some(strength_code) = entry.encryption_strength {
735 extra_field.extend_from_slice(&[0x01, 0x99]); extra_field.extend_from_slice(&[7, 0]); extra_field.extend_from_slice(&[2, 0]); extra_field.extend_from_slice(&[0x41, 0x45]); extra_field.extend_from_slice(&strength_code.to_le_bytes()); extra_field.extend_from_slice(&entry.compression_method.to_le_bytes());
742 }
744
745 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 extra_field.extend_from_slice(&0x0001u16.to_le_bytes());
752 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())?; self.output.write_all(&0u16.to_le_bytes())?; self.output.write_all(&0u16.to_le_bytes())?; self.output.write_all(&0u16.to_le_bytes())?; self.output.write_all(&0u32.to_le_bytes())?; 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 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 self.output.write_all(&[0x50, 0x4b, 0x06, 0x06])?; let zip64_eocd_size: u64 = 44;
802 self.output.write_all(&zip64_eocd_size.to_le_bytes())?;
803 self.output.write_all(&[20, 0])?;
805 self.output.write_all(&[20, 0])?;
806 self.output.write_all(&0u32.to_le_bytes())?;
808 self.output.write_all(&0u32.to_le_bytes())?;
809 self.output
811 .write_all(&(self.entries.len() as u64).to_le_bytes())?;
812 self.output
814 .write_all(&(self.entries.len() as u64).to_le_bytes())?;
815 self.output.write_all(¢ral_dir_size.to_le_bytes())?;
817 self.output.write_all(¢ral_dir_offset.to_le_bytes())?;
819
820 self.output.write_all(&[0x50, 0x4b, 0x06, 0x07])?; self.output.write_all(&0u32.to_le_bytes())?;
825 let zip64_eocd_pos = central_dir_offset + central_dir_size; self.output.write_all(&zip64_eocd_pos.to_le_bytes())?;
828 self.output.write_all(&0u32.to_le_bytes())?;
830 }
831
832 self.output.write_all(&[0x50, 0x4b, 0x05, 0x06])?;
834 self.output.write_all(&0u16.to_le_bytes())?; self.output.write_all(&0u16.to_le_bytes())?; 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 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())?; self.output.flush()?;
866 Ok(self.output)
867 }
868}