zzz_arc/
encryption.rs

1//! Encryption module for authenticated encryption of archives
2//!
3//! This module provides authenticated encryption for ZSTD archives using:
4//! - Argon2id for key derivation from passwords
5//! - AES-256-GCM for streaming authenticated encryption
6//! - Chunked processing for memory efficiency
7
8use crate::Result;
9use anyhow::{anyhow, bail, Context};
10use argon2::Argon2;
11use rand::RngCore;
12use std::io::{self, ErrorKind, Read};
13
14// AES-GCM imports
15use aes_gcm::{
16    aead::{Aead, KeyInit, OsRng as AeadOsRng},
17    Aes256Gcm, Nonce,
18};
19
20// AES-256-GCM constants
21pub const AES_KEY_SIZE: usize = 32; // 256 bits
22pub const NONCE_SIZE: usize = 12; // 96 bits, standard for GCM
23pub const TAG_SIZE: usize = 16; // 128 bits, standard for GCM
24
25// Argon2id parameters
26pub const ARGON2_SALT_LEN: usize = 16; // Bytes
27
28// Default Argon2 parameters (secure but reasonable performance)
29pub const ARGON2_TIME_COST: u32 = 3;
30pub const ARGON2_MEM_COST: u32 = 65536; // 64 MiB
31pub const ARGON2_LANES: u32 = 4;
32
33// Magic string to identify encrypted ZSTD files
34pub const ENCRYPTED_ZSTD_MAGIC: &[u8; 11] = b"ZSTDECRYPT1";
35
36// Size of the header for encrypted files = magic + salt_len
37pub const ENCRYPTION_HEADER_SIZE: usize = ENCRYPTED_ZSTD_MAGIC.len() + ARGON2_SALT_LEN;
38
39// Default chunk size for encryption (plaintext)
40pub const DEFAULT_ENCRYPTION_CHUNK_SIZE: usize = 64 * 1024;
41// Guardrail for corrupted ciphertext length fields (ciphertext + tag).
42pub const MAX_ENCRYPTED_CHUNK_LEN: usize = DEFAULT_ENCRYPTION_CHUNK_SIZE * 16 + TAG_SIZE;
43
44/// Derive an AES-256 key from a password using Argon2id
45///
46/// Returns (derived_key, salt) where salt is either the provided salt or a new random salt
47pub fn derive_key(password: &str, salt_opt: Option<&[u8]>) -> Result<(Vec<u8>, Vec<u8>)> {
48    let salt = match salt_opt {
49        Some(s) => {
50            if s.len() != ARGON2_SALT_LEN {
51                return Err(anyhow!(
52                    "Provided salt has incorrect length. Expected {}, got {}.",
53                    ARGON2_SALT_LEN,
54                    s.len()
55                ));
56            }
57            s.to_vec()
58        }
59        None => {
60            let mut new_salt = vec![0u8; ARGON2_SALT_LEN];
61            rand::rngs::OsRng
62                .try_fill_bytes(&mut new_salt)
63                .context("Failed to generate random salt for Argon2")?;
64            new_salt
65        }
66    };
67
68    // Use Argon2id with specified parameters
69    use argon2::{Algorithm, Params, Version};
70    let params = Params::new(
71        ARGON2_MEM_COST,
72        ARGON2_TIME_COST,
73        ARGON2_LANES,
74        Some(AES_KEY_SIZE),
75    )
76    .map_err(|e| anyhow!("Failed to create Argon2 parameters: {e}"))?;
77    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
78
79    let mut derived_key = vec![0u8; AES_KEY_SIZE];
80    argon2
81        .hash_password_into(password.as_bytes(), &salt, &mut derived_key)
82        .map_err(|e| anyhow!("Argon2 key derivation failed: {e}"))?;
83
84    Ok((derived_key, salt))
85}
86
87/// A writer that encrypts data in chunks using AES-256-GCM before writing to the underlying writer
88///
89/// Each chunk is encrypted with a random nonce and written as:
90/// [nonce: 12 bytes][ciphertext_length: 4 bytes BE][ciphertext_with_tag: variable]
91pub struct EncryptingWriter<W: std::io::Write> {
92    inner: W,
93    cipher: Aes256Gcm,
94    buffer: Vec<u8>,
95    chunk_size: usize,
96}
97
98impl<W: std::io::Write> EncryptingWriter<W> {
99    pub fn new(inner: W, key: &[u8], chunk_size: usize) -> Result<Self> {
100        if key.len() != AES_KEY_SIZE {
101            return Err(anyhow!(
102                "Invalid key size for EncryptingWriter. Expected {}, got {}",
103                AES_KEY_SIZE,
104                key.len()
105            ));
106        }
107        let cipher = Aes256Gcm::new_from_slice(key)
108            .map_err(|e| anyhow!("Failed to initialize AES-GCM cipher: {e}"))?;
109        Ok(Self {
110            inner,
111            cipher,
112            buffer: Vec::with_capacity(chunk_size),
113            chunk_size,
114        })
115    }
116
117    fn encrypt_and_write_chunk(&mut self) -> Result<()> {
118        if self.buffer.is_empty() {
119            return Ok(());
120        }
121
122        // Generate random nonce for this chunk
123        let mut nonce_bytes = [0u8; NONCE_SIZE];
124        AeadOsRng
125            .try_fill_bytes(&mut nonce_bytes)
126            .context("Failed to generate nonce for encryption")?;
127        let nonce = Nonce::from_slice(&nonce_bytes);
128
129        // Encrypt the chunk (includes authentication tag)
130        let ciphertext_with_tag = self
131            .cipher
132            .encrypt(nonce, self.buffer.as_slice())
133            .map_err(|e| anyhow!("AES-GCM encryption failed: {e}"))?;
134
135        // Write: nonce + length + ciphertext_with_tag
136        self.inner
137            .write_all(&nonce_bytes)
138            .context("Failed to write nonce to inner writer")?;
139
140        let len_bytes = (ciphertext_with_tag.len() as u32).to_be_bytes();
141        self.inner
142            .write_all(&len_bytes)
143            .context("Failed to write ciphertext length to inner writer")?;
144
145        self.inner
146            .write_all(&ciphertext_with_tag)
147            .context("Failed to write ciphertext with tag to inner writer")?;
148
149        self.buffer.clear();
150        Ok(())
151    }
152}
153
154impl<W: std::io::Write> std::io::Write for EncryptingWriter<W> {
155    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
156        let mut bytes_written_total = 0;
157        let mut input_remaining = buf;
158
159        while !input_remaining.is_empty() {
160            let space_in_buffer = self.chunk_size - self.buffer.len();
161            let bytes_to_buffer = std::cmp::min(space_in_buffer, input_remaining.len());
162
163            self.buffer
164                .extend_from_slice(&input_remaining[..bytes_to_buffer]);
165            input_remaining = &input_remaining[bytes_to_buffer..];
166            bytes_written_total += bytes_to_buffer;
167
168            // If buffer is full, encrypt and write the chunk
169            if self.buffer.len() == self.chunk_size {
170                self.encrypt_and_write_chunk()
171                    .map_err(|e| std::io::Error::other(e.to_string()))?;
172            }
173        }
174        Ok(bytes_written_total)
175    }
176
177    fn flush(&mut self) -> std::io::Result<()> {
178        // Encrypt and write any remaining data in buffer
179        if !self.buffer.is_empty() {
180            self.encrypt_and_write_chunk()
181                .map_err(|e| std::io::Error::other(e.to_string()))?;
182        }
183        self.inner.flush()
184    }
185}
186
187impl<W: std::io::Write> Drop for EncryptingWriter<W> {
188    fn drop(&mut self) {
189        if !self.buffer.is_empty() {
190            if let Err(e) = self.encrypt_and_write_chunk() {
191                eprintln!("Error encrypting remaining chunk during drop: {e}");
192            }
193        }
194    }
195}
196
197/// A reader that decrypts chunked AES-256-GCM encrypted data
198///
199/// Reads chunks in format: [nonce: 12 bytes][ciphertext_length: 4 bytes BE][ciphertext_with_tag: variable]
200pub struct DecryptingReader<R: Read> {
201    inner: R,
202    cipher: Aes256Gcm,
203    buffer: Vec<u8>,
204    buffer_pos: usize,
205    eof_reached: bool,
206}
207
208impl<R: Read> DecryptingReader<R> {
209    pub fn new(inner: R, key: &[u8]) -> Result<Self> {
210        if key.len() != AES_KEY_SIZE {
211            return Err(anyhow!(
212                "Invalid key size for DecryptingReader. Expected {}, got {}",
213                AES_KEY_SIZE,
214                key.len()
215            ));
216        }
217        let cipher = Aes256Gcm::new_from_slice(key)
218            .map_err(|e| anyhow!("Failed to initialize AES-GCM cipher: {e}"))?;
219        Ok(Self {
220            inner,
221            cipher,
222            buffer: Vec::new(),
223            buffer_pos: 0,
224            eof_reached: false,
225        })
226    }
227
228    fn read_and_decrypt_chunk(&mut self) -> Result<bool> {
229        self.buffer.clear();
230        self.buffer_pos = 0;
231
232        // Read nonce
233        let mut nonce_bytes = [0u8; NONCE_SIZE];
234        match self.inner.read_exact(&mut nonce_bytes) {
235            Ok(()) => {}
236            Err(e) if e.kind() == ErrorKind::UnexpectedEof => {
237                self.eof_reached = true;
238                return Ok(false);
239            }
240            Err(e) => return Err(e).context("Failed to read nonce from inner reader"),
241        }
242        let nonce = Nonce::from_slice(&nonce_bytes);
243
244        // Read ciphertext length
245        let mut len_bytes = [0u8; 4];
246        match self.inner.read_exact(&mut len_bytes) {
247            Ok(()) => {}
248            Err(e) if e.kind() == ErrorKind::UnexpectedEof => {
249                return Err(e)
250                    .context("Failed to read ciphertext length (unexpected EOF after nonce)");
251            }
252            Err(e) => return Err(e).context("Failed to read ciphertext length from inner reader"),
253        }
254        let ct_with_tag_len = u32::from_be_bytes(len_bytes) as usize;
255
256        if ct_with_tag_len < TAG_SIZE {
257            bail!("Invalid ciphertext length: {ct_with_tag_len} (must be at least TAG_SIZE {TAG_SIZE})");
258        }
259        if ct_with_tag_len > MAX_ENCRYPTED_CHUNK_LEN {
260            bail!("Invalid ciphertext length: {ct_with_tag_len} (exceeds max {MAX_ENCRYPTED_CHUNK_LEN})");
261        }
262
263        // Read ciphertext with tag
264        let mut ciphertext_with_tag = vec![0u8; ct_with_tag_len];
265        self.inner
266            .read_exact(&mut ciphertext_with_tag)
267            .context("Failed to read ciphertext with tag from inner reader")?;
268
269        // Decrypt and verify
270        let plaintext = self
271            .cipher
272            .decrypt(nonce, ciphertext_with_tag.as_slice())
273            .map_err(|e| anyhow!("AES-GCM decryption failed (data integrity or key error): {e}"))?;
274
275        self.buffer = plaintext;
276        Ok(true)
277    }
278}
279
280impl<R: Read> Read for DecryptingReader<R> {
281    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
282        // If we've consumed all buffered data, try to read and decrypt the next chunk
283        if self.buffer_pos == self.buffer.len() {
284            if self.eof_reached {
285                return Ok(0);
286            }
287            match self.read_and_decrypt_chunk() {
288                Ok(true) => {}
289                Ok(false) => {
290                    self.eof_reached = true;
291                    return Ok(0);
292                }
293                Err(e) => {
294                    let io_err_kind = if e.to_string().contains("AES-GCM decryption failed") {
295                        ErrorKind::InvalidData
296                    } else {
297                        ErrorKind::Other
298                    };
299                    return Err(io::Error::new(
300                        io_err_kind,
301                        format!("Decryption error: {e}"),
302                    ));
303                }
304            }
305        }
306
307        let bytes_to_read = std::cmp::min(buf.len(), self.buffer.len() - self.buffer_pos);
308
309        if bytes_to_read == 0 && self.buffer_pos == self.buffer.len() && self.eof_reached {
310            return Ok(0);
311        }
312
313        buf[..bytes_to_read]
314            .copy_from_slice(&self.buffer[self.buffer_pos..self.buffer_pos + bytes_to_read]);
315        self.buffer_pos += bytes_to_read;
316
317        Ok(bytes_to_read)
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324    use std::io::{Cursor, Read, Write};
325
326    #[test]
327    fn test_derive_key_new_salt() {
328        let password = "test_password";
329        let (key1, salt1) = derive_key(password, None).unwrap();
330        assert_eq!(key1.len(), AES_KEY_SIZE);
331        assert_eq!(salt1.len(), ARGON2_SALT_LEN);
332
333        let (key2, salt2) = derive_key(password, None).unwrap();
334        assert_eq!(key2.len(), AES_KEY_SIZE);
335        assert_eq!(salt2.len(), ARGON2_SALT_LEN);
336        assert_ne!(salt1, salt2);
337        assert_ne!(key1, key2);
338    }
339
340    #[test]
341    fn test_derive_key_with_provided_salt() {
342        let password = "test_password";
343        let mut salt = vec![0u8; ARGON2_SALT_LEN];
344        rand::rngs::OsRng.try_fill_bytes(&mut salt).unwrap();
345
346        let (key1, used_salt1) = derive_key(password, Some(&salt)).unwrap();
347        assert_eq!(key1.len(), AES_KEY_SIZE);
348        assert_eq!(used_salt1, salt);
349
350        let (key2, used_salt2) = derive_key(password, Some(&salt)).unwrap();
351        assert_eq!(key2, key1);
352        assert_eq!(used_salt2, salt);
353    }
354
355    #[test]
356    fn test_derive_key_invalid_salt_length() {
357        let password = "test_password";
358        let invalid_salt = vec![0u8; ARGON2_SALT_LEN - 1];
359        let result = derive_key(password, Some(&invalid_salt));
360        assert!(result.is_err());
361    }
362
363    #[test]
364    fn test_encrypting_writer_simple() -> Result<()> {
365        let key = [0u8; AES_KEY_SIZE];
366        let mut output_buffer = Vec::new();
367        let chunk_size = 16;
368
369        {
370            let mut writer = EncryptingWriter::new(&mut output_buffer, &key, chunk_size)?;
371            writer.write_all(b"test data that is longer than one chunk")?; // 38 bytes
372            writer.flush()?;
373        }
374        // Should have encrypted data (nonce + length + ciphertext_with_tag for each chunk)
375        assert!(!output_buffer.is_empty());
376        Ok(())
377    }
378
379    #[test]
380    fn test_decrypting_reader_simple_roundtrip() -> Result<()> {
381        let key = [3u8; AES_KEY_SIZE];
382        let original_data = b"Hello, world! This is a test of the decryption system.";
383        let chunk_size = 16;
384
385        let mut encrypted_data = Vec::new();
386        {
387            let mut enc_writer = EncryptingWriter::new(&mut encrypted_data, &key, chunk_size)?;
388            enc_writer.write_all(original_data)?;
389            enc_writer.flush()?;
390        }
391
392        assert!(!encrypted_data.is_empty());
393
394        let mut dec_reader = DecryptingReader::new(Cursor::new(&encrypted_data), &key)?;
395        let mut decrypted_data = Vec::new();
396        dec_reader.read_to_end(&mut decrypted_data)?;
397
398        assert_eq!(original_data.to_vec(), decrypted_data);
399        Ok(())
400    }
401
402    #[test]
403    fn test_decrypting_reader_wrong_key() -> Result<()> {
404        let key_encrypt = [7u8; AES_KEY_SIZE];
405        let key_decrypt = [8u8; AES_KEY_SIZE];
406        let original_data = b"secret message";
407        let chunk_size = 128;
408
409        let mut encrypted_data = Vec::new();
410        {
411            let mut enc_writer =
412                EncryptingWriter::new(&mut encrypted_data, &key_encrypt, chunk_size)?;
413            enc_writer.write_all(original_data)?;
414        }
415
416        let mut dec_reader = DecryptingReader::new(Cursor::new(&encrypted_data), &key_decrypt)?;
417        let mut decrypted_data = Vec::new();
418        let result = dec_reader.read_to_end(&mut decrypted_data);
419
420        assert!(result.is_err());
421        if let Err(e) = result {
422            assert!(
423                e.to_string().contains("AES-GCM decryption failed")
424                    || e.to_string().contains("Decryption error")
425            );
426        }
427        Ok(())
428    }
429
430    #[test]
431    fn test_decrypting_reader_rejects_large_chunk_len() -> Result<()> {
432        let key = [0u8; AES_KEY_SIZE];
433        let mut encrypted = Vec::new();
434        encrypted.extend_from_slice(&[0u8; NONCE_SIZE]);
435        let oversize = (MAX_ENCRYPTED_CHUNK_LEN + 1) as u32;
436        encrypted.extend_from_slice(&oversize.to_be_bytes());
437
438        let mut reader = DecryptingReader::new(Cursor::new(encrypted), &key)?;
439        let mut out = Vec::new();
440        let err = reader.read_to_end(&mut out).unwrap_err();
441        assert!(err.to_string().contains("Invalid ciphertext length"));
442        Ok(())
443    }
444}
zzz_arc/encryption.rs

zzz_arc/
encryption.rs
