1use arkhe_forge_core::pii::{
31 compute_aad, AeadKind, DekId, DekMessageCounter, PiiError, PiiType, RotationTrigger,
32};
33use bytes::Bytes;
34use serde::{Deserialize, Serialize};
35use std::cell::Cell;
36use std::marker::PhantomData;
37use zeroize::{Zeroize, ZeroizeOnDrop};
38
39#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
47pub struct DekConfig {
48 pub replica_id: u32,
53}
54
55#[derive(Zeroize, ZeroizeOnDrop)]
90pub struct Dek {
91 material: [u8; 32],
92 #[zeroize(skip)]
95 counter: Cell<u64>,
96 #[zeroize(skip)]
99 replica_id: u32,
100}
101
102impl Dek {
103 #[inline]
108 #[must_use]
109 pub fn from_bytes(material: [u8; 32]) -> Self {
110 Self::with_config(material, DekConfig::default())
111 }
112
113 #[inline]
117 #[must_use]
118 pub fn with_config(material: [u8; 32], config: DekConfig) -> Self {
119 Self {
120 material,
121 counter: Cell::new(0),
122 replica_id: config.replica_id,
123 }
124 }
125
126 pub fn try_from_slice(bytes: &[u8]) -> Result<Self, PiiError> {
136 if bytes.len() != 32 {
137 return Err(PiiError::InvalidKeyLength);
138 }
139 let mut material = [0u8; 32];
140 material.copy_from_slice(bytes);
141 Ok(Self {
142 material,
143 counter: Cell::new(0),
144 replica_id: 0,
145 })
146 }
147
148 #[inline]
152 #[must_use]
153 #[cfg_attr(
154 not(any(feature = "tier-1-kms", feature = "tier-2-multi-kms")),
155 allow(dead_code)
156 )]
157 pub(crate) fn as_bytes(&self) -> &[u8; 32] {
158 &self.material
159 }
160
161 #[cfg_attr(not(feature = "tier-2-multi-kms"), allow(dead_code))]
167 fn advance_counter(&self) -> Result<u64, PiiError> {
168 let n = self.counter.get();
169 if n == u64::MAX {
170 return Err(PiiError::DekExhausted);
171 }
172 self.counter.set(n.wrapping_add(1));
173 Ok(n)
174 }
175
176 #[cfg(all(test, feature = "tier-2-multi-kms"))]
179 pub(crate) fn set_counter_for_test(&self, n: u64) {
180 self.counter.set(n);
181 }
182
183 #[cfg(all(test, feature = "tier-2-multi-kms"))]
185 pub(crate) fn get_counter_for_test(&self) -> u64 {
186 self.counter.get()
187 }
188}
189
190#[cfg_attr(not(feature = "tier-2-multi-kms"), allow(dead_code))]
203#[inline]
204fn aes_gcm_nonce_from_counter(replica_id: u32, counter: u64) -> [u8; 12] {
205 let mut n = [0u8; 12];
206 n[0..4].copy_from_slice(&replica_id.to_be_bytes());
207 n[4..12].copy_from_slice(&counter.to_be_bytes());
208 n
209}
210
211impl core::fmt::Debug for Dek {
212 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
213 f.debug_struct("Dek").finish_non_exhaustive()
215 }
216}
217
218#[non_exhaustive]
225#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
226pub enum NonceBytes {
227 X24([u8; 24]),
229 Short12([u8; 12]),
231}
232
233impl NonceBytes {
234 #[inline]
238 #[must_use]
239 pub fn expected_len(kind: AeadKind) -> usize {
240 match kind {
241 AeadKind::XChaCha20Poly1305 => 24,
242 AeadKind::Aes256Gcm | AeadKind::Aes256GcmSiv => 12,
243 _ => 0,
244 }
245 }
246
247 #[inline]
249 #[must_use]
250 pub fn as_slice(&self) -> &[u8] {
251 match self {
252 Self::X24(b) => b,
253 Self::Short12(b) => b,
254 }
255 }
256}
257
258#[derive(Debug, PartialEq, Eq)]
271pub struct EncryptedPii<T: PiiType> {
272 pub dek_id: DekId,
274 pub pii_code: u16,
277 pub aead_kind: AeadKind,
279 pub nonce: NonceBytes,
281 pub ciphertext: Bytes,
283 pub(crate) _marker: PhantomData<fn() -> T>,
284}
285
286impl<T: PiiType> Clone for EncryptedPii<T> {
289 fn clone(&self) -> Self {
290 Self {
291 dek_id: self.dek_id,
292 pii_code: self.pii_code,
293 aead_kind: self.aead_kind,
294 nonce: self.nonce.clone(),
295 ciphertext: self.ciphertext.clone(),
296 _marker: PhantomData,
297 }
298 }
299}
300
301#[derive(Serialize, Deserialize)]
304struct EncryptedPiiWire {
305 dek_id: DekId,
306 pii_code: u16,
307 aead_kind: AeadKind,
308 nonce: NonceBytes,
309 ciphertext: Bytes,
310}
311
312impl<T: PiiType> Serialize for EncryptedPii<T> {
313 fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
314 EncryptedPiiWire {
315 dek_id: self.dek_id,
316 pii_code: self.pii_code,
317 aead_kind: self.aead_kind,
318 nonce: self.nonce.clone(),
319 ciphertext: self.ciphertext.clone(),
320 }
321 .serialize(serializer)
322 }
323}
324
325impl<'de, T: PiiType> Deserialize<'de> for EncryptedPii<T> {
326 fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
327 let wire = EncryptedPiiWire::deserialize(deserializer)?;
328 Ok(Self {
329 dek_id: wire.dek_id,
330 pii_code: wire.pii_code,
331 aead_kind: wire.aead_kind,
332 nonce: wire.nonce,
333 ciphertext: wire.ciphertext,
334 _marker: PhantomData,
335 })
336 }
337}
338
339impl<T: PiiType> EncryptedPii<T> {
340 #[inline]
345 #[must_use]
346 pub fn new(dek_id: DekId, aead_kind: AeadKind, nonce: NonceBytes, ciphertext: Bytes) -> Self {
347 Self {
348 dek_id,
349 pii_code: T::PII_CODE,
350 aead_kind,
351 nonce,
352 ciphertext,
353 _marker: PhantomData,
354 }
355 }
356
357 fn into_raw(self) -> RawEncryptedPii {
361 RawEncryptedPii {
362 dek_id: self.dek_id,
363 pii_code: self.pii_code,
364 aead_kind: self.aead_kind,
365 nonce: self.nonce,
366 ciphertext: self.ciphertext,
367 }
368 }
369}
370
371#[derive(Debug, Clone)]
375struct RawEncryptedPii {
376 dek_id: DekId,
377 pii_code: u16,
378 aead_kind: AeadKind,
379 nonce: NonceBytes,
380 ciphertext: Bytes,
381}
382
383#[derive(Debug)]
394pub struct CryptoCoordinator<N: NonceSource = OsNonceSource> {
395 manifest_cipher: AeadKind,
396 #[cfg_attr(not(feature = "tier-1-kms"), allow(dead_code))]
400 nonce_source: N,
401}
402
403pub trait NonceSource {
407 fn fill(&self, out: &mut [u8]);
412}
413
414#[derive(Debug, Default, Clone, Copy)]
420pub struct OsNonceSource;
421
422impl NonceSource for OsNonceSource {
423 #[cfg(feature = "tier-1-kms")]
424 fn fill(&self, out: &mut [u8]) {
425 use chacha20poly1305::aead::{rand_core::RngCore, OsRng};
426 OsRng.fill_bytes(out);
427 }
428
429 #[cfg(not(feature = "tier-1-kms"))]
430 fn fill(&self, out: &mut [u8]) {
431 for byte in out.iter_mut() {
434 *byte = 0;
435 }
436 }
437}
438
439impl<N: NonceSource> CryptoCoordinator<N> {
440 #[inline]
443 #[must_use]
444 pub fn new(manifest_cipher: AeadKind, nonce_source: N) -> Self {
445 Self {
446 manifest_cipher,
447 nonce_source,
448 }
449 }
450
451 #[inline]
453 #[must_use]
454 pub fn manifest_cipher(&self) -> AeadKind {
455 self.manifest_cipher
456 }
457
458 pub fn encrypt<T: PiiType>(
463 &self,
464 plaintext: &T,
465 dek: &Dek,
466 dek_id: DekId,
467 ) -> Result<EncryptedPii<T>, PiiError> {
468 let aad = compute_aad(&dek_id, T::PII_CODE, self.manifest_cipher);
469 let pt_bytes = postcard::to_stdvec(plaintext).map_err(|_| PiiError::EncryptFailed)?;
470 let (nonce, ciphertext) = self.encrypt_raw(dek, &aad, &pt_bytes)?;
471 Ok(EncryptedPii::new(
472 dek_id,
473 self.manifest_cipher,
474 nonce,
475 Bytes::from(ciphertext),
476 ))
477 }
478
479 pub fn decrypt<T: PiiType>(
483 &self,
484 envelope: &EncryptedPii<T>,
485 dek: &Dek,
486 ) -> Result<T, PiiError> {
487 if envelope.pii_code != T::PII_CODE {
488 return Err(PiiError::TypeMismatch);
489 }
490 if envelope.aead_kind != self.manifest_cipher {
491 return Err(PiiError::CipherDowngrade);
492 }
493 let aad = compute_aad(&envelope.dek_id, envelope.pii_code, envelope.aead_kind);
494 let pt = self.decrypt_raw(
495 dek,
496 envelope.aead_kind,
497 &envelope.nonce,
498 &aad,
499 &envelope.ciphertext,
500 )?;
501 postcard::from_bytes::<T>(&pt).map_err(|_| PiiError::DecodeFailed)
502 }
503
504 fn decrypt_raw_under(&self, dek: &Dek, raw: &RawEncryptedPii) -> Result<Vec<u8>, PiiError> {
511 let aad = compute_aad(&raw.dek_id, raw.pii_code, raw.aead_kind);
512 self.decrypt_raw(dek, raw.aead_kind, &raw.nonce, &aad, &raw.ciphertext)
513 }
514
515 fn encrypt_raw(
516 &self,
517 dek: &Dek,
518 aad: &[u8; 19],
519 plaintext: &[u8],
520 ) -> Result<(NonceBytes, Vec<u8>), PiiError> {
521 match self.manifest_cipher {
522 AeadKind::XChaCha20Poly1305 => self.encrypt_xchacha(dek, aad, plaintext),
523 AeadKind::Aes256Gcm => self.encrypt_aes_gcm(dek, aad, plaintext),
524 AeadKind::Aes256GcmSiv => self.encrypt_aes_gcm_siv(dek, aad, plaintext),
525 _ => Err(PiiError::UnsupportedAead),
526 }
527 }
528
529 fn decrypt_raw(
530 &self,
531 dek: &Dek,
532 kind: AeadKind,
533 nonce: &NonceBytes,
534 aad: &[u8; 19],
535 ciphertext: &[u8],
536 ) -> Result<Vec<u8>, PiiError> {
537 match kind {
538 AeadKind::XChaCha20Poly1305 => self.decrypt_xchacha(dek, nonce, aad, ciphertext),
539 AeadKind::Aes256Gcm => self.decrypt_aes_gcm(dek, nonce, aad, ciphertext),
540 AeadKind::Aes256GcmSiv => self.decrypt_aes_gcm_siv(dek, nonce, aad, ciphertext),
541 _ => Err(PiiError::UnsupportedAead),
542 }
543 }
544
545 #[cfg(feature = "tier-1-kms")]
548 fn encrypt_xchacha(
549 &self,
550 dek: &Dek,
551 aad: &[u8; 19],
552 plaintext: &[u8],
553 ) -> Result<(NonceBytes, Vec<u8>), PiiError> {
554 use chacha20poly1305::aead::{Aead, KeyInit, Payload};
555 use chacha20poly1305::{Key, XChaCha20Poly1305, XNonce};
556
557 let key = Key::from_slice(dek.as_bytes());
558 let cipher = XChaCha20Poly1305::new(key);
559 let mut nonce_buf = [0u8; 24];
560 self.nonce_source.fill(&mut nonce_buf);
561 let nonce = XNonce::from_slice(&nonce_buf);
562 let ciphertext = cipher
563 .encrypt(
564 nonce,
565 Payload {
566 msg: plaintext,
567 aad,
568 },
569 )
570 .map_err(|_| PiiError::EncryptFailed)?;
571 Ok((NonceBytes::X24(nonce_buf), ciphertext))
572 }
573
574 #[cfg(not(feature = "tier-1-kms"))]
575 fn encrypt_xchacha(
576 &self,
577 _dek: &Dek,
578 _aad: &[u8; 19],
579 _plaintext: &[u8],
580 ) -> Result<(NonceBytes, Vec<u8>), PiiError> {
581 Err(PiiError::TierTooLow)
582 }
583
584 #[cfg(feature = "tier-1-kms")]
585 fn decrypt_xchacha(
586 &self,
587 dek: &Dek,
588 nonce: &NonceBytes,
589 aad: &[u8; 19],
590 ciphertext: &[u8],
591 ) -> Result<Vec<u8>, PiiError> {
592 use chacha20poly1305::aead::{Aead, KeyInit, Payload};
593 use chacha20poly1305::{Key, XChaCha20Poly1305, XNonce};
594
595 let bytes_24 = match nonce {
596 NonceBytes::X24(b) => b,
597 NonceBytes::Short12(_) => return Err(PiiError::AadMismatch),
598 };
599 let key = Key::from_slice(dek.as_bytes());
600 let cipher = XChaCha20Poly1305::new(key);
601 let nonce = XNonce::from_slice(bytes_24);
602 cipher
603 .decrypt(
604 nonce,
605 Payload {
606 msg: ciphertext,
607 aad,
608 },
609 )
610 .map_err(|_| PiiError::AadMismatch)
611 }
612
613 #[cfg(not(feature = "tier-1-kms"))]
614 fn decrypt_xchacha(
615 &self,
616 _dek: &Dek,
617 _nonce: &NonceBytes,
618 _aad: &[u8; 19],
619 _ciphertext: &[u8],
620 ) -> Result<Vec<u8>, PiiError> {
621 Err(PiiError::TierTooLow)
622 }
623
624 #[cfg(feature = "tier-2-multi-kms")]
627 fn encrypt_aes_gcm(
628 &self,
629 dek: &Dek,
630 aad: &[u8; 19],
631 plaintext: &[u8],
632 ) -> Result<(NonceBytes, Vec<u8>), PiiError> {
633 use aes_gcm::aead::{Aead, KeyInit, Payload};
634 use aes_gcm::{Aes256Gcm, Key, Nonce};
635
636 let counter = dek.advance_counter()?;
640 let nonce_buf = aes_gcm_nonce_from_counter(dek.replica_id, counter);
641 let key = Key::<Aes256Gcm>::from_slice(dek.as_bytes());
642 let cipher = Aes256Gcm::new(key);
643 let nonce = Nonce::from_slice(&nonce_buf);
644 let ciphertext = cipher
645 .encrypt(
646 nonce,
647 Payload {
648 msg: plaintext,
649 aad,
650 },
651 )
652 .map_err(|_| PiiError::EncryptFailed)?;
653 Ok((NonceBytes::Short12(nonce_buf), ciphertext))
654 }
655
656 #[cfg(not(feature = "tier-2-multi-kms"))]
657 fn encrypt_aes_gcm(
658 &self,
659 _dek: &Dek,
660 _aad: &[u8; 19],
661 _plaintext: &[u8],
662 ) -> Result<(NonceBytes, Vec<u8>), PiiError> {
663 Err(PiiError::UnsupportedAead)
664 }
665
666 #[cfg(feature = "tier-2-multi-kms")]
667 fn decrypt_aes_gcm(
668 &self,
669 dek: &Dek,
670 nonce: &NonceBytes,
671 aad: &[u8; 19],
672 ciphertext: &[u8],
673 ) -> Result<Vec<u8>, PiiError> {
674 use aes_gcm::aead::{Aead, KeyInit, Payload};
675 use aes_gcm::{Aes256Gcm, Key, Nonce};
676
677 let bytes_12 = match nonce {
678 NonceBytes::Short12(b) => b,
679 NonceBytes::X24(_) => return Err(PiiError::AadMismatch),
680 };
681 let key = Key::<Aes256Gcm>::from_slice(dek.as_bytes());
682 let cipher = Aes256Gcm::new(key);
683 let nonce = Nonce::from_slice(bytes_12);
684 cipher
685 .decrypt(
686 nonce,
687 Payload {
688 msg: ciphertext,
689 aad,
690 },
691 )
692 .map_err(|_| PiiError::AadMismatch)
693 }
694
695 #[cfg(not(feature = "tier-2-multi-kms"))]
696 fn decrypt_aes_gcm(
697 &self,
698 _dek: &Dek,
699 _nonce: &NonceBytes,
700 _aad: &[u8; 19],
701 _ciphertext: &[u8],
702 ) -> Result<Vec<u8>, PiiError> {
703 Err(PiiError::UnsupportedAead)
704 }
705
706 #[cfg(feature = "tier-2-multi-kms")]
709 fn encrypt_aes_gcm_siv(
710 &self,
711 dek: &Dek,
712 aad: &[u8; 19],
713 plaintext: &[u8],
714 ) -> Result<(NonceBytes, Vec<u8>), PiiError> {
715 use aes_gcm_siv::aead::{Aead, KeyInit, Payload};
716 use aes_gcm_siv::{Aes256GcmSiv, Key, Nonce};
717
718 let counter = dek.advance_counter()?;
722 let nonce_buf = aes_gcm_nonce_from_counter(dek.replica_id, counter);
723 let key = Key::<Aes256GcmSiv>::from_slice(dek.as_bytes());
724 let cipher = Aes256GcmSiv::new(key);
725 let nonce = Nonce::from_slice(&nonce_buf);
726 let ciphertext = cipher
727 .encrypt(
728 nonce,
729 Payload {
730 msg: plaintext,
731 aad,
732 },
733 )
734 .map_err(|_| PiiError::EncryptFailed)?;
735 Ok((NonceBytes::Short12(nonce_buf), ciphertext))
736 }
737
738 #[cfg(not(feature = "tier-2-multi-kms"))]
739 fn encrypt_aes_gcm_siv(
740 &self,
741 _dek: &Dek,
742 _aad: &[u8; 19],
743 _plaintext: &[u8],
744 ) -> Result<(NonceBytes, Vec<u8>), PiiError> {
745 Err(PiiError::UnsupportedAead)
746 }
747
748 #[cfg(feature = "tier-2-multi-kms")]
749 fn decrypt_aes_gcm_siv(
750 &self,
751 dek: &Dek,
752 nonce: &NonceBytes,
753 aad: &[u8; 19],
754 ciphertext: &[u8],
755 ) -> Result<Vec<u8>, PiiError> {
756 use aes_gcm_siv::aead::{Aead, KeyInit, Payload};
757 use aes_gcm_siv::{Aes256GcmSiv, Key, Nonce};
758
759 let bytes_12 = match nonce {
760 NonceBytes::Short12(b) => b,
761 NonceBytes::X24(_) => return Err(PiiError::AadMismatch),
762 };
763 let key = Key::<Aes256GcmSiv>::from_slice(dek.as_bytes());
764 let cipher = Aes256GcmSiv::new(key);
765 let nonce = Nonce::from_slice(bytes_12);
766 cipher
767 .decrypt(
768 nonce,
769 Payload {
770 msg: ciphertext,
771 aad,
772 },
773 )
774 .map_err(|_| PiiError::AadMismatch)
775 }
776
777 #[cfg(not(feature = "tier-2-multi-kms"))]
778 fn decrypt_aes_gcm_siv(
779 &self,
780 _dek: &Dek,
781 _nonce: &NonceBytes,
782 _aad: &[u8; 19],
783 _ciphertext: &[u8],
784 ) -> Result<Vec<u8>, PiiError> {
785 Err(PiiError::UnsupportedAead)
786 }
787}
788
789pub fn rotate_dek<T: PiiType>(
801 coordinator: &CryptoCoordinator<impl NonceSource>,
802 old_dek: &Dek,
803 new_dek: &Dek,
804 new_dek_id: DekId,
805 ciphertexts: &mut [EncryptedPii<T>],
806 counter: &mut DekMessageCounter,
807) -> Result<(), PiiError> {
808 let originals: Vec<EncryptedPii<T>> = ciphertexts.to_vec();
809 for slot in ciphertexts.iter_mut() {
810 let raw = slot.clone().into_raw();
811 let plaintext_bytes = match coordinator.decrypt_raw_under(old_dek, &raw) {
812 Ok(v) => v,
813 Err(err) => {
814 for (target, backup) in ciphertexts.iter_mut().zip(originals.iter()) {
815 *target = backup.clone();
816 }
817 return Err(err);
818 }
819 };
820 let aad = compute_aad(&new_dek_id, T::PII_CODE, coordinator.manifest_cipher);
821 let (nonce, new_ct) = match coordinator.encrypt_raw(new_dek, &aad, &plaintext_bytes) {
822 Ok(v) => v,
823 Err(err) => {
824 for (target, backup) in ciphertexts.iter_mut().zip(originals.iter()) {
825 *target = backup.clone();
826 }
827 return Err(err);
828 }
829 };
830 *slot = EncryptedPii::new(
831 new_dek_id,
832 coordinator.manifest_cipher,
833 nonce,
834 Bytes::from(new_ct),
835 );
836 counter.record_message();
837 }
838 Ok(())
839}
840
841#[inline]
843#[must_use]
844pub fn rotation_advice(counter: &DekMessageCounter) -> RotationTrigger {
845 counter.rotation_trigger()
846}
847
848#[cfg(test)]
851#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
852mod tests {
853 use super::*;
854 use arkhe_forge_core::pii::ActorHandle;
855
856 #[derive(Clone, Copy, Default)]
857 struct FixedNonce;
858
859 impl NonceSource for FixedNonce {
860 fn fill(&self, out: &mut [u8]) {
861 for (i, byte) in out.iter_mut().enumerate() {
862 *byte = (i & 0xFF) as u8;
863 }
864 }
865 }
866
867 fn make_dek(byte: u8) -> Dek {
868 Dek::from_bytes([byte; 32])
869 }
870
871 fn make_dek_id(byte: u8) -> DekId {
872 DekId([byte; 16])
873 }
874
875 #[test]
876 fn dek_from_bytes_exposes_material_via_crate_accessor() {
877 let d = make_dek(0x42);
878 assert_eq!(d.as_bytes(), &[0x42u8; 32]);
879 }
880
881 #[test]
882 fn dek_try_from_slice_rejects_short_key() {
883 let err = Dek::try_from_slice(&[0u8; 16]).unwrap_err();
884 assert!(matches!(err, PiiError::InvalidKeyLength));
885 }
886
887 #[test]
888 fn dek_try_from_slice_accepts_32_bytes() {
889 let key = [0x77u8; 32];
890 let dek = Dek::try_from_slice(&key).unwrap();
891 assert_eq!(dek.as_bytes(), &key);
892 }
893
894 #[test]
895 fn dek_debug_does_not_expose_material() {
896 let d = make_dek(0xAB);
897 let s = format!("{:?}", d);
898 assert!(!s.contains("AB"), "Debug output must not leak key bytes");
899 assert!(!s.contains("ab"));
900 }
901
902 #[test]
903 fn nonce_bytes_expected_len_matches_kind() {
904 assert_eq!(NonceBytes::expected_len(AeadKind::XChaCha20Poly1305), 24);
905 assert_eq!(NonceBytes::expected_len(AeadKind::Aes256Gcm), 12);
906 assert_eq!(NonceBytes::expected_len(AeadKind::Aes256GcmSiv), 12);
907 }
908
909 #[test]
910 fn encrypted_pii_wire_layout_roundtrips_through_postcard() {
911 let envelope = EncryptedPii::<ActorHandle>::new(
912 make_dek_id(0x11),
913 AeadKind::XChaCha20Poly1305,
914 NonceBytes::X24([0x22; 24]),
915 Bytes::from_static(&[0x33; 48]),
916 );
917 let bytes = postcard::to_stdvec(&envelope).unwrap();
918 let back: EncryptedPii<ActorHandle> = postcard::from_bytes(&bytes).unwrap();
919 assert_eq!(envelope, back);
920 assert_eq!(back.pii_code, ActorHandle::PII_CODE);
921 }
922
923 #[cfg(not(feature = "tier-1-kms"))]
926 #[test]
927 fn tier0_default_rejects_encryption() {
928 let coord = CryptoCoordinator::new(AeadKind::XChaCha20Poly1305, FixedNonce);
929 let err = coord
930 .encrypt::<ActorHandle>(
931 &ActorHandle(b"alice".to_vec()),
932 &make_dek(0x00),
933 make_dek_id(0x11),
934 )
935 .unwrap_err();
936 assert!(matches!(err, PiiError::TierTooLow));
937 }
938
939 #[cfg(feature = "tier-1-kms")]
942 #[test]
943 fn tier1_xchacha_encrypt_decrypt_roundtrip() {
944 let coord = CryptoCoordinator::new(AeadKind::XChaCha20Poly1305, FixedNonce);
945 let handle = ActorHandle(b"alice".to_vec());
946 let dek = make_dek(0xA5);
947 let dek_id = make_dek_id(0x11);
948 let env = coord.encrypt(&handle, &dek, dek_id).unwrap();
949 assert_eq!(env.aead_kind, AeadKind::XChaCha20Poly1305);
950 assert_eq!(env.pii_code, ActorHandle::PII_CODE);
951 let back: ActorHandle = coord.decrypt(&env, &dek).unwrap();
952 assert_eq!(back, handle);
953 }
954
955 #[cfg(feature = "tier-1-kms")]
956 #[test]
957 fn tier1_xchacha_aad_tamper_fails_tag() {
958 let coord = CryptoCoordinator::new(AeadKind::XChaCha20Poly1305, FixedNonce);
959 let handle = ActorHandle(b"alice".to_vec());
960 let dek = make_dek(0x01);
961 let mut env = coord.encrypt(&handle, &dek, make_dek_id(0x11)).unwrap();
962 env.dek_id = make_dek_id(0x12);
965 let err = coord.decrypt::<ActorHandle>(&env, &dek).unwrap_err();
966 assert!(matches!(err, PiiError::AadMismatch));
967 }
968
969 #[cfg(feature = "tier-1-kms")]
970 #[test]
971 fn tier1_ciphertext_tamper_fails_tag() {
972 let coord = CryptoCoordinator::new(AeadKind::XChaCha20Poly1305, FixedNonce);
973 let handle = ActorHandle(b"alice".to_vec());
974 let dek = make_dek(0x03);
975 let env = coord.encrypt(&handle, &dek, make_dek_id(0x11)).unwrap();
976 let mut ct = env.ciphertext.to_vec();
977 if let Some(first) = ct.first_mut() {
978 *first ^= 0x01;
979 }
980 let tampered =
981 EncryptedPii::<ActorHandle>::new(env.dek_id, env.aead_kind, env.nonce, Bytes::from(ct));
982 let err = coord.decrypt::<ActorHandle>(&tampered, &dek).unwrap_err();
983 assert!(matches!(err, PiiError::AadMismatch));
984 }
985
986 #[cfg(feature = "tier-1-kms")]
987 #[test]
988 fn tier1_wrong_pii_code_rejected_as_type_mismatch() {
989 let coord = CryptoCoordinator::new(AeadKind::XChaCha20Poly1305, FixedNonce);
990 let handle = ActorHandle(b"alice".to_vec());
991 let dek = make_dek(0x07);
992 let env = coord.encrypt(&handle, &dek, make_dek_id(0x11)).unwrap();
993 let wrong = EncryptedPii::<ActorHandle> {
996 dek_id: env.dek_id,
997 pii_code: arkhe_forge_core::pii::EntryBody::PII_CODE,
998 aead_kind: env.aead_kind,
999 nonce: env.nonce,
1000 ciphertext: env.ciphertext,
1001 _marker: PhantomData,
1002 };
1003 let err = coord.decrypt::<ActorHandle>(&wrong, &dek).unwrap_err();
1004 assert!(matches!(err, PiiError::TypeMismatch));
1005 }
1006
1007 #[cfg(feature = "tier-1-kms")]
1008 #[test]
1009 fn tier1_aead_downgrade_rejected_by_coordinator_manifest() {
1010 let coord = CryptoCoordinator::new(AeadKind::XChaCha20Poly1305, FixedNonce);
1013 let env = EncryptedPii::<ActorHandle>::new(
1014 make_dek_id(0x11),
1015 AeadKind::Aes256Gcm,
1016 NonceBytes::Short12([0u8; 12]),
1017 Bytes::from_static(&[0u8; 48]),
1018 );
1019 let err = coord
1020 .decrypt::<ActorHandle>(&env, &make_dek(0x00))
1021 .unwrap_err();
1022 assert!(matches!(err, PiiError::CipherDowngrade));
1023 }
1024
1025 #[cfg(feature = "tier-1-kms")]
1026 #[test]
1027 fn tier1_aes_gcm_without_tier2_is_unsupported() {
1028 let coord = CryptoCoordinator::new(AeadKind::Aes256Gcm, FixedNonce);
1031 let handle = ActorHandle(b"alice".to_vec());
1032 let out = coord.encrypt(&handle, &make_dek(0x00), make_dek_id(0x11));
1033 #[cfg(feature = "tier-2-multi-kms")]
1034 assert!(out.is_ok());
1035 #[cfg(not(feature = "tier-2-multi-kms"))]
1036 assert!(matches!(out, Err(PiiError::UnsupportedAead)));
1037 }
1038
1039 #[cfg(feature = "tier-2-multi-kms")]
1040 #[test]
1041 fn tier2_aes_gcm_roundtrip() {
1042 let coord = CryptoCoordinator::new(AeadKind::Aes256Gcm, FixedNonce);
1043 let handle = ActorHandle(b"aes-user".to_vec());
1044 let dek = make_dek(0x5A);
1045 let env = coord.encrypt(&handle, &dek, make_dek_id(0x21)).unwrap();
1046 assert_eq!(env.aead_kind, AeadKind::Aes256Gcm);
1047 assert!(matches!(env.nonce, NonceBytes::Short12(_)));
1048 let back: ActorHandle = coord.decrypt(&env, &dek).unwrap();
1049 assert_eq!(back, handle);
1050 }
1051
1052 #[cfg(feature = "tier-2-multi-kms")]
1053 #[test]
1054 fn tier2_aes_gcm_siv_roundtrip() {
1055 let coord = CryptoCoordinator::new(AeadKind::Aes256GcmSiv, FixedNonce);
1056 let handle = ActorHandle(b"aes-siv-user".to_vec());
1057 let dek = make_dek(0x7B);
1058 let env = coord.encrypt(&handle, &dek, make_dek_id(0x22)).unwrap();
1059 assert_eq!(env.aead_kind, AeadKind::Aes256GcmSiv);
1060 let back: ActorHandle = coord.decrypt(&env, &dek).unwrap();
1061 assert_eq!(back, handle);
1062 }
1063
1064 #[cfg(feature = "tier-2-multi-kms")]
1065 #[test]
1066 fn aes_gcm_nonce_is_deterministic_counter() {
1067 let coord = CryptoCoordinator::new(AeadKind::Aes256Gcm, FixedNonce);
1071 let dek = make_dek(0x5A);
1072 let handle = ActorHandle(b"alice".to_vec());
1073
1074 let env1 = coord.encrypt(&handle, &dek, make_dek_id(0x11)).unwrap();
1075 let env2 = coord.encrypt(&handle, &dek, make_dek_id(0x11)).unwrap();
1076
1077 let NonceBytes::Short12(n1) = &env1.nonce else {
1078 panic!("AES-GCM always returns Short12");
1079 };
1080 let NonceBytes::Short12(n2) = &env2.nonce else {
1081 panic!("AES-GCM always returns Short12");
1082 };
1083 assert_eq!(&n1[0..4], &[0u8; 4], "invocation field zeros");
1084 assert_eq!(&n1[4..12], &0u64.to_be_bytes());
1085 assert_eq!(&n2[4..12], &1u64.to_be_bytes());
1086 assert_ne!(n1, n2);
1087
1088 assert_eq!(coord.decrypt::<ActorHandle>(&env1, &dek).unwrap(), handle);
1090 assert_eq!(coord.decrypt::<ActorHandle>(&env2, &dek).unwrap(), handle);
1091 }
1092
1093 #[cfg(feature = "tier-2-multi-kms")]
1094 #[test]
1095 fn aes_gcm_nonce_honours_dek_replica_id() {
1096 let coord = CryptoCoordinator::new(AeadKind::Aes256Gcm, FixedNonce);
1101 let dek_a = Dek::with_config([0xC3; 32], DekConfig { replica_id: 0 });
1102 let dek_b = Dek::with_config(
1103 [0xC3; 32],
1104 DekConfig {
1105 replica_id: 0xDEAD_BEEF,
1106 },
1107 );
1108 let handle = ActorHandle(b"alice".to_vec());
1109
1110 let env_a = coord.encrypt(&handle, &dek_a, make_dek_id(0x11)).unwrap();
1111 let env_b = coord.encrypt(&handle, &dek_b, make_dek_id(0x11)).unwrap();
1112
1113 let NonceBytes::Short12(na) = &env_a.nonce else {
1114 panic!("AES-GCM returns Short12");
1115 };
1116 let NonceBytes::Short12(nb) = &env_b.nonce else {
1117 panic!("AES-GCM returns Short12");
1118 };
1119 assert_eq!(&na[0..4], &0u32.to_be_bytes());
1120 assert_eq!(&nb[0..4], &0xDEAD_BEEFu32.to_be_bytes());
1121 assert_eq!(&na[4..12], &0u64.to_be_bytes());
1124 assert_eq!(&nb[4..12], &0u64.to_be_bytes());
1125 assert_ne!(na, nb);
1126 }
1127
1128 #[cfg(feature = "tier-2-multi-kms")]
1129 #[test]
1130 fn aes_gcm_siv_nonce_is_deterministic_counter() {
1131 let coord = CryptoCoordinator::new(AeadKind::Aes256GcmSiv, FixedNonce);
1134 let dek = make_dek(0x7B);
1135 let handle = ActorHandle(b"siv".to_vec());
1136
1137 let env1 = coord.encrypt(&handle, &dek, make_dek_id(0x22)).unwrap();
1138 let NonceBytes::Short12(n1) = &env1.nonce else {
1139 panic!("AES-GCM-SIV returns Short12");
1140 };
1141 assert_eq!(&n1[4..12], &0u64.to_be_bytes());
1142 assert_eq!(dek.get_counter_for_test(), 1);
1143 }
1144
1145 #[cfg(feature = "tier-2-multi-kms")]
1146 #[test]
1147 fn dek_counter_exhaustion_errors() {
1148 let coord = CryptoCoordinator::new(AeadKind::Aes256Gcm, FixedNonce);
1151 let dek = make_dek(0xA5);
1152 dek.set_counter_for_test(u64::MAX);
1153
1154 let handle = ActorHandle(b"alice".to_vec());
1155 let err = coord.encrypt(&handle, &dek, make_dek_id(0x11)).unwrap_err();
1156 assert!(matches!(err, PiiError::DekExhausted));
1157 assert_eq!(dek.get_counter_for_test(), u64::MAX);
1159 }
1160
1161 #[cfg(feature = "tier-2-multi-kms")]
1162 #[test]
1163 fn rotate_dek_starts_new_counter_from_zero() {
1164 let coord = CryptoCoordinator::new(AeadKind::Aes256Gcm, FixedNonce);
1168 let old = make_dek(0x10);
1169 let new = make_dek(0x20);
1170 let new_id = make_dek_id(0x02);
1171
1172 let plaintexts: Vec<ActorHandle> = (0..3u8).map(|i| ActorHandle(vec![i; 8])).collect();
1173 let mut envs: Vec<EncryptedPii<ActorHandle>> = plaintexts
1174 .iter()
1175 .map(|pt| coord.encrypt(pt, &old, make_dek_id(0x01)).unwrap())
1176 .collect();
1177 assert_eq!(old.get_counter_for_test(), 3);
1178 assert_eq!(new.get_counter_for_test(), 0);
1179
1180 let mut rotation_metric = DekMessageCounter::new(new_id);
1181 rotate_dek(&coord, &old, &new, new_id, &mut envs, &mut rotation_metric).unwrap();
1182
1183 assert_eq!(new.get_counter_for_test(), 3);
1184 assert_eq!(rotation_metric.count(), 3);
1185 for (i, env) in envs.iter().enumerate() {
1186 let NonceBytes::Short12(n) = &env.nonce else {
1187 panic!("AES-GCM returns Short12");
1188 };
1189 assert_eq!(
1190 &n[4..12],
1191 &(i as u64).to_be_bytes(),
1192 "counter values run 0,1,2 under new DEK"
1193 );
1194 }
1195 }
1196
1197 #[cfg(feature = "tier-1-kms")]
1198 #[test]
1199 fn dek_rotate_preserves_plaintext() {
1200 let coord = CryptoCoordinator::new(AeadKind::XChaCha20Poly1305, FixedNonce);
1201 let old = make_dek(0x10);
1202 let new = make_dek(0x20);
1203 let new_id = make_dek_id(0x02);
1204 let plaintexts: Vec<ActorHandle> = (0..4u8).map(|i| ActorHandle(vec![i; 8])).collect();
1205 let mut envelopes: Vec<EncryptedPii<ActorHandle>> = plaintexts
1206 .iter()
1207 .map(|pt| coord.encrypt(pt, &old, make_dek_id(0x01)).unwrap())
1208 .collect();
1209 let mut counter = DekMessageCounter::new(make_dek_id(0x02));
1210 rotate_dek(&coord, &old, &new, new_id, &mut envelopes, &mut counter).unwrap();
1211 assert_eq!(counter.count(), 4);
1212 for (env, pt) in envelopes.iter().zip(plaintexts.iter()) {
1213 assert_eq!(env.dek_id, new_id);
1214 assert_eq!(&coord.decrypt::<ActorHandle>(env, &new).unwrap(), pt);
1215 }
1216 }
1217
1218 #[cfg(feature = "tier-1-kms")]
1219 #[test]
1220 fn dek_rotate_with_wrong_old_key_rolls_back() {
1221 let coord = CryptoCoordinator::new(AeadKind::XChaCha20Poly1305, FixedNonce);
1222 let real_old = make_dek(0x10);
1223 let wrong_old = make_dek(0xFF);
1224 let new = make_dek(0x20);
1225 let original_envelope = coord
1226 .encrypt(
1227 &ActorHandle(b"alice".to_vec()),
1228 &real_old,
1229 make_dek_id(0x01),
1230 )
1231 .unwrap();
1232 let mut envelopes = vec![original_envelope.clone()];
1233 let mut counter = DekMessageCounter::new(make_dek_id(0x02));
1234 let err = rotate_dek(
1235 &coord,
1236 &wrong_old,
1237 &new,
1238 make_dek_id(0x02),
1239 &mut envelopes,
1240 &mut counter,
1241 )
1242 .unwrap_err();
1243 assert!(matches!(err, PiiError::AadMismatch));
1244 assert_eq!(envelopes[0], original_envelope);
1246 }
1247}