Skip to main content

arkhe_forge_platform/
crypto.rs

1//! Crypto-erasure coordinator — Tier-1+ AEAD envelope encryption.
2//!
3//! Provides:
4//!
5//! * [`Dek`] — 32-byte key material with zeroise-on-drop semantics. The
6//!   runtime never handles wrapped key material; that lives in the
7//!   HSM/KMS backend.
8//! * [`EncryptedPii`] — generic wrapper over an opaque ciphertext + tag +
9//!   nonce + `DekId` + `AeadKind`. The wire tag binds the PII marker
10//!   (`T::PII_CODE`) via AAD, defeating the type-confused-deputy path.
11//! * [`CryptoCoordinator`] — stateful entry-point that dispatches
12//!   `encrypt` / `decrypt` by the shell manifest's declared `AeadKind`
13//!   and compliance tier. Under the default (Tier-0) feature set the
14//!   coordinator refuses encryption with [`PiiError::TierTooLow`].
15//! * [`rotate_dek`] — slice-level DEK rotation helper. Callers must hold
16//!   a single-writer lock; the helper is atomic per-element and rolls
17//!   the whole slice back on the first failure.
18//!
19//! Feature matrix:
20//!
21//! | Feature                | `XChaCha20-Poly1305` | `AES-256-GCM` | `AES-256-GCM-SIV` |
22//! |------------------------|----------------------|---------------|-------------------|
23//! | *(default — Tier-0)*   | rejected             | rejected      | rejected          |
24//! | `tier-1-kms`           | ✓                    | rejected      | rejected          |
25//! | `tier-2-multi-kms`     | ✓                    | ✓             | ✓                 |
26//!
27//! The coordinator's public surface is stable. HSM / KMS wrap-unwrap
28//! integration and the Sigstore transparency anchor route through `hf2_kms`.
29
30use 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// ===================== Dek =====================
40
41/// Construction-time configuration for a [`Dek`]. Single-writer
42/// deployments use the default (all fields zero); federation builds
43/// populate `replica_id` from the per-instance manifest anchor so two
44/// regions sharing the same DEK material cannot collide their
45/// deterministic nonces (the F6 invocation-field reservation).
46#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
47pub struct DekConfig {
48    /// 4-byte invocation field for the AES-GCM(-SIV) nonce. Must be
49    /// instance-pinned immutable (L0 A11 pure compute); Runtime
50    /// reconfig **must not** mutate this value without a full
51    /// RuntimeBootstrap re-emit and DEK rotation cycle.
52    pub replica_id: u32,
53}
54
55/// Per-user 32-byte DEK material. The byte buffer is
56/// wiped on `Drop` via the `zeroize` crate; callers obtain a `Dek` from
57/// an HSM unwrap — the runtime never derives key material directly
58/// (envelope encryption).
59///
60/// A per-DEK monotonic counter drives the deterministic 96-bit nonce
61/// for AES-GCM / AES-GCM-SIV under the NIST SP 800-38D §8.2.1
62/// construction (4-byte invocation field = [`DekConfig::replica_id`] +
63/// 8-byte counter, big-endian). XChaCha20-Poly1305 uses a 192-bit
64/// random nonce and leaves the counter alone. Counter exhaustion at
65/// `u64::MAX` surfaces [`PiiError::DekExhausted`] so the operator
66/// rotates the DEK before any nonce reuse.
67///
68/// `replica_id` is held by the DEK itself and is **immutable after
69/// construction** — changing it post-hoc would violate L0 A11 pure
70/// compute (replay determinism depends on stable nonce bytes). Default
71/// single-writer deployments use `replica_id = 0`; federation builds
72/// populate a per-replica id via [`Dek::with_config`].
73///
74/// `Dek` is intentionally **not** `Clone` — two copies of the same key
75/// material with independent counters would collide their nonces under
76/// AES-GCM (catastrophic integrity loss). Callers must hold a single
77/// owner per key; rotation yields a fresh `Dek` via [`Dek::from_bytes`]
78/// whose counter starts at `0`.
79///
80/// The interior `Cell<u64>` counter makes `Dek` implicitly `!Sync`, so
81/// the compiler refuses naive cross-thread sharing. If a deployment
82/// genuinely needs to share a DEK across task boundaries (e.g., an
83/// L2 async runtime feeding multiple encrypt handlers), wrap it in
84/// `Arc<Mutex<Dek>>` or `Arc<parking_lot::Mutex<Dek>>` at the caller
85/// site — the mutex guards the counter advance and keeps the
86/// deterministic-nonce invariant intact. Single-writer deployments
87/// (L0 A2 single-thread) hold a plain owned `Dek` without
88/// synchronisation.
89#[derive(Zeroize, ZeroizeOnDrop)]
90pub struct Dek {
91    material: [u8; 32],
92    /// Monotonic message counter used by the AES-GCM(-SIV) nonce
93    /// construction. Not sensitive; skipped by zeroize.
94    #[zeroize(skip)]
95    counter: Cell<u64>,
96    /// Invocation field for the nonce construction. Immutable after
97    /// construction — see [`DekConfig::replica_id`].
98    #[zeroize(skip)]
99    replica_id: u32,
100}
101
102impl Dek {
103    /// Construct from a 32-byte buffer using default configuration
104    /// (`replica_id = 0`, single-writer). The input is copied; callers
105    /// remain responsible for wiping their own buffer. Counter starts
106    /// at `0`.
107    #[inline]
108    #[must_use]
109    pub fn from_bytes(material: [u8; 32]) -> Self {
110        Self::with_config(material, DekConfig::default())
111    }
112
113    /// Construct from a 32-byte buffer with an explicit [`DekConfig`].
114    /// Federation path consumes this entry point with a non-zero
115    /// `replica_id`; single-writer deployments use [`Dek::from_bytes`].
116    #[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    /// Construct from a byte slice. Rejects the call with
127    /// [`PiiError::InvalidKeyLength`] when `bytes.len() != 32`.
128    /// Counter starts at `0` with default [`DekConfig`].
129    ///
130    /// The length check is a single `usize` compare against the
131    /// constant `32` — no byte-by-byte value comparison happens on
132    /// the reject path, so there is no timing side-channel the
133    /// `subtle` crate would mitigate. `copy_from_slice` runs in
134    /// time dependent on the buffer length, not its contents.
135    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    /// Borrow the underlying 32-byte key. Intentionally crate-visible so
149    /// downstream crypto primitives can feed the AEAD `Key` type without
150    /// leaking the buffer through the public surface.
151    #[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    /// Return the current counter value and advance to the next. Used
162    /// by the AES-GCM(-SIV) encrypt paths to produce deterministic
163    /// nonces. Returns [`PiiError::DekExhausted`] when the counter has
164    /// already reached `u64::MAX` — a rotation is required before any
165    /// further encryption.
166    #[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    /// Test-only accessor — set the counter to an arbitrary value to
177    /// exercise exhaustion and rotation paths.
178    #[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    /// Test-only accessor — read the current counter.
184    #[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/// NIST SP 800-38D §8.2.1 deterministic nonce: 4-byte invocation
191/// field (`replica_id`) + 8-byte big-endian counter. The 96-bit
192/// layout is the native nonce size for both AES-GCM and AES-GCM-SIV.
193///
194/// `replica_id = 0` is the single-writer default. Federation builds
195/// supply a per-replica id via [`DekConfig::replica_id`] so two
196/// regions sharing the same DEK material cannot collide their
197/// counters. The value is captured at `Dek` construction and is
198/// immutable for the lifetime of the key (L0 A11 pure compute);
199/// changing it later would surface in the WAL as different ciphertext
200/// bytes and therefore must ship behind a manifest-anchored policy
201/// event when the federation activation lands.
202#[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        // Never print key material.
214        f.debug_struct("Dek").finish_non_exhaustive()
215    }
216}
217
218// ===================== EncryptedPii =====================
219
220/// Per-AEAD-kind nonce carrier — XChaCha20-Poly1305 uses 24 bytes, the
221/// AES-256-GCM family uses 12. A single variant enum keeps the on-wire
222/// layout round-trip stable under postcard. Variants use postcard's
223/// default untagged discriminant (enum index byte).
224#[non_exhaustive]
225#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
226pub enum NonceBytes {
227    /// 24-byte nonce — `XChaCha20-Poly1305`.
228    X24([u8; 24]),
229    /// 12-byte nonce — `AES-256-GCM` / `AES-256-GCM-SIV`.
230    Short12([u8; 12]),
231}
232
233impl NonceBytes {
234    /// Expected length for the given `AeadKind`. Forward-compat unknown
235    /// variants map to `0` so callers can catch the mismatch as
236    /// [`PiiError::UnsupportedAead`].
237    #[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    /// Returns the nonce bytes as a slice regardless of variant.
248    #[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/// Per-PII-marker ciphertext envelope.
259///
260/// The wire shape is
261/// `(dek_id, pii_code, aead_kind, nonce, ciphertext_with_tag)` —
262/// every input to the AEAD AAD is mirrored on the envelope so the
263/// receiver can recompute the 19-byte AAD exactly. `ciphertext`
264/// includes the 16-byte Poly1305 / GCM tag appended by the AEAD
265/// primitive.
266///
267/// The generic parameter `T` is a *phantom* — the wire layout is purely
268/// data-bearing, and a manual (de)serialize impl threads around the
269/// `PhantomData` so postcard can round-trip the struct.
270#[derive(Debug, PartialEq, Eq)]
271pub struct EncryptedPii<T: PiiType> {
272    /// HSM/KMS key reference.
273    pub dek_id: DekId,
274    /// Wire tag. Validated against `T::PII_CODE` at
275    /// decrypt time.
276    pub pii_code: u16,
277    /// AEAD family used for the ciphertext.
278    pub aead_kind: AeadKind,
279    /// Nonce — length varies per AEAD kind.
280    pub nonce: NonceBytes,
281    /// Ciphertext with the 16-byte AEAD tag appended.
282    pub ciphertext: Bytes,
283    pub(crate) _marker: PhantomData<fn() -> T>,
284}
285
286// Manual Clone — `T` carries no data on the envelope (the PhantomData is
287// `fn() -> T` for variance), so cloning never calls into `T::clone`.
288impl<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/// Data-only mirror of [`EncryptedPii`] — used as the serde surface so
302/// the generic phantom does not reach the wire.
303#[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    /// Construct from components — intended for deserialization paths /
341    /// tests that need to reassemble a postcard-decoded envelope into
342    /// its typed form. Encryption path callers use
343    /// [`CryptoCoordinator::encrypt`].
344    #[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    /// Non-typed copy of the envelope — used when the caller must move
358    /// between two `T`s during DEK rotation without re-encoding the
359    /// wire layout.
360    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/// Type-erased [`EncryptedPii`] view — used internally by the DEK
372/// rotation helper to avoid a second generic bound during the
373/// decrypt-then-re-encrypt swap.
374#[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// ===================== CryptoCoordinator =====================
384
385/// Tier-1+ AEAD coordinator.
386///
387/// The coordinator carries the shell-declared `AeadKind` (from
388/// `[audit.pii_cipher]`) plus an opaque `nonce_source` the caller
389/// supplies so tests and production paths can share the same
390/// dispatcher. Tier-0 (default feature set) consumers may still
391/// instantiate a coordinator — every mutating call returns
392/// [`PiiError::TierTooLow`].
393#[derive(Debug)]
394pub struct CryptoCoordinator<N: NonceSource = OsNonceSource> {
395    manifest_cipher: AeadKind,
396    // Only consumed by the XChaCha20-Poly1305 (tier-1-kms) path — the
397    // AES-GCM(-SIV) paths use a DEK-local deterministic counter (NIST
398    // SP 800-38D §8.2.1) rather than the nonce source.
399    #[cfg_attr(not(feature = "tier-1-kms"), allow(dead_code))]
400    nonce_source: N,
401}
402
403/// Per-kind nonce generator. Production uses [`OsNonceSource`] (the
404/// `getrandom` syscall wrapper baked into the AEAD crates). Tests plug
405/// in a fixed value for bit-identical fixtures.
406pub trait NonceSource {
407    /// Fill `out` with `len` fresh nonce bytes. `len` matches
408    /// [`NonceBytes::expected_len`] for the active AEAD kind; panics /
409    /// short writes are illegal — return an error via the caller's
410    /// wrapper if needed.
411    fn fill(&self, out: &mut [u8]);
412}
413
414/// OS-backed nonce source. On `tier-1-kms` / `tier-2-multi-kms` this
415/// pulls from the underlying AEAD crate's default RNG hook (itself
416/// `getrandom`-backed). On the default feature set it is still
417/// constructible — encryption paths reject before the nonce source is
418/// consulted.
419#[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        // No crypto feature active — zero-fill. Encrypt paths reject
432        // before the nonce is read, so the value is never observed.
433        for byte in out.iter_mut() {
434            *byte = 0;
435        }
436    }
437}
438
439impl<N: NonceSource> CryptoCoordinator<N> {
440    /// Construct a coordinator that enforces `manifest_cipher` at the
441    /// encrypt / decrypt boundary.
442    #[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    /// Borrow the manifest-declared cipher.
452    #[inline]
453    #[must_use]
454    pub fn manifest_cipher(&self) -> AeadKind {
455        self.manifest_cipher
456    }
457
458    /// AEAD-encrypt a PII payload under the given DEK. The 19-byte AAD
459    /// is computed from `(dek_id, T::PII_CODE, manifest_cipher)` so a
460    /// downstream decrypt call with a tampered envelope fails the tag
461    /// check.
462    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    /// Inverse of [`CryptoCoordinator::encrypt`]. Checks the wire PII
480    /// marker, refuses cipher downgrades, then verifies the AEAD tag
481    /// against a recomputed AAD.
482    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    /// Variant of [`CryptoCoordinator::decrypt`] that tolerates a
505    /// legacy `AeadKind` — used by the DEK rotation path so a slice of
506    /// ciphertexts written under an older manifest `pii_cipher` can be
507    /// migrated in place. Caller supplies the historical kind; the
508    /// manifest downgrade check is bypassed (equivalence relation
509    /// anchored to the envelope's own `aead_kind`).
510    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    // ----- XChaCha20-Poly1305 — tier-1-kms gated -----
546
547    #[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    // ----- AES-256-GCM — tier-2-multi-kms gated -----
625
626    #[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        // Deterministic counter nonce (NIST SP 800-38D §8.2.1). Counter
637        // exhaustion is surfaced by `advance_counter` before any AEAD
638        // work starts, so a failed call never consumes a nonce value.
639        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    // ----- AES-256-GCM-SIV — tier-2-multi-kms gated -----
707
708    #[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        // Deterministic counter nonce — AES-GCM-SIV already tolerates
719        // nonce reuse but the counter construction removes the chance
720        // entirely under single-writer (L0 A2) semantics.
721        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
789// ===================== DEK rotation =====================
790
791/// Re-wrap every element of `ciphertexts` under `new_dek` using a fresh
792/// `new_dek_id`. Decrypts under `old_dek` first, re-encrypts under the
793/// new key material, and rolls the slice back if any element fails
794/// (atomic-per-call semantics).
795///
796/// Callers must hold a single-writer lock across the slice while this
797/// helper runs — the coordinator does not expose its own synchronisation.
798/// The update also ticks `counter` (by the slice length) so the operator
799/// can observe the per-DEK rotation metric.
800pub 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/// Helper — extract the rotation trigger for a post-rotation counter.
842#[inline]
843#[must_use]
844pub fn rotation_advice(counter: &DekMessageCounter) -> RotationTrigger {
845    counter.rotation_trigger()
846}
847
848// ===================== Tests =====================
849
850#[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    // --- Default (Tier-0) — encryption is rejected across the board. ---
924
925    #[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    // --- tier-1-kms — XChaCha20-Poly1305 round trip. ---
940
941    #[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        // Tamper with the dek_id on the envelope — AAD recompute changes,
963        // tag verification must fail.
964        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        // Overwrite pii_code with a different marker. decrypt::<ActorHandle>
994        // spots the mismatch without touching AEAD.
995        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        // Coordinator is pinned to XChaCha20-Poly1305; envelope is written
1011        // with AES-GCM. decrypt must refuse.
1012        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        // Coordinator is pinned to AES-GCM; under tier-1 only, encrypt
1029        // must surface UnsupportedAead.
1030        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        // First two encrypts under the same DEK must produce nonces 0
1068        // and 1 in the 8-byte big-endian counter tail; the 4-byte
1069        // invocation field is zero for the single-writer deployment.
1070        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        // Round-trip still holds under the new construction.
1089        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        // Federation path — a non-zero `replica_id` must appear as the
1097        // 4-byte invocation field prefix of the nonce. Two DEKs with
1098        // identical material but distinct replica ids produce
1099        // disjoint nonce spaces (0-counter slot differs).
1100        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        // Counter portion is the same (both at position 0) but the
1122        // full nonce differs because of the replica id prefix.
1123        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        // SIV path also consumes the DEK counter — defence-in-depth
1132        // despite AES-GCM-SIV's own nonce-reuse resistance.
1133        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        // Counter at u64::MAX rejects further encryption — no nonce is
1149        // consumed, operator must rotate.
1150        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        // Counter unchanged — failed call did not advance.
1158        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        // `rotate_dek` decrypts under the old DEK (no counter use) and
1165        // re-encrypts under the new DEK (counter advances 0..N). After
1166        // rotating N elements the new DEK's counter reflects exactly N.
1167        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        // Slice rolled back — original envelope intact.
1245        assert_eq!(envelopes[0], original_envelope);
1246    }
1247}