1#[cfg(feature = "alloc")]
19use alloc::boxed::Box;
20use alloc::string::ToString;
21use alloc::vec;
22use alloc::vec::Vec;
23
24use lib_q_core::{
25 Aead,
26 AeadDecryptSemantic,
27 AeadKey,
28 Algorithm,
29 DecryptSemanticOutcome,
30 Error,
31 Nonce,
32 Result,
33};
34use lib_q_sha3::digest::{
35 Update,
36 XofReader,
37};
38use zeroize::Zeroizing;
39
40use crate::metadata::{
42 AeadMetadata,
43 AeadWithMetadata,
44};
45use crate::security::memory::secure_zero_slice;
46use crate::security::stack_buffer::UninitStackBuffer;
47
48pub struct Shake256Aead {
50 metadata: &'static AeadMetadata,
51}
52
53impl Shake256Aead {
54 pub fn new() -> Self {
56 Self {
57 metadata: crate::metadata::get_metadata(Algorithm::Shake256Aead)
58 .expect("SHAKE256 AEAD metadata not found"),
59 }
60 }
61
62 const DOMAIN_ENC_KEY: &'static [u8] = b"LIBQ-SHAKE256-AEAD-ENC-KEY";
64 const DOMAIN_MAC_KEY: &'static [u8] = b"LIBQ-SHAKE256-AEAD-MAC-KEY";
65 const DOMAIN_IV: &'static [u8] = b"LIBQ-SHAKE256-AEAD-IV";
66 const DOMAIN_TAG: &'static [u8] = b"LIBQ-SHAKE256-AEAD-TAG";
67
68 fn derive_encryption_key(
70 &self,
71 key: &[u8],
72 nonce: &[u8],
73 associated_data: &[u8],
74 ) -> Result<Zeroizing<[u8; 32]>> {
75 use lib_q_sha3::Shake256;
76 use lib_q_sha3::digest::ExtendableOutput;
77
78 let mut hasher = Shake256::default();
79 hasher.update(Self::DOMAIN_ENC_KEY);
80 hasher.update(key);
81 hasher.update(nonce);
82 hasher.update(associated_data);
83 hasher.update(&(associated_data.len() as u64).to_le_bytes());
84
85 let mut enc_key = Zeroizing::new([0u8; 32]);
86 let mut reader = hasher.finalize_xof();
87 reader.read(&mut enc_key[..]);
88 Ok(enc_key)
89 }
90
91 fn derive_mac_key(
93 &self,
94 key: &[u8],
95 nonce: &[u8],
96 associated_data: &[u8],
97 ) -> Result<Zeroizing<[u8; 32]>> {
98 use lib_q_sha3::Shake256;
99 use lib_q_sha3::digest::ExtendableOutput;
100
101 let mut hasher = Shake256::default();
102 hasher.update(Self::DOMAIN_MAC_KEY);
103 hasher.update(key);
104 hasher.update(nonce);
105 hasher.update(associated_data);
106 hasher.update(&(associated_data.len() as u64).to_le_bytes());
107
108 let mut mac_key = Zeroizing::new([0u8; 32]);
109 let mut reader = hasher.finalize_xof();
110 reader.read(&mut mac_key[..]);
111 Ok(mac_key)
112 }
113
114 fn generate_iv(
116 &self,
117 key: &[u8],
118 nonce: &[u8],
119 associated_data: &[u8],
120 ) -> Result<Zeroizing<[u8; 16]>> {
121 use lib_q_sha3::Shake256;
122 use lib_q_sha3::digest::ExtendableOutput;
123
124 let mut hasher = Shake256::default();
125 hasher.update(Self::DOMAIN_IV);
126 hasher.update(key);
127 hasher.update(nonce);
128 hasher.update(associated_data);
129 hasher.update(&(associated_data.len() as u64).to_le_bytes());
130
131 let mut iv = Zeroizing::new([0u8; 16]);
132 let mut reader = hasher.finalize_xof();
133 reader.read(&mut iv[..]);
134 Ok(iv)
135 }
136
137 fn generate_tag(
139 &self,
140 mac_key: &[u8],
141 associated_data: &[u8],
142 ciphertext: &[u8],
143 ) -> Result<Zeroizing<[u8; 32]>> {
144 use lib_q_sha3::Shake256;
145 use lib_q_sha3::digest::ExtendableOutput;
146
147 let mut hasher = Shake256::default();
148 hasher.update(Self::DOMAIN_TAG);
149 hasher.update(mac_key);
150 hasher.update(associated_data);
151 hasher.update(&(associated_data.len() as u64).to_le_bytes());
152 hasher.update(ciphertext);
153 hasher.update(&(ciphertext.len() as u64).to_le_bytes());
154
155 let mut tag = Zeroizing::new([0u8; 32]);
156 let mut reader = hasher.finalize_xof();
157 reader.read(&mut tag[..]);
158 Ok(tag)
159 }
160
161 fn ctr_encrypt(&self, enc_key: &[u8], iv: &[u8], plaintext: &mut [u8]) -> Result<()> {
163 use lib_q_sha3::Shake256;
164 use lib_q_sha3::digest::ExtendableOutput;
165
166 let mut hasher = Shake256::default();
168 hasher.update(enc_key);
169 hasher.update(iv);
170 hasher.update(&(plaintext.len() as u64).to_le_bytes());
171
172 if plaintext.len() <= 4096 {
175 let mut keystream = UninitStackBuffer::<4096>::new();
176 keystream
177 .resize(plaintext.len())
178 .map_err(|_| Error::BufferTooSmall {
179 capacity: keystream.capacity(),
180 requested: plaintext.len(),
181 })?;
182 let mut reader = hasher.finalize_xof();
183 reader.read(keystream.as_mut_slice());
184
185 for (i, byte) in plaintext.iter_mut().enumerate() {
187 *byte ^= keystream.as_slice()[i];
188 }
189 } else {
190 let mut keystream = vec![0u8; plaintext.len()];
192 let mut reader = hasher.finalize_xof();
193 reader.read(&mut keystream);
194
195 for (i, byte) in plaintext.iter_mut().enumerate() {
197 *byte ^= keystream[i];
198 }
199 secure_zero_slice(&mut keystream);
200 }
201
202 Ok(())
203 }
204
205 fn ctr_decrypt(&self, enc_key: &[u8], iv: &[u8], ciphertext: &mut [u8]) -> Result<()> {
207 self.ctr_encrypt(enc_key, iv, ciphertext)
209 }
210
211 fn validate_key(&self, key: &AeadKey) -> Result<()> {
213 AeadWithMetadata::validate_key(self, key)
215 }
216
217 fn validate_nonce(&self, nonce: &Nonce) -> Result<()> {
219 AeadWithMetadata::validate_nonce(self, nonce)
221 }
222
223 fn validate_ciphertext_size(&self, size: usize) -> Result<()> {
225 AeadWithMetadata::validate_ciphertext_size(self, size)
227 }
228
229 fn encrypt_internal(
231 &self,
232 key: &AeadKey,
233 nonce: &Nonce,
234 plaintext: &[u8],
235 associated_data: Option<&[u8]>,
236 ) -> Result<Vec<u8>> {
237 self.validate_key(key)?;
239 self.validate_nonce(nonce)?;
240 crate::security::validation::validate_plaintext(plaintext)?;
241 crate::security::validation::validate_key(key.as_bytes())?;
242 crate::security::validation::validate_nonce(nonce.as_bytes())?;
243
244 let associated_data = associated_data.unwrap_or(&[]);
245 crate::security::validation::validate_associated_data(associated_data)?;
246
247 #[cfg(feature = "shake256")]
249 {
250 let mut key_staged = Zeroizing::new([0u8; 32]);
251 key_staged.copy_from_slice(key.as_bytes());
252 let mut nonce_staged = Zeroizing::new([0u8; 16]);
253 nonce_staged.copy_from_slice(nonce.as_bytes());
254 let kb = key_staged.as_slice();
255 let nb = nonce_staged.as_slice();
256
257 let enc_key = self.derive_encryption_key(kb, nb, associated_data)?;
259 let mac_key = self.derive_mac_key(kb, nb, associated_data)?;
260 let iv = self.generate_iv(kb, nb, associated_data)?;
261
262 let mut ciphertext = plaintext.to_vec();
264 self.ctr_encrypt(enc_key.as_slice(), iv.as_slice(), &mut ciphertext)?;
265
266 let tag = self.generate_tag(mac_key.as_slice(), associated_data, &ciphertext)?;
268
269 ciphertext.extend_from_slice(tag.as_slice());
271
272 Ok(ciphertext)
273 }
274
275 #[cfg(not(feature = "shake256"))]
276 {
277 Err(Error::NotImplemented {
278 feature: "SHAKE256 AEAD implementation requires 'shake256' feature",
279 })
280 }
281 }
282
283 fn decrypt_semantic_core(
286 &self,
287 key: &AeadKey,
288 nonce: &Nonce,
289 ciphertext: &[u8],
290 associated_data: Option<&[u8]>,
291 ) -> Result<DecryptSemanticOutcome> {
292 self.validate_key(key)?;
293 self.validate_nonce(nonce)?;
294 self.validate_ciphertext_size(ciphertext.len())?;
295 crate::security::validation::validate_ciphertext(ciphertext)?;
296 crate::security::validation::validate_key(key.as_bytes())?;
297 crate::security::validation::validate_nonce(nonce.as_bytes())?;
298
299 let associated_data = associated_data.unwrap_or(&[]);
300 crate::security::validation::validate_associated_data(associated_data)?;
301
302 #[cfg(feature = "shake256")]
303 {
304 let (ciphertext_data, tag) = ciphertext.split_at(ciphertext.len() - self.tag_size());
305
306 let mut key_staged = Zeroizing::new([0u8; 32]);
307 key_staged.copy_from_slice(key.as_bytes());
308 let mut nonce_staged = Zeroizing::new([0u8; 16]);
309 nonce_staged.copy_from_slice(nonce.as_bytes());
310 let kb = key_staged.as_slice();
311 let nb = nonce_staged.as_slice();
312
313 let enc_key = self.derive_encryption_key(kb, nb, associated_data)?;
314 let mac_key = self.derive_mac_key(kb, nb, associated_data)?;
315 let iv = self.generate_iv(kb, nb, associated_data)?;
316
317 let computed_tag =
318 self.generate_tag(mac_key.as_slice(), associated_data, ciphertext_data)?;
319 let tag_valid =
320 crate::security::constant_time::constant_time_eq(tag, computed_tag.as_slice());
321
322 let mut plaintext = ciphertext_data.to_vec();
323 self.ctr_decrypt(enc_key.as_slice(), iv.as_slice(), &mut plaintext)?;
324
325 crate::security::constant_time::constant_time_zero(!tag_valid, &mut plaintext);
326
327 if tag_valid {
328 Ok(DecryptSemanticOutcome::Success(Zeroizing::new(plaintext)))
329 } else {
330 Ok(DecryptSemanticOutcome::AuthenticationFailed)
331 }
332 }
333
334 #[cfg(not(feature = "shake256"))]
335 {
336 Err(Error::NotImplemented {
337 feature: "SHAKE256 AEAD implementation requires 'shake256' feature",
338 })
339 }
340 }
341
342 fn decrypt_internal(
344 &self,
345 key: &AeadKey,
346 nonce: &Nonce,
347 ciphertext: &[u8],
348 associated_data: Option<&[u8]>,
349 ) -> Result<Vec<u8>> {
350 match self.decrypt_semantic_core(key, nonce, ciphertext, associated_data)? {
351 DecryptSemanticOutcome::Success(pt) => Ok(Vec::clone(&*pt)),
352 DecryptSemanticOutcome::AuthenticationFailed => Err(Error::AuthenticationFailed {
353 operation: "Tag verification failed".to_string(),
354 }),
355 }
356 }
357}
358
359impl Aead for Shake256Aead {
360 fn encrypt(
361 &self,
362 key: &AeadKey,
363 nonce: &Nonce,
364 plaintext: &[u8],
365 associated_data: Option<&[u8]>,
366 ) -> Result<Vec<u8>> {
367 self.encrypt_internal(key, nonce, plaintext, associated_data)
368 }
369
370 fn decrypt(
371 &self,
372 key: &AeadKey,
373 nonce: &Nonce,
374 ciphertext: &[u8],
375 associated_data: Option<&[u8]>,
376 ) -> Result<Vec<u8>> {
377 self.decrypt_internal(key, nonce, ciphertext, associated_data)
378 }
379}
380
381impl AeadDecryptSemantic for Shake256Aead {
382 fn decrypt_semantic(
383 &self,
384 key: &AeadKey,
385 nonce: &Nonce,
386 ciphertext: &[u8],
387 associated_data: Option<&[u8]>,
388 ) -> Result<DecryptSemanticOutcome> {
389 self.decrypt_semantic_core(key, nonce, ciphertext, associated_data)
390 }
391}
392
393impl AeadWithMetadata for Shake256Aead {
394 fn metadata(&self) -> &'static AeadMetadata {
395 self.metadata
396 }
397}
398
399impl Default for Shake256Aead {
400 fn default() -> Self {
401 Self::new()
402 }
403}
404
405impl crate::plugin::AeadPlugin for Shake256Aead {
407 fn algorithm(&self) -> Algorithm {
408 Algorithm::Shake256Aead
409 }
410
411 fn create(&self) -> Result<Box<dyn AeadWithMetadata>> {
412 Ok(Box::new(Self::new()))
413 }
414
415 fn metadata(&self) -> &'static AeadMetadata {
416 crate::metadata::get_metadata(Algorithm::Shake256Aead)
417 .expect("Metadata not found for algorithm")
418 }
419
420 fn name(&self) -> &'static str {
421 "SHAKE256 AEAD"
422 }
423
424 fn version(&self) -> &'static str {
425 "1.0.0"
426 }
427
428 fn description(&self) -> &'static str {
429 "SHAKE256-based AEAD construction using post-quantum hash function with proper domain separation"
430 }
431}
432
433#[cfg(test)]
434mod tests {
435 use lib_q_core::{
436 AeadDecryptSemantic,
437 DecryptSemanticOutcome,
438 };
439
440 use super::*;
441
442 fn create_test_key() -> AeadKey {
444 AeadKey::new(vec![
445 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, 0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54,
446 0x32, 0x10, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC,
447 0xDD, 0xEE, 0xFF, 0x00,
448 ])
449 }
450
451 fn create_test_nonce() -> Nonce {
453 Nonce::new(vec![
454 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, 0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54,
455 0x32, 0x10,
456 ])
457 }
458
459 fn create_test_key2() -> AeadKey {
461 AeadKey::new(vec![
462 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE,
463 0xFF, 0x00, 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, 0xFE, 0xDC, 0xBA, 0x98,
464 0x76, 0x54, 0x32, 0x10,
465 ])
466 }
467
468 fn create_test_nonce2() -> Nonce {
470 Nonce::new(vec![
471 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE,
472 0xFF, 0x00,
473 ])
474 }
475
476 #[cfg(feature = "shake256")]
477 #[test]
478 fn test_shake256_decrypt_semantic_matches_decrypt() {
479 let aead = Shake256Aead::new();
480 let key = create_test_key();
481 let nonce = create_test_nonce();
482 let ad = b"metadata";
483 let pt = b"semantic path";
484 let ct = aead.encrypt(&key, &nonce, pt, Some(ad)).expect("encrypt");
485 let layer_a = aead.decrypt(&key, &nonce, &ct, Some(ad)).expect("decrypt");
486 match aead
487 .decrypt_semantic(&key, &nonce, &ct, Some(ad))
488 .expect("decrypt_semantic")
489 {
490 DecryptSemanticOutcome::Success(got) => {
491 assert_eq!(got.as_slice(), layer_a.as_slice())
492 }
493 DecryptSemanticOutcome::AuthenticationFailed => {
494 panic!("expected Success")
495 }
496 }
497 }
498
499 #[cfg(feature = "shake256")]
500 #[test]
501 fn test_shake256_decrypt_semantic_tampered_tag() {
502 let aead = Shake256Aead::new();
503 let key = create_test_key();
504 let nonce = create_test_nonce();
505 let mut ct = aead.encrypt(&key, &nonce, b"x", None).unwrap();
506 *ct.last_mut().unwrap() ^= 1;
507 let out = aead.decrypt_semantic(&key, &nonce, &ct, None).unwrap();
508 assert_eq!(out, DecryptSemanticOutcome::AuthenticationFailed);
509 }
510
511 #[test]
512 fn test_shake256_creation() {
513 let aead = Shake256Aead::new();
514 assert_eq!(aead.algorithm(), Algorithm::Shake256Aead);
515 assert_eq!(aead.key_size(), 32);
516 assert_eq!(aead.nonce_size(), 16);
517 assert_eq!(aead.tag_size(), 32);
518 assert_eq!(aead.security_level(), 1);
519 }
520
521 #[test]
522 fn test_shake256_metadata() {
523 let aead = Shake256Aead::new();
524 let metadata = aead.metadata();
525
526 assert_eq!(metadata.algorithm, Algorithm::Shake256Aead);
527 assert_eq!(metadata.name, "SHAKE256-AEAD");
528 assert_eq!(metadata.key_size, 32);
529 assert_eq!(metadata.nonce_size, 16);
530 assert_eq!(metadata.tag_size, 32);
531 assert_eq!(metadata.security_level, 1);
532 }
533
534 #[test]
535 fn test_shake256_validation() {
536 let aead = Shake256Aead::new();
537
538 let key = AeadKey::new(vec![0u8; 32]);
540 assert!(aead.validate_key(&key).is_ok());
541
542 let invalid_key = AeadKey::new(vec![0u8; 16]);
544 assert!(aead.validate_key(&invalid_key).is_err());
545
546 let nonce = Nonce::new(vec![0u8; 16]);
548 assert!(aead.validate_nonce(&nonce).is_ok());
549
550 let invalid_nonce = Nonce::new(vec![0u8; 12]);
552 assert!(aead.validate_nonce(&invalid_nonce).is_err());
553 }
554
555 #[cfg(feature = "shake256")]
556 #[test]
557 fn test_shake256_encrypt_decrypt() {
558 let aead = Shake256Aead::new();
559
560 let key = create_test_key();
561 let nonce = create_test_nonce();
562 let plaintext = b"Hello, World!";
563 let associated_data = b"metadata";
564
565 let ciphertext = aead.encrypt(&key, &nonce, plaintext, Some(associated_data.as_slice()));
567 assert!(ciphertext.is_ok());
568
569 let ciphertext = ciphertext.unwrap();
570 assert_eq!(ciphertext.len(), plaintext.len() + aead.tag_size());
571
572 let decrypted = aead.decrypt(&key, &nonce, &ciphertext, Some(associated_data.as_slice()));
574 assert!(decrypted.is_ok());
575 assert_eq!(decrypted.unwrap(), plaintext);
576 }
577
578 #[cfg(feature = "shake256")]
579 #[test]
580 fn test_shake256_authentication_failure() {
581 let aead = Shake256Aead::new();
582
583 let key = create_test_key();
584 let nonce = create_test_nonce();
585 let plaintext = b"Hello, World!";
586
587 let ciphertext = aead.encrypt(&key, &nonce, plaintext, None).unwrap();
589
590 let mut tampered = ciphertext.clone();
592 tampered[0] ^= 0xFF;
593
594 let result = aead.decrypt(&key, &nonce, &tampered, None);
596 assert!(result.is_err());
597
598 if let Err(Error::AuthenticationFailed { operation }) = result {
599 assert!(operation.contains("Tag verification failed"));
600 } else {
601 panic!("Expected AuthenticationFailed error");
602 }
603 }
604
605 #[cfg(feature = "shake256")]
606 #[test]
607 fn test_shake256_wrong_key() {
608 let aead = Shake256Aead::new();
609
610 let key1 = create_test_key();
611 let key2 = create_test_key2();
612 let nonce = create_test_nonce();
613 let plaintext = b"Hello, World!";
614
615 let ciphertext = aead.encrypt(&key1, &nonce, plaintext, None).unwrap();
617
618 let result = aead.decrypt(&key2, &nonce, &ciphertext, None);
620 assert!(result.is_err());
621 }
622
623 #[cfg(feature = "shake256")]
624 #[test]
625 fn test_shake256_wrong_nonce() {
626 let aead = Shake256Aead::new();
627
628 let key = create_test_key();
629 let nonce1 = create_test_nonce();
630 let nonce2 = create_test_nonce2();
631 let plaintext = b"Hello, World!";
632
633 let ciphertext = aead.encrypt(&key, &nonce1, plaintext, None).unwrap();
635
636 let result = aead.decrypt(&key, &nonce2, &ciphertext, None);
638 assert!(result.is_err());
639 }
640
641 #[cfg(feature = "shake256")]
642 #[test]
643 fn test_shake256_empty_plaintext() {
644 let aead = Shake256Aead::new();
645
646 let key = create_test_key();
647 let nonce = create_test_nonce();
648 let plaintext = b"";
649
650 let ciphertext = aead.encrypt(&key, &nonce, plaintext, None);
652 assert!(ciphertext.is_ok());
653
654 let ciphertext = ciphertext.unwrap();
655 assert_eq!(ciphertext.len(), aead.tag_size());
656
657 let decrypted = aead.decrypt(&key, &nonce, &ciphertext, None);
659 assert!(decrypted.is_ok());
660 assert_eq!(decrypted.unwrap(), plaintext);
661 }
662
663 #[cfg(feature = "shake256")]
664 #[test]
665 fn test_shake256_domain_separation() {
666 let aead = Shake256Aead::new();
667
668 let key = create_test_key();
669 let nonce = create_test_nonce();
670 let plaintext = b"Test message";
671
672 let ciphertext1 = aead.encrypt(&key, &nonce, plaintext, Some(b"ad1")).unwrap();
674 let ciphertext2 = aead.encrypt(&key, &nonce, plaintext, Some(b"ad2")).unwrap();
675
676 assert_ne!(ciphertext1, ciphertext2);
678
679 let decrypted1 = aead
681 .decrypt(&key, &nonce, &ciphertext1, Some(b"ad1"))
682 .unwrap();
683 let decrypted2 = aead
684 .decrypt(&key, &nonce, &ciphertext2, Some(b"ad2"))
685 .unwrap();
686
687 assert_eq!(decrypted1, plaintext);
688 assert_eq!(decrypted2, plaintext);
689 }
690
691 #[cfg(feature = "shake256")]
692 #[test]
693 fn test_shake256_nonce_uniqueness() {
694 let aead = Shake256Aead::new();
695
696 let key = create_test_key();
697 let plaintext = b"Test message";
698
699 let nonce1 = create_test_nonce();
701 let nonce2 = create_test_nonce2();
702
703 let ciphertext1 = aead.encrypt(&key, &nonce1, plaintext, None).unwrap();
704 let ciphertext2 = aead.encrypt(&key, &nonce2, plaintext, None).unwrap();
705
706 assert_ne!(ciphertext1, ciphertext2);
708
709 let decrypted1 = aead.decrypt(&key, &nonce1, &ciphertext1, None).unwrap();
711 let decrypted2 = aead.decrypt(&key, &nonce2, &ciphertext2, None).unwrap();
712
713 assert_eq!(decrypted1, plaintext);
714 assert_eq!(decrypted2, plaintext);
715 }
716}