1use rand::rngs::OsRng;
35use std::collections::HashSet;
36use std::fmt;
37
38use anubis_core::{
39 format::{FileKey, Stanza, FILE_KEY_BYTES},
40 primitives::{aead_decrypt, aead_encrypt, hkdf},
41 secrecy::ExposeSecret,
42};
43use base64::{prelude::BASE64_STANDARD_NO_PAD, Engine};
44use oqs::kem::{Algorithm, Kem};
45use zeroize::{Zeroize, Zeroizing};
46
47use crate::{
48 error::{DecryptError, EncryptError},
49 pqc::{mlkem, x25519},
50};
51
52const HYBRID_RECIPIENT_TAG: &str = "hybrid";
53const HYBRID_LABEL: &str = "postquantum"; fn mlkem() -> Kem {
56 oqs::init();
57 Kem::new(Algorithm::MlKem1024).expect("ML-KEM-1024 algorithm available")
58}
59
60fn hybrid_combiner(
69 x25519_ss: &[u8; 32],
70 mlkem_ss: &[u8; 32],
71 x25519_epk: &[u8; 32],
72 mlkem_ct: &[u8; 1568],
73) -> [u8; 32] {
74 let mut ikm = Vec::with_capacity(64);
76 ikm.extend_from_slice(x25519_ss);
77 ikm.extend_from_slice(mlkem_ss);
78
79 let mut salt = Vec::with_capacity(1600);
81 salt.extend_from_slice(x25519_epk);
82 salt.extend_from_slice(mlkem_ct);
83
84 hkdf(&salt, b"anubis-hybrid-v2/X25519+MLKEM-1024", &ikm)
86}
87
88pub struct Identity {
93 x25519: x25519::Identity,
94 mlkem: mlkem::Identity,
95}
96
97impl Identity {
98 pub fn generate() -> Self {
100 Identity {
101 x25519: x25519::Identity::generate(),
102 mlkem: mlkem::Identity::generate(),
103 }
104 }
105
106 pub fn to_public(&self) -> Recipient {
108 Recipient {
109 x25519: self.x25519.to_public(),
110 mlkem: self.mlkem.to_public(),
111 }
112 }
113
114 fn unwrap_stanza(&self, stanza: &Stanza) -> Option<Result<FileKey, DecryptError>> {
118 if stanza.tag != HYBRID_RECIPIENT_TAG {
119 return None;
120 }
121
122 if stanza.args.len() != 2 {
129 return Some(Err(DecryptError::InvalidHeader));
130 }
131
132 let x25519_epk_bytes = match BASE64_STANDARD_NO_PAD.decode(&stanza.args[0]) {
134 Ok(bytes) if bytes.len() == 32 => {
135 let mut arr = [0u8; 32];
136 arr.copy_from_slice(&bytes);
137 arr
138 }
139 _ => return Some(Err(DecryptError::InvalidHeader)),
140 };
141
142 let mlkem_ct_bytes = match BASE64_STANDARD_NO_PAD.decode(&stanza.args[1]) {
144 Ok(bytes) if bytes.len() == 1568 => {
145 let mut arr = [0u8; 1568];
146 arr.copy_from_slice(&bytes);
147 arr
148 }
149 _ => return Some(Err(DecryptError::InvalidHeader)),
150 };
151
152 let x25519_ss = match self.x25519.diffie_hellman(&x25519_epk_bytes) {
154 Ok(ss) => ss,
155 Err(_) => return Some(Err(DecryptError::DecryptionFailed)),
156 };
157
158 let mlkem_ss = match self.mlkem.decapsulate(&mlkem_ct_bytes) {
160 Ok(ss) => ss,
161 Err(_) => return Some(Err(DecryptError::DecryptionFailed)),
162 };
163
164 let wrap_key = hybrid_combiner(&x25519_ss, &mlkem_ss, &x25519_epk_bytes, &mlkem_ct_bytes);
166
167 const ENCRYPTED_FILE_KEY_BYTES: usize = FILE_KEY_BYTES + 16;
169 if stanza.body.len() != ENCRYPTED_FILE_KEY_BYTES {
170 return Some(Err(DecryptError::InvalidHeader));
171 }
172
173 aead_decrypt(&Zeroizing::new(wrap_key), FILE_KEY_BYTES, &stanza.body)
174 .ok()
175 .map(|mut plaintext| {
176 Ok(FileKey::init_with_mut(|file_key| {
177 file_key.copy_from_slice(&plaintext);
178 plaintext.zeroize();
179 }))
180 })
181 }
182}
183
184impl crate::Identity for Identity {
185 fn unwrap_stanza(&self, stanza: &Stanza) -> Option<Result<FileKey, DecryptError>> {
186 Identity::unwrap_stanza(self, stanza)
187 }
188}
189
190#[derive(Clone)]
195pub struct Recipient {
196 x25519: x25519::Recipient,
197 mlkem: mlkem::Recipient,
198}
199
200impl Recipient {
201 fn wrap_file_key(&self, file_key: &FileKey) -> Result<Vec<Stanza>, EncryptError> {
203 let mut rng = OsRng;
204
205 let x25519_esk = x25519_dalek::EphemeralSecret::random_from_rng(&mut rng);
207 let x25519_epk = x25519_dalek::PublicKey::from(&x25519_esk);
208
209 let x25519_ss = x25519_esk.diffie_hellman(self.x25519.public_key());
211 let x25519_ss_bytes: [u8; 32] = *x25519_ss.as_bytes();
212
213 let (mlkem_ct, mlkem_ss) = self.mlkem.encapsulate(&mut rng)?;
215 let mlkem_ss_bytes: [u8; 32] = mlkem_ss[..32].try_into().unwrap();
216
217 let x25519_epk_bytes: [u8; 32] = *x25519_epk.as_bytes();
219 let wrap_key = hybrid_combiner(&x25519_ss_bytes, &mlkem_ss_bytes, &x25519_epk_bytes, &mlkem_ct);
220
221 let encrypted_file_key = aead_encrypt(&Zeroizing::new(wrap_key), file_key.expose_secret());
223
224 Ok(vec![Stanza {
225 tag: HYBRID_RECIPIENT_TAG.to_string(),
226 args: vec![
227 BASE64_STANDARD_NO_PAD.encode(&x25519_epk_bytes),
228 BASE64_STANDARD_NO_PAD.encode(&mlkem_ct),
229 ],
230 body: encrypted_file_key,
231 }])
232 }
233}
234
235impl crate::Recipient for Recipient {
236 fn wrap_file_key(
237 &self,
238 file_key: &FileKey,
239 ) -> Result<(Vec<Stanza>, HashSet<String>), EncryptError> {
240 let mut labels = HashSet::new();
241 labels.insert(HYBRID_LABEL.to_string());
242 Ok((self.wrap_file_key(file_key)?, labels))
243 }
244}
245
246impl fmt::Display for Recipient {
247 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
248 write!(f, "anubis1hybrid{}{}", self.x25519, self.mlkem)
249 }
250}
251
252impl fmt::Display for Identity {
253 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
254 write!(
255 f,
256 "ANUBIS-HYBRID-SECRET-KEY-1{}\n{}",
257 self.x25519, self.mlkem
258 )
259 }
260}
261
262#[cfg(test)]
263mod tests {
264 use super::*;
265
266 #[test]
267 fn hybrid_combiner_deterministic() {
268 let x25519_ss = [1u8; 32];
269 let mlkem_ss = [2u8; 32];
270 let x25519_epk = [3u8; 32];
271 let mlkem_ct = [4u8; 1568];
272
273 let key1 = hybrid_combiner(&x25519_ss, &mlkem_ss, &x25519_epk, &mlkem_ct);
274 let key2 = hybrid_combiner(&x25519_ss, &mlkem_ss, &x25519_epk, &mlkem_ct);
275
276 assert_eq!(key1, key2, "Combiner should be deterministic");
277 }
278
279 #[test]
280 fn hybrid_combiner_different_inputs() {
281 let x25519_ss = [1u8; 32];
282 let mlkem_ss = [2u8; 32];
283 let x25519_epk = [3u8; 32];
284 let mlkem_ct = [4u8; 1568];
285
286 let key1 = hybrid_combiner(&x25519_ss, &mlkem_ss, &x25519_epk, &mlkem_ct);
287
288 let x25519_ss2 = [5u8; 32];
290 let key2 = hybrid_combiner(&x25519_ss2, &mlkem_ss, &x25519_epk, &mlkem_ct);
291 assert_ne!(key1, key2, "Different X25519 SS should produce different key");
292
293 let mlkem_ss2 = [6u8; 32];
295 let key3 = hybrid_combiner(&x25519_ss, &mlkem_ss2, &x25519_epk, &mlkem_ct);
296 assert_ne!(key1, key3, "Different ML-KEM SS should produce different key");
297 }
298
299 #[test]
300 fn hybrid_round_trip() {
301 let identity = Identity::generate();
302 let recipient = identity.to_public();
303
304 let file_key = FileKey::new(Box::new([42; 16]));
305
306 let stanzas = recipient.wrap_file_key(&file_key).unwrap();
308 assert_eq!(stanzas.len(), 1);
309 assert_eq!(stanzas[0].tag, HYBRID_RECIPIENT_TAG);
310 assert_eq!(stanzas[0].args.len(), 2); let decrypted = identity.unwrap_stanza(&stanzas[0]).unwrap().unwrap();
314 assert_eq!(decrypted.expose_secret(), file_key.expose_secret());
315 }
316
317 #[test]
318 fn hybrid_labels() {
319 let recipient = Identity::generate().to_public();
320 let file_key = FileKey::new(Box::new([42; 16]));
321
322 let (_, labels) = <Recipient as crate::Recipient>::wrap_file_key(&recipient, &file_key)
323 .unwrap();
324
325 assert!(labels.contains(HYBRID_LABEL));
326 assert_eq!(labels.len(), 1);
327 }
328}