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#[derive(Clone, Serialize, Deserialize, Zeroize, ZeroizeOnDrop)]
20pub struct CachedMaterial {
21 pub file_key: [u8; 16],
23 pub wrapping_key: [u8; 32],
25 pub salt: [u8; 16],
27 #[zeroize(skip)]
29 pub params: Argon2Params,
30}
31
32pub struct CachedIdentity {
36 file_key: [u8; 16],
37}
38
39impl CachedIdentity {
40 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 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
71pub struct CachedRecipient {
76 wrapping_key: [u8; 32],
77 salt: [u8; 16],
78 params: Argon2Params,
79}
80
81impl CachedRecipient {
82 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 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 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![], body: vec![], };
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 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 let passphrase = b"test-password";
225 let params = fast_params();
226
227 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 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 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 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 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}