Skip to main content

age_plugin_argon2/
cached.rs

1use std::collections::HashSet;
2
3use age::{DecryptError, EncryptError};
4use age_core::format::{FileKey, Stanza};
5use base64::engine::general_purpose::STANDARD_NO_PAD;
6use base64::Engine;
7use secrecy::ExposeSecret;
8use serde::{Deserialize, Serialize};
9use zeroize::{Zeroize, ZeroizeOnDrop};
10
11use crate::params::Argon2Params;
12
13const STANZA_TAG: &str = "thesis.co/argon2";
14
15/// Cached key material from a successful Argon2id decryption.
16///
17/// Contains everything needed to re-encrypt and re-decrypt without
18/// running the KDF again. Stored in the OS keychain during a session.
19#[derive(Clone, Serialize, Deserialize, Zeroize, ZeroizeOnDrop)]
20pub struct CachedMaterial {
21    /// The age FileKey (16 bytes)
22    pub file_key: [u8; 16],
23    /// The Argon2id-derived wrapping key (32 bytes)
24    pub wrapping_key: [u8; 32],
25    /// The salt used for key derivation (16 bytes)
26    pub salt: [u8; 16],
27    /// The Argon2id parameters used
28    #[zeroize(skip)]
29    pub params: Argon2Params,
30}
31
32/// Zero-KDF identity for session reads.
33///
34/// Returns the cached FileKey directly without any cryptographic operations.
35pub struct CachedIdentity {
36    file_key: [u8; 16],
37}
38
39impl CachedIdentity {
40    /// Create a new cached identity from previously captured key material.
41    pub fn new(material: &CachedMaterial) -> Self {
42        Self {
43            file_key: material.file_key,
44        }
45    }
46}
47
48impl Drop for CachedIdentity {
49    fn drop(&mut self) {
50        self.file_key.zeroize();
51    }
52}
53
54impl age::Identity for CachedIdentity {
55    /// Return the cached FileKey for any matching stanza.
56    ///
57    /// This intentionally skips body verification: the cached file key was
58    /// already authenticated during the initial full-KDF `Argon2idIdentity`
59    /// decryption. If the file key is wrong (e.g. corrupted keychain), the
60    /// age STREAM layer will detect it via its per-chunk Poly1305 MAC and
61    /// return a decryption error — no silent data corruption is possible.
62    fn unwrap_stanza(&self, stanza: &Stanza) -> Option<Result<FileKey, DecryptError>> {
63        if stanza.tag != STANZA_TAG {
64            return None;
65        }
66
67        Some(Ok(FileKey::new(Box::new(self.file_key))))
68    }
69}
70
71/// Zero-KDF recipient for session writes.
72///
73/// Wraps the FileKey using the cached wrapping key and salt,
74/// avoiding any Argon2id computation.
75pub struct CachedRecipient {
76    wrapping_key: [u8; 32],
77    salt: [u8; 16],
78    params: Argon2Params,
79}
80
81impl CachedRecipient {
82    /// Create a new cached recipient from previously captured key material.
83    pub fn new(material: &CachedMaterial) -> Self {
84        Self {
85            wrapping_key: material.wrapping_key,
86            salt: material.salt,
87            params: material.params,
88        }
89    }
90}
91
92impl Drop for CachedRecipient {
93    fn drop(&mut self) {
94        self.wrapping_key.zeroize();
95        self.salt.zeroize();
96    }
97}
98
99impl age::Recipient for CachedRecipient {
100    fn wrap_file_key(
101        &self,
102        file_key: &FileKey,
103    ) -> Result<(Vec<Stanza>, HashSet<String>), EncryptError> {
104        // Wrap FileKey using the cached wrapping key
105        let body = age_core::primitives::aead_encrypt(&self.wrapping_key, file_key.expose_secret());
106
107        let stanza = Stanza {
108            tag: STANZA_TAG.to_string(),
109            args: vec![
110                STANDARD_NO_PAD.encode(self.salt),
111                self.params.m_cost().to_string(),
112                self.params.t_cost().to_string(),
113                self.params.p_cost().to_string(),
114            ],
115            body,
116        };
117
118        // Same label pattern as Argon2idRecipient: must be only recipient
119        let mut labels = HashSet::new();
120        labels.insert(uuid::Uuid::new_v4().to_string());
121
122        Ok((vec![stanza], labels))
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use crate::identity::Argon2idIdentity;
130    use crate::recipient::Argon2idRecipient;
131    use age::{Identity, Recipient};
132
133    fn fast_params() -> Argon2Params {
134        Argon2Params::new(256, 1, 1).unwrap()
135    }
136
137    #[test]
138    fn test_cached_identity_returns_key_without_kdf() {
139        let material = CachedMaterial {
140            file_key: [42u8; 16],
141            wrapping_key: [0u8; 32],
142            salt: [0u8; 16],
143            params: fast_params(),
144        };
145
146        let identity = CachedIdentity::new(&material);
147        let stanza = Stanza {
148            tag: "thesis.co/argon2".to_string(),
149            args: vec![], // args not checked by CachedIdentity
150            body: vec![], // body not checked by CachedIdentity
151        };
152
153        let result = identity.unwrap_stanza(&stanza).unwrap().unwrap();
154        assert_eq!(result.expose_secret(), &[42u8; 16]);
155    }
156
157    #[test]
158    fn test_cached_identity_ignores_wrong_tag() {
159        let material = CachedMaterial {
160            file_key: [42u8; 16],
161            wrapping_key: [0u8; 32],
162            salt: [0u8; 16],
163            params: fast_params(),
164        };
165
166        let identity = CachedIdentity::new(&material);
167        let stanza = Stanza {
168            tag: "X25519".to_string(),
169            args: vec![],
170            body: vec![],
171        };
172
173        assert!(identity.unwrap_stanza(&stanza).is_none());
174    }
175
176    #[test]
177    fn test_cached_recipient_produces_valid_stanza() {
178        let material = CachedMaterial {
179            file_key: [42u8; 16],
180            wrapping_key: [99u8; 32],
181            salt: [1u8; 16],
182            params: fast_params(),
183        };
184
185        let recipient = CachedRecipient::new(&material);
186        let file_key = FileKey::new(Box::new([42u8; 16]));
187
188        let (stanzas, labels) = recipient.wrap_file_key(&file_key).unwrap();
189        assert_eq!(stanzas.len(), 1);
190        assert_eq!(stanzas[0].tag, "thesis.co/argon2");
191        assert_eq!(stanzas[0].args.len(), 4);
192        assert_eq!(stanzas[0].body.len(), 32);
193        assert_eq!(labels.len(), 1);
194    }
195
196    #[test]
197    fn test_cached_material_serde_roundtrip() {
198        let material = CachedMaterial {
199            file_key: [42u8; 16],
200            wrapping_key: [99u8; 32],
201            salt: [1u8; 16],
202            params: fast_params(),
203        };
204
205        let json = serde_json::to_string(&material).expect("serialize");
206        let restored: CachedMaterial = serde_json::from_str(&json).expect("deserialize");
207
208        assert_eq!(restored.file_key, material.file_key);
209        assert_eq!(restored.wrapping_key, material.wrapping_key);
210        assert_eq!(restored.salt, material.salt);
211        assert_eq!(restored.params, material.params);
212    }
213
214    #[test]
215    fn test_cached_material_serde_rejects_invalid_params() {
216        // Valid material but with p_cost=0 should fail deserialization
217        let json = r#"{"file_key":[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],"wrapping_key":[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],"salt":[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],"params":{"m_cost":256,"t_cost":1,"p_cost":0}}"#;
218        assert!(serde_json::from_str::<CachedMaterial>(json).is_err());
219    }
220
221    #[test]
222    fn test_cached_roundtrip_with_real_material() {
223        // Full roundtrip: encrypt with KDF → capture material → decrypt with cached
224        let passphrase = b"test-password";
225        let params = fast_params();
226
227        // Step 1: Encrypt with full KDF
228        let recipient = Argon2idRecipient::new(passphrase, params);
229        let original_key = FileKey::new(Box::new([42u8; 16]));
230        let (stanzas, _) = recipient.wrap_file_key(&original_key).unwrap();
231
232        // Step 2: Decrypt with full KDF and capture material
233        let identity = Argon2idIdentity::new(passphrase);
234        let decrypted = identity.unwrap_stanza(&stanzas[0]).unwrap().unwrap();
235        assert_eq!(decrypted.expose_secret(), &[42u8; 16]);
236
237        let material = identity.captured_material().unwrap();
238
239        // Step 3: Re-encrypt with cached recipient
240        let cached_recipient = CachedRecipient::new(&material);
241        let rekey = FileKey::new(Box::new(material.file_key));
242        let (new_stanzas, _) = cached_recipient.wrap_file_key(&rekey).unwrap();
243
244        // Step 4: Decrypt with cached identity
245        let cached_identity = CachedIdentity::new(&material);
246        let result = cached_identity
247            .unwrap_stanza(&new_stanzas[0])
248            .unwrap()
249            .unwrap();
250        assert_eq!(result.expose_secret(), &[42u8; 16]);
251
252        // Step 5: Also verify Argon2idIdentity can decrypt the cached stanza
253        let full_identity = Argon2idIdentity::new(passphrase);
254        let result2 = full_identity
255            .unwrap_stanza(&new_stanzas[0])
256            .unwrap()
257            .unwrap();
258        assert_eq!(result2.expose_secret(), &[42u8; 16]);
259    }
260}