Skip to main content

zzz_arc/formats/
zstd.rs

1//! zstd compression format implementation
2
3use crate::encryption::{
4    self, DecryptingReader, EncryptingWriter, ARGON2_SALT_LEN, DEFAULT_ENCRYPTION_CHUNK_SIZE,
5    ENCRYPTED_ZSTD_MAGIC,
6};
7use crate::filter::FileFilter;
8use crate::formats::{
9    tarball, ArchiveEntry, CompressionFormat, CompressionOptions, CompressionStats,
10    ExtractionOptions,
11};
12use crate::progress::{Progress, ProgressReader};
13use crate::Result;
14use anyhow::{anyhow, bail, Context};
15use std::fs::File;
16use std::io::{Read, Seek, SeekFrom, Write};
17use std::path::Path;
18use zstd::stream::raw::CParameter;
19
20pub struct ZstdFormat;
21
22impl CompressionFormat for ZstdFormat {
23    fn compress(
24        input_path: &Path,
25        output_path: &Path,
26        options: &CompressionOptions,
27        filter: &FileFilter,
28        progress: Option<&Progress>,
29    ) -> Result<CompressionStats> {
30        // calculate input size for progress and stats
31        let input_size = crate::utils::calculate_directory_size(
32            input_path,
33            filter,
34            options.follow_symlinks,
35            options.allow_symlink_escape,
36        )?;
37
38        // create output file
39        let mut underlying_file = File::create(output_path)
40            .with_context(|| format!("failed to create output file: {}", output_path.display()))?;
41
42        let mut key_material: Option<Vec<u8>> = None;
43
44        // Handle password-based encryption
45        if let Some(password) = &options.password {
46            if !password.is_empty() {
47                let (derived_key, salt) = encryption::derive_key(password, None)
48                    .context("Failed to derive encryption key for ZSTD compression")?;
49
50                // Write magic header and salt
51                underlying_file
52                    .write_all(ENCRYPTED_ZSTD_MAGIC)
53                    .context("Failed to write encryption magic header")?;
54                underlying_file
55                    .write_all(&salt)
56                    .context("Failed to write encryption salt")?;
57
58                key_material = Some(derived_key);
59            }
60        }
61
62        // Set up compression parameters
63        let zstd_level = if options.level == 0 { 3 } else { options.level };
64        let thread_count = if options.threads == 0 {
65            num_cpus::get() as u32
66        } else {
67            options.threads
68        };
69
70        let build_options = tarball::BuildOptions {
71            normalize_ownership: options.normalize_ownership,
72            apply_filter_to_single_file: true,
73            directory_slash: true,
74            set_mtime_for_single_file: true,
75        };
76
77        // Handle encrypted vs unencrypted compression differently
78        if let Some(key) = key_material {
79            // Encrypted compression pipeline
80            let encrypting_writer =
81                EncryptingWriter::new(underlying_file, &key, DEFAULT_ENCRYPTION_CHUNK_SIZE)
82                    .context("Failed to create EncryptingWriter for ZSTD")?;
83            let mut zstd_encoder = zstd::Encoder::new(encrypting_writer, zstd_level)
84                .context("Failed to create ZSTD encoder for encrypted stream")?;
85
86            // Configure threading
87            if thread_count > 1 {
88                let _ = zstd_encoder.set_parameter(CParameter::NbWorkers(thread_count));
89            }
90
91            let zstd_encoder = tarball::build_tarball(
92                zstd_encoder,
93                input_path,
94                options,
95                filter,
96                progress,
97                build_options,
98            )?;
99            let _inner = zstd_encoder.finish()?;
100        } else {
101            // Standard unencrypted compression pipeline
102            let mut zstd_encoder = zstd::Encoder::new(underlying_file, zstd_level)
103                .context("Failed to create ZSTD encoder for unencrypted stream")?;
104
105            // Configure threading
106            if thread_count > 1 {
107                let _ = zstd_encoder.set_parameter(CParameter::NbWorkers(thread_count));
108            }
109
110            let zstd_encoder = tarball::build_tarball(
111                zstd_encoder,
112                input_path,
113                options,
114                filter,
115                progress,
116                build_options,
117            )?;
118            let output_file = zstd_encoder.finish()?;
119            let _ = output_file; // The file is already written
120        }
121
122        let output_size = std::fs::metadata(output_path)?.len();
123
124        // finalize progress
125        if let Some(progress) = progress {
126            progress.update(input_size);
127        }
128
129        Ok(CompressionStats::new(input_size, output_size))
130    }
131
132    fn extract(
133        archive_path: &Path,
134        output_dir: &Path,
135        options: &ExtractionOptions,
136        progress: Option<&crate::progress::Progress>,
137    ) -> Result<()> {
138        // open archive file
139        let mut archive_file = File::open(archive_path)
140            .with_context(|| format!("failed to open archive file: {}", archive_path.display()))?;
141        let archive_size = std::fs::metadata(archive_path)
142            .with_context(|| {
143                format!(
144                    "failed to read archive metadata: {}",
145                    archive_path.display()
146                )
147            })?
148            .len();
149
150        // Check for encryption magic header
151        let mut magic_buffer = [0u8; ENCRYPTED_ZSTD_MAGIC.len()];
152        let bytes_read = archive_file
153            .read(&mut magic_buffer)
154            .context("Failed to read initial bytes from archive for encryption check")?;
155
156        // Determine if this is an encrypted archive and set up the input stream
157        let (input_stream, bytes_offset): (Box<dyn Read>, u64) = if bytes_read
158            == ENCRYPTED_ZSTD_MAGIC.len()
159            && magic_buffer == *ENCRYPTED_ZSTD_MAGIC
160        {
161            // This is an encrypted archive
162            let password = options.password.as_deref().ok_or_else(|| {
163                anyhow!(
164                    "Encrypted archive '{}' requires a password.",
165                    archive_path.display()
166                )
167            })?;
168
169            if password.is_empty() {
170                bail!(
171                    "Password cannot be empty for encrypted archive '{}'.",
172                    archive_path.display()
173                );
174            }
175
176            // Read the salt
177            let mut salt = vec![0u8; ARGON2_SALT_LEN];
178            archive_file
179                .read_exact(&mut salt)
180                .context("Failed to read salt from encrypted archive")?;
181
182            // Derive the decryption key
183            let (derived_key, _used_salt) = encryption::derive_key(password, Some(&salt))
184                .context("Failed to derive decryption key")?;
185
186            // Create decrypting reader
187            let decrypting_reader =
188                DecryptingReader::new(ProgressReader::new(archive_file, progress), &derived_key)
189                    .context("Failed to create DecryptingReader for ZSTD")?;
190
191            (
192                Box::new(decrypting_reader),
193                (ENCRYPTED_ZSTD_MAGIC.len() + ARGON2_SALT_LEN) as u64,
194            )
195        } else {
196            // This is a standard (unencrypted) archive
197            archive_file
198                .seek(SeekFrom::Start(0))
199                .context("Failed to rewind archive file for standard processing")?;
200
201            if options.password.is_some()
202                && !options.password.as_deref().unwrap_or_default().is_empty()
203            {
204                eprintln!(
205                    "warning: Password provided, but archive '{}' does not appear to be in the expected encrypted format. Attempting standard extraction.",
206                    archive_path.display()
207                );
208            }
209
210            (Box::new(ProgressReader::new(archive_file, progress)), 0)
211        };
212
213        if let Some(progress) = progress {
214            progress.set_length(archive_size.saturating_sub(bytes_offset));
215        }
216
217        // create zstd decoder with the appropriate input stream
218        let decoder = zstd::Decoder::new(input_stream).with_context(|| {
219            format!(
220                "failed to create zstd decoder for: {}",
221                archive_path.display()
222            )
223        })?;
224
225        tarball::extract_tarball(decoder, output_dir, options, progress)
226    }
227
228    fn list(archive_path: &Path) -> Result<Vec<ArchiveEntry>> {
229        // open archive file
230        let mut archive_file = File::open(archive_path)
231            .with_context(|| format!("failed to open archive file: {}", archive_path.display()))?;
232
233        // Check for encryption magic header
234        let mut magic_buffer = [0u8; ENCRYPTED_ZSTD_MAGIC.len()];
235        let bytes_read = archive_file
236            .read(&mut magic_buffer)
237            .context("Failed to read initial bytes from archive for encryption check")?;
238
239        // If this is an encrypted archive, we can't list it without a password
240        if bytes_read == ENCRYPTED_ZSTD_MAGIC.len() && magic_buffer == *ENCRYPTED_ZSTD_MAGIC {
241            return Err(anyhow!(
242                "Cannot list encrypted ZSTD archive '{}' - password required. Use the extract command with --password to access contents.",
243                archive_path.display()
244            ));
245        }
246
247        // This is a standard archive, proceed normally
248        archive_file
249            .seek(SeekFrom::Start(0))
250            .context("Failed to rewind archive file for standard processing")?;
251
252        // create zstd decoder
253        let decoder = zstd::Decoder::new(archive_file).with_context(|| {
254            format!(
255                "failed to create zstd decoder for: {}",
256                archive_path.display()
257            )
258        })?;
259
260        tarball::list_tarball(decoder)
261    }
262
263    fn extension() -> &'static str {
264        "zst"
265    }
266
267    fn test_integrity(archive_path: &Path) -> Result<()> {
268        // open archive file
269        let mut archive_file = File::open(archive_path)
270            .with_context(|| format!("failed to open archive file: {}", archive_path.display()))?;
271
272        // Check for encryption magic header
273        let mut magic_buffer = [0u8; ENCRYPTED_ZSTD_MAGIC.len()];
274        let bytes_read = archive_file
275            .read(&mut magic_buffer)
276            .context("Failed to read initial bytes from archive for encryption check")?;
277
278        // If this is an encrypted archive, we can only verify the header format
279        if bytes_read == ENCRYPTED_ZSTD_MAGIC.len() && magic_buffer == *ENCRYPTED_ZSTD_MAGIC {
280            // For encrypted archives, we can check if the salt is readable
281            let mut salt = vec![0u8; ARGON2_SALT_LEN];
282            archive_file
283                .read_exact(&mut salt)
284                .context("Failed to read salt from encrypted archive")?;
285
286            // If we got here, the header format is valid
287            // Full integrity testing would require a password, but basic format is OK
288            return Ok(());
289        }
290
291        // This is a standard archive, proceed with full integrity testing
292        archive_file
293            .seek(SeekFrom::Start(0))
294            .context("Failed to rewind archive file for standard processing")?;
295
296        // Check if this is a tar.zst or single .zst file
297        if archive_path
298            .extension()
299            .is_some_and(|ext| ext == "tzst" || ext == "tar.zst")
300            || archive_path
301                .file_name()
302                .is_some_and(|name| name.to_string_lossy().ends_with(".tar.zst"))
303        {
304            // This is a tar.zst archive - test by reading all tar entries
305            let zstd_decoder = zstd::stream::read::Decoder::new(archive_file)?;
306            let mut archive = tar::Archive::new(zstd_decoder);
307            for entry in archive.entries()? {
308                let _entry = entry?; // This will fail if data is corrupted
309            }
310        } else {
311            // Single .zst file - test by decompressing fully
312            let mut zstd_decoder = zstd::stream::read::Decoder::new(archive_file)?;
313            let mut buffer = Vec::new();
314            zstd_decoder.read_to_end(&mut buffer)?;
315        }
316
317        Ok(())
318    }
319}