Skip to main content

age_plugin_argon2/
encrypt.rs

1//! Custom age-format writer that accepts a known FileKey.
2//!
3//! Since `age::Encryptor` generates a random FileKey internally and doesn't
4//! expose it, we need a custom encrypt function to reuse a cached FileKey.
5//! This produces spec-compliant age v1 binary files.
6
7use std::io::Write;
8
9use age::Recipient;
10use age_core::format::FileKey;
11use base64::Engine;
12use chacha20poly1305::aead::{AeadInPlace, KeyInit};
13use chacha20poly1305::ChaCha20Poly1305;
14use hmac::Mac;
15use rand::rngs::OsRng;
16use rand::RngCore;
17
18/// Encrypt plaintext using a known FileKey + recipient.
19///
20/// Produces a standard age v1 binary format file that can be decrypted
21/// by any age-compatible tool with the matching identity.
22///
23/// # Safety
24///
25/// Reusing the same FileKey + salt across writes is safe because:
26/// - The stanza wrapping is deterministic (same inputs, same output)
27/// - The STREAM layer gets a fresh random 16-byte nonce each write
28pub fn encrypt_with_file_key(
29    file_key: &[u8; 16],
30    recipient: &impl Recipient,
31    plaintext: &[u8],
32) -> Result<Vec<u8>, EncryptWithFileKeyError> {
33    let fk = FileKey::new(Box::new(*file_key));
34
35    // 1. Wrap the file key with the recipient
36    let (stanzas, _labels) = recipient
37        .wrap_file_key(&fk)
38        .map_err(|e| EncryptWithFileKeyError::Wrap(e.to_string()))?;
39
40    // 2. Build the header (everything covered by the MAC)
41    // Per age spec: MAC covers from "age-encryption.org/v1\n" through "---" inclusive
42    let mut header = Vec::new();
43    header.extend_from_slice(b"age-encryption.org/v1\n");
44
45    for stanza in &stanzas {
46        // Write stanza: -> tag arg1 arg2 ...
47        write!(header, "-> {}", stanza.tag)
48            .map_err(|e| EncryptWithFileKeyError::Io(e.to_string()))?;
49        for arg in &stanza.args {
50            write!(header, " {}", arg).map_err(|e| EncryptWithFileKeyError::Io(e.to_string()))?;
51        }
52        header.push(b'\n');
53
54        // Write body lines (base64, 64 chars per line)
55        let encoded = base64::engine::general_purpose::STANDARD_NO_PAD.encode(&stanza.body);
56        let mut remaining = encoded.as_str();
57        loop {
58            if remaining.len() >= 64 {
59                let (line, rest) = remaining.split_at(64);
60                header.extend_from_slice(line.as_bytes());
61                header.push(b'\n');
62                remaining = rest;
63            } else {
64                // Final short line (may be empty if encoded was multiple of 64)
65                header.extend_from_slice(remaining.as_bytes());
66                header.push(b'\n');
67                break;
68            }
69        }
70    }
71
72    // Append "---" to the MAC input (per spec: MAC covers through "---" inclusive)
73    header.extend_from_slice(b"---");
74
75    // 3. Compute header MAC
76    // MAC key = HKDF-SHA256(ikm=file_key, salt="", info="header")
77    let mac_key = age_core::primitives::hkdf(&[], b"header", file_key);
78
79    let mut mac = <hmac::Hmac<sha2::Sha256> as Mac>::new_from_slice(&mac_key)
80        .map_err(|_| EncryptWithFileKeyError::Crypto("invalid MAC key length".to_string()))?;
81    mac.update(&header);
82
83    let mac_result = mac.finalize().into_bytes();
84    let mac_b64 = base64::engine::general_purpose::STANDARD_NO_PAD.encode(mac_result);
85
86    // 4. Build complete output
87    // The actual file has: header_without_dashes + "--- " + mac_b64 + "\n" + payload
88    let mut output = Vec::new();
89    // Write everything except the trailing "---" that we added for MAC computation
90    output.extend_from_slice(&header[..header.len() - 3]);
91    // Write the MAC line: "--- <mac>\n"
92    writeln!(output, "--- {}", mac_b64).map_err(|e| EncryptWithFileKeyError::Io(e.to_string()))?;
93
94    // 5. Generate 16-byte random nonce for the payload
95    let mut nonce = [0u8; 16];
96    OsRng.fill_bytes(&mut nonce);
97    output.extend_from_slice(&nonce);
98
99    // 6. Derive payload key: HKDF-SHA256(salt=nonce, info="payload", ikm=file_key)
100    let payload_key = age_core::primitives::hkdf(&nonce, b"payload", file_key);
101
102    // 7. STREAM encryption — single final chunk
103    // age uses a custom STREAM construction:
104    // - ChaCha20-Poly1305 with a 12-byte nonce
105    // - Nonce = 11 bytes counter (big-endian) + 1 byte flag
106    // - For a single final chunk: counter=0, flag=0x01 (final)
107    let mut stream_nonce = [0u8; 12];
108    stream_nonce[11] = 0x01; // last_chunk flag
109
110    let cipher = ChaCha20Poly1305::new_from_slice(&payload_key)
111        .map_err(|e| EncryptWithFileKeyError::Crypto(e.to_string()))?;
112
113    let mut ciphertext = plaintext.to_vec();
114    cipher
115        .encrypt_in_place(
116            chacha20poly1305::Nonce::from_slice(&stream_nonce),
117            &[],
118            &mut ciphertext,
119        )
120        .map_err(|e| EncryptWithFileKeyError::Crypto(e.to_string()))?;
121
122    output.extend_from_slice(&ciphertext);
123
124    Ok(output)
125}
126
127/// Error returned by [`encrypt_with_file_key`].
128#[derive(Debug, thiserror::Error)]
129pub enum EncryptWithFileKeyError {
130    /// The recipient failed to wrap the file key.
131    #[error("failed to wrap file key: {0}")]
132    Wrap(String),
133    /// An I/O error occurred while building the output.
134    #[error("I/O error: {0}")]
135    Io(String),
136    /// A cryptographic operation failed.
137    #[error("cryptographic error: {0}")]
138    Crypto(String),
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::cached::{CachedIdentity, CachedRecipient};
145    use crate::identity::Argon2idIdentity;
146    use crate::params::Argon2Params;
147    use crate::recipient::Argon2idRecipient;
148
149    fn fast_params() -> Argon2Params {
150        Argon2Params::new(256, 1, 1).unwrap()
151    }
152
153    #[test]
154    fn test_encrypt_decrypt_with_full_kdf() {
155        let passphrase = b"test-password";
156        let params = fast_params();
157
158        // Encrypt with full KDF recipient
159        let recipient = Argon2idRecipient::new(passphrase, params);
160        let file_key = [42u8; 16];
161        let plaintext = b"hello, world!";
162
163        let ciphertext = encrypt_with_file_key(&file_key, &recipient, plaintext).unwrap();
164
165        // Decrypt with full KDF identity
166        let identity = Argon2idIdentity::new(passphrase);
167        let decrypted = age::decrypt(&identity, &ciphertext).unwrap();
168
169        assert_eq!(decrypted, plaintext);
170    }
171
172    #[test]
173    fn test_encrypt_decrypt_with_cached_material() {
174        let passphrase = b"test-password";
175        let params = fast_params();
176
177        // First: normal encrypt+decrypt to capture material
178        let recipient = Argon2idRecipient::new(passphrase, params);
179        let file_key_bytes = [42u8; 16];
180        let plaintext = b"sensitive data here";
181
182        let ciphertext = encrypt_with_file_key(&file_key_bytes, &recipient, plaintext).unwrap();
183
184        let identity = Argon2idIdentity::new(passphrase);
185        let decrypted = age::decrypt(&identity, &ciphertext).unwrap();
186        assert_eq!(decrypted, plaintext);
187
188        let material = identity.captured_material().unwrap();
189
190        // Now: re-encrypt with cached recipient + custom writer
191        let cached_recipient = CachedRecipient::new(&material);
192        let new_plaintext = b"updated data";
193
194        let new_ciphertext =
195            encrypt_with_file_key(&material.file_key, &cached_recipient, new_plaintext).unwrap();
196
197        // Decrypt with cached identity
198        let cached_identity = CachedIdentity::new(&material);
199        let result = age::decrypt(&cached_identity, &new_ciphertext).unwrap();
200        assert_eq!(result, new_plaintext);
201    }
202
203    #[test]
204    fn test_output_is_valid_age_format() {
205        let passphrase = b"test";
206        let params = fast_params();
207        let recipient = Argon2idRecipient::new(passphrase, params);
208        let file_key = [1u8; 16];
209
210        let ciphertext = encrypt_with_file_key(&file_key, &recipient, b"test data").unwrap();
211
212        // The output should start with the age header
213        assert!(ciphertext.starts_with(b"age-encryption.org/v1\n"));
214
215        // And be decryptable by the standard age library
216        let identity = Argon2idIdentity::new(passphrase);
217        let result = age::decrypt(&identity, &ciphertext).unwrap();
218        assert_eq!(result, b"test data");
219    }
220
221    #[test]
222    fn test_encrypt_empty_plaintext() {
223        let passphrase = b"test";
224        let params = fast_params();
225        let recipient = Argon2idRecipient::new(passphrase, params);
226        let file_key = [1u8; 16];
227
228        let ciphertext = encrypt_with_file_key(&file_key, &recipient, b"").unwrap();
229
230        let identity = Argon2idIdentity::new(passphrase);
231        let result = age::decrypt(&identity, &ciphertext).unwrap();
232        assert_eq!(result, b"");
233    }
234}