diff --git a/packages/aead/src/lib.rs b/packages/aead/src/lib.rs index 2718bb7..f5909a6 100644 --- a/packages/aead/src/lib.rs +++ b/packages/aead/src/lib.rs @@ -9,6 +9,7 @@ mod encrypt; #[cfg(feature = "hlist")] pub mod hlist; mod nonce; +mod tree; #[doc(inline)] pub use aad::{Aad, IntoAad}; @@ -24,3 +25,8 @@ pub use decipher::{Decipher, DecipherVisitor, Decrypt, MapAccess, SeqAccess}; pub use encrypt::Encrypt; #[doc(inline)] pub use nonce::{Nonce, NonceGenerator, RandomNonceGenerator}; +#[doc(inline)] +pub use tree::{ + CipherTree, LeafOpener, PendingLeaf, TreeCipher, TreeDecipher, TreeMap, TreeMapAccess, TreeSeq, + TreeSeqAccess, +}; diff --git a/packages/aead/src/nonce.rs b/packages/aead/src/nonce.rs index 9d5ea35..fe38f70 100644 --- a/packages/aead/src/nonce.rs +++ b/packages/aead/src/nonce.rs @@ -19,6 +19,15 @@ impl Nonce { pub fn into_inner(self) -> [u8; N] { self.0 } + + /// Borrow the nonce as a fixed-size array. Unlike + /// [`into_inner`](Self::into_inner) this does not consume the nonce, so the + /// same nonce can be used for the AEAD call and then appended to the + /// ciphertext. Copy the result (`*nonce.as_array()`) when an owned array is + /// needed — no fallible slice conversion required. + pub fn as_array(&self) -> &[u8; N] { + &self.0 + } } // TODO: Make this a trait that can be implemented for a cipher rather than an associated type on the Cipher diff --git a/packages/aead/src/tree.rs b/packages/aead/src/tree.rs new file mode 100644 index 0000000..9f55ba6 --- /dev/null +++ b/packages/aead/src/tree.rs @@ -0,0 +1,527 @@ +//! A reusable recursive ciphertext tree, a structural encrypt codec, and the +//! generic [`Decipher`] machinery that drives it. +//! +//! Most ciphers that produce a structured ciphertext (a single sealed value, a +//! sequence, a map, …) share the *same* structural traversal — only the per-leaf +//! sealing/opening differs. This module factors that traversal out: +//! +//! - [`CipherTree`] is the recursive container, generic over the per-leaf +//! payload `L`. +//! - [`TreeCipher`] implements [`Cipher`]. It does **not** seal anything: it +//! walks the plaintext structure and collects each byte leaf — together with +//! the AAD bound at that position — into a [`CipherTree`]. Sealing +//! is a separate, cipher-specific step. +//! - [`TreeDecipher`] implements [`Decipher`] for any [`LeafOpener`] `O`; it +//! drives the [`DecipherVisitor`] walk and calls [`LeafOpener::open`] at each +//! leaf. +//! +//! ## Why sealing is a separate step (encrypt side) +//! +//! Splitting *structure* from *sealing* is what lets one codec serve both an +//! eager, single-key cipher and a batching one: +//! +//! - `vitaminc_encrypt::Aes256Cipher` runs [`TreeCipher`] to build the pending +//! tree, then seals it leaf-by-leaf under fresh random nonces (see its +//! `seal_tree`). +//! - A ZeroKMS-backed cipher can build the same pending tree, call +//! [`CipherTree::leaf_count`] to size a single batched data-key request, then +//! [`CipherTree::try_map_leaves`] to seal every leaf with the fetched keys. +//! +//! Both reuse [`TreeCipher`] verbatim; only the seal step differs. The earlier +//! `LeafSealer` design called a seal hook eagerly at each leaf, which a batching +//! cipher could not satisfy without buffering plaintext behind a no-op `seal` +//! and a second pass anyway — so the buffering is made explicit here instead. +//! +//! The trade-off of deferring the seal is that the whole plaintext structure is +//! resident (inside [`Protected`], so still zeroized on drop) as a +//! [`CipherTree`] before sealing, rather than each leaf being sealed +//! the instant it is visited. For the small structured values this is used for, +//! that transient residency is negligible. +//! +//! ## Leaf ordering contract +//! +//! [`CipherTree::leaf_count`], [`CipherTree::for_each_leaf`], and +//! [`CipherTree::try_map_leaves`] all visit keyed leaves in the **same** +//! depth-first, left-to-right order. A batching cipher relies on this: it +//! collects per-leaf material with `for_each_leaf` (or sizes a request with +//! `leaf_count`), fetches keys in that order, then re-seals positionally with +//! `try_map_leaves`. The alignment is pinned by a test (`leaf_order_is_stable`). +//! +//! ## Coherence +//! +//! [`TreeDecipher`]'s generic impl is on the concrete local type, **not** a +//! blanket `impl Decipher for &C`. A conditional blanket would conflict +//! (E0119) with any direct `Decipher` impl, and that trait is a public extension +//! point meant to be implemented directly. Concrete ciphers wrap themselves in a +//! [`LeafOpener`] instead. + +use std::any::Any; + +use vitaminc_protected::Protected; + +use crate::{ + Cipher, Decipher, DecipherVisitor, Decrypt, Encrypt, IntoAad, MapAccess, MapCipher, SeqAccess, + SeqCipher, Unspecified, +}; + +/// A recursive ciphertext container whose shape mirrors the encrypted plaintext. +/// +/// Generic over the per-leaf payload `L` (e.g. a [`PendingLeaf`] awaiting a data +/// key, or a sealed `LocalCipherText`). +#[derive(Debug)] +pub enum CipherTree { + /// A single sealed value. + Single(L), + /// A sequence of ciphertexts (from a `Vec`-shaped plaintext). + Sequence(Vec>), + /// A map of (cleartext key, ciphertext value) pairs. Keys are not encrypted. + Map(Vec<(String, CipherTree)>), + /// The authenticated absent marker produced by [`Cipher::encrypt_none`]. + None(L), + /// A typed value passed through unencrypted via [`Cipher::passthrough`]. + /// + /// The value is type-erased into a boxed [`Any`]; the decrypt side recovers + /// it with a runtime downcast (see [`Decipher::decrypt_passthrough`]). This + /// representation is **not** serializable or inspectable: a cipher that needs + /// to transmit the tree over a wire (e.g. for batch transport) will need a + /// typed/serializable passthrough variant instead. Passthrough is for + /// non-sensitive in-process data only — see [`Cipher::passthrough`]. + Passthrough(Box), +} + +impl CipherTree { + /// Number of keyed leaves ([`Single`](Self::Single) + [`None`](Self::None)); + /// passthrough nodes carry no leaf payload. + pub fn leaf_count(&self) -> usize { + match self { + CipherTree::Single(_) | CipherTree::None(_) => 1, + CipherTree::Sequence(items) => items.iter().map(Self::leaf_count).sum(), + CipherTree::Map(entries) => entries.iter().map(|(_, v)| v.leaf_count()).sum(), + CipherTree::Passthrough(_) => 0, + } + } + + /// Visit every keyed leaf in depth-first, left-to-right order. This is the + /// canonical leaf order — [`try_map_leaves`](Self::try_map_leaves) consumes + /// leaves in the same order, so collectors and sealers line up (see the + /// module-level "Leaf ordering contract"). + pub fn for_each_leaf<'s>(&'s self, f: &mut impl FnMut(&'s L)) { + match self { + CipherTree::Single(leaf) | CipherTree::None(leaf) => f(leaf), + CipherTree::Sequence(items) => items.iter().for_each(|i| i.for_each_leaf(f)), + CipherTree::Map(entries) => entries.iter().for_each(|(_, v)| v.for_each_leaf(f)), + CipherTree::Passthrough(_) => {} + } + } + + /// Fallibly transform every keyed leaf, preserving structure (and passthrough + /// nodes). Leaves are visited in [`for_each_leaf`](Self::for_each_leaf) order, + /// so a closure that draws from an ordered key source stays aligned. The + /// first `Err` short-circuits. + pub fn try_map_leaves( + self, + f: &mut impl FnMut(L) -> Result, + ) -> Result, E> { + Ok(match self { + CipherTree::Single(leaf) => CipherTree::Single(f(leaf)?), + CipherTree::None(leaf) => CipherTree::None(f(leaf)?), + CipherTree::Sequence(items) => CipherTree::Sequence( + items + .into_iter() + .map(|item| item.try_map_leaves(f)) + .collect::, E>>()?, + ), + CipherTree::Map(entries) => CipherTree::Map( + entries + .into_iter() + .map(|(k, v)| v.try_map_leaves(f).map(|v| (k, v))) + .collect::, E>>()?, + ), + CipherTree::Passthrough(value) => CipherTree::Passthrough(value), + }) + } +} + +/// A byte leaf collected by [`TreeCipher`], awaiting sealing. +/// +/// Carries the plaintext (kept inside [`Protected`] so the zeroize-on-drop +/// guarantee survives until the leaf is sealed or dropped) and the AAD bytes +/// bound at this position in the structure. A cipher's seal step consumes a +/// `PendingLeaf` and produces its own sealed leaf type. +pub struct PendingLeaf { + /// The plaintext bytes to seal. For an [`encrypt_none`](Cipher::encrypt_none) + /// marker this is the empty vector (sealed so its tag still binds the AAD). + pub plaintext: Protected>, + /// The fully-encoded AAD bytes the leaf must be authenticated against. + pub aad: Vec, +} + +impl PendingLeaf { + fn new<'a, A: IntoAad<'a>>(plaintext: Protected>, aad: A) -> Self { + Self { + plaintext, + aad: aad.into_aad().as_bytes().to_vec(), + } + } +} + +// ============================================================================= +// Encrypt side — a structural codec that collects, but does not seal. +// ============================================================================= + +/// A [`Cipher`] that walks a plaintext structure into a [`CipherTree`] +/// **without sealing**. Run this to get the tree shape and per-leaf plaintext+AAD, +/// then seal the leaves with a cipher-specific step (see the module docs). +pub struct TreeCipher; + +impl Cipher for TreeCipher { + type Ok = CipherTree; + type Error = Unspecified; + type SeqCipher = TreeSeq; + type MapCipher = TreeMap; + + fn encrypt_bytes_vec<'a, A>( + self, + data: Protected>, + aad: A, + ) -> Result + where + A: IntoAad<'a>, + { + Ok(CipherTree::Single(PendingLeaf::new(data, aad))) + } + + fn encrypt_seq(self, size_hint: Option) -> Self::SeqCipher { + TreeSeq { + items: Vec::with_capacity(size_hint.unwrap_or(0)), + } + } + + fn encrypt_map(self) -> Self::MapCipher { + TreeMap { + entries: Vec::new(), + current_key: None, + } + } + + fn encrypt_none<'a, A>(self, aad: A) -> Result + where + A: IntoAad<'a>, + { + // The marker seals an empty plaintext so its tag still binds the AAD. + Ok(CipherTree::None(PendingLeaf::new( + Protected::new(Vec::new()), + aad, + ))) + } + + fn passthrough(self, value: T) -> Result + where + T: Any + Send + 'static, + { + Ok(CipherTree::Passthrough(Box::new(value))) + } +} + +/// [`SeqCipher`] for [`TreeCipher`]: accumulates a pending [`CipherTree`] per element. +pub struct TreeSeq { + items: Vec>, +} + +impl SeqCipher for TreeSeq { + type Ok = CipherTree; + type Error = Unspecified; + + fn encrypt_next<'a, T, A>(mut self, data: T, aad: A) -> Result + where + T: Encrypt, + A: IntoAad<'a>, + { + // Recurse with a fresh structural codec — the element is collected, not sealed. + self.items.push(data.encrypt_with_aad(TreeCipher, aad)?); + Ok(self) + } + + fn passthrough_next(mut self, value: T) -> Result + where + T: Any + Send + 'static, + { + self.items.push(CipherTree::Passthrough(Box::new(value))); + Ok(self) + } + + fn end(self) -> Result { + Ok(CipherTree::Sequence(self.items)) + } +} + +/// [`MapCipher`] for [`TreeCipher`]: keys are stored in the clear; values become +/// sub-trees. Enforces the key→value call contract like a hand-written impl. +pub struct TreeMap { + entries: Vec<(String, CipherTree)>, + current_key: Option<&'static str>, +} + +impl MapCipher for TreeMap { + type Ok = CipherTree; + type Error = Unspecified; + + fn encrypt_key(mut self, key: &'static str) -> Result { + // Two keys with no intervening value is a contract violation. + if self.current_key.is_some() { + return Err(Unspecified); + } + self.current_key = Some(key); + Ok(self) + } + + fn encrypt_value<'a, T, A>(mut self, value: T, aad: A) -> Result + where + T: Encrypt, + A: IntoAad<'a>, + { + let key = self.current_key.take().ok_or(Unspecified)?; + let value = value.encrypt_with_aad(TreeCipher, aad)?; + self.entries.push((key.to_string(), value)); + Ok(self) + } + + fn passthrough_entry(mut self, key: &'static str, value: T) -> Result + where + T: Any + Send + 'static, + { + if self.current_key.is_some() { + return Err(Unspecified); + } + self.entries + .push((key.to_string(), CipherTree::Passthrough(Box::new(value)))); + Ok(self) + } + + fn end(self) -> Result { + if self.current_key.is_some() { + return Err(Unspecified); + } + Ok(CipherTree::Map(self.entries)) + } +} + +// ============================================================================= +// Decrypt side +// ============================================================================= + +/// Opens one leaf back to plaintext. The single per-cipher operation on the +/// decrypt side. +/// +/// `Copy` because openers are stateless lookups — each leaf self-identifies the +/// key it needs (e.g. via a stored nonce, or an iv+tag a batching cipher resolves +/// against a pre-fetched key map), so the same opener can be threaded through +/// every leaf of the structure without ownership churn. +pub trait LeafOpener: Copy { + /// The per-leaf payload this opener consumes (a cipher's sealed leaf type). + type Leaf; + + /// Open `leaf`, authenticating against `aad`, returning the plaintext bytes. + fn open(self, leaf: Self::Leaf, aad: &[u8]) -> Result>, Unspecified>; +} + +/// A [`Decipher`] over a [`CipherTree`] that delegates each leaf to a +/// [`LeafOpener`]. Synchronous: `Ok = Result`. +pub struct TreeDecipher { + tree: CipherTree, + opener: O, +} + +impl TreeDecipher { + /// Build a decipher that opens `tree`'s leaves with `opener`. + pub fn new(opener: O, tree: CipherTree) -> Self { + Self { tree, opener } + } +} + +impl<'c, O: LeafOpener> Decipher<'c> for TreeDecipher { + type Ok + = Result + where + T: Send + 'c; + + fn map_ok(ok: Self::Ok, f: F) -> Self::Ok + where + T: Send + 'c, + U: Send + 'c, + F: FnOnce(T) -> U, + { + ok.map(f) + } + + fn decrypt_bytes<'a, V, A>(self, visitor: V, aad: A) -> Self::Ok + where + V: DecipherVisitor<'c> + Send + 'c, + A: IntoAad<'a>, + { + match self.tree { + CipherTree::Single(leaf) => { + let aad = aad.into_aad(); + visitor.visit_bytes_vec(self.opener.open(leaf, aad.as_bytes())?) + } + _ => Err(Unspecified), + } + } + + fn decrypt_seq<'a, V, A>(self, visitor: V, aad: A) -> Self::Ok + where + V: DecipherVisitor<'c> + Send + 'c, + A: IntoAad<'a>, + { + match self.tree { + CipherTree::Sequence(items) => visitor.visit_seq(TreeSeqAccess { + items: items.into_iter(), + opener: self.opener, + aad: aad.into_aad(), + }), + _ => Err(Unspecified), + } + } + + fn decrypt_map<'a, V, A>(self, visitor: V, aad: A) -> Self::Ok + where + V: DecipherVisitor<'c> + Send + 'c, + A: IntoAad<'a>, + { + match self.tree { + CipherTree::Map(entries) => visitor.visit_map(TreeMapAccess { + entries: entries.into_iter(), + opener: self.opener, + aad: aad.into_aad(), + }), + _ => Err(Unspecified), + } + } + + fn decrypt_passthrough(self) -> Self::Ok + where + T: Any + Send + 'static, + { + match self.tree { + CipherTree::Passthrough(boxed) => { + boxed.downcast::().map(|b| *b).map_err(|_| Unspecified) + } + _ => Err(Unspecified), + } + } + + fn decrypt_option<'a, T, A>(self, aad: A) -> Self::Ok> + where + T: Decrypt<'c> + 'c, + A: IntoAad<'a>, + { + match self.tree { + CipherTree::None(leaf) => { + // Authenticate the absent marker (verifies the tag) and discard. + let aad = aad.into_aad(); + self.opener.open(leaf, aad.as_bytes())?; + Ok(None) + } + CipherTree::Passthrough(_) => Err(Unspecified), + // Any other shape is the `Some` payload — recurse into `T`. + other => T::decrypt_with_aad(TreeDecipher::new(self.opener, other), aad).map(Some), + } + } +} + +/// [`SeqAccess`] for [`TreeDecipher`]. Re-applies the sequence AAD to every +/// element by borrowing the stored bytes (no per-element allocation). +pub struct TreeSeqAccess<'a, O: LeafOpener> { + items: std::vec::IntoIter>, + opener: O, + aad: crate::Aad<'a>, +} + +impl<'c, 'a, O: LeafOpener> SeqAccess<'c> for TreeSeqAccess<'a, O> { + type Error = Unspecified; + + fn next_element + 'c>(&mut self) -> Result, Self::Error> { + match self.items.next() { + Some(tree) => { + T::decrypt_with_aad(TreeDecipher::new(self.opener, tree), self.aad.as_bytes()) + .map(Some) + } + None => Ok(None), + } + } +} + +/// [`MapAccess`] for [`TreeDecipher`]. Keys are returned in the clear; each value +/// is decrypted under the stored AAD. +pub struct TreeMapAccess<'a, O: LeafOpener> { + entries: std::vec::IntoIter<(String, CipherTree)>, + opener: O, + aad: crate::Aad<'a>, +} + +impl<'c, 'a, O: LeafOpener> MapAccess<'c> for TreeMapAccess<'a, O> { + type Error = Unspecified; + + fn next_entry + 'c>(&mut self) -> Result, Self::Error> { + match self.entries.next() { + Some((key, tree)) => { + let value = + T::decrypt_with_aad(TreeDecipher::new(self.opener, tree), self.aad.as_bytes())?; + Ok(Some((key, value))) + } + None => Ok(None), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Pins the module's "Leaf ordering contract": `for_each_leaf`, + /// `try_map_leaves`, and `leaf_count` must agree on which leaves exist and in + /// what order. A batching cipher collects per-leaf material in `for_each_leaf` + /// order and re-seals positionally with `try_map_leaves`; if the two ever + /// diverged, leaf *i* would be sealed under the key fetched for leaf *j*. This + /// test would catch that before it became a silent key/leaf misalignment. + #[test] + fn leaf_order_is_stable() { + // A mixed tree: a map containing a passthrough (no leaf), a single, and a + // nested sequence — exercises every variant of the traversal. + let tree: CipherTree = CipherTree::Map(vec![ + ("pt".to_string(), CipherTree::Passthrough(Box::new(0u8))), + ("a".to_string(), CipherTree::Single(1)), + ( + "seq".to_string(), + CipherTree::Sequence(vec![ + CipherTree::Single(2), + CipherTree::None(3), + CipherTree::Single(4), + ]), + ), + ]); + + // Passthrough carries no leaf, so 4 keyed leaves: 1, 2, 3, 4. + assert_eq!(tree.leaf_count(), 4); + + let mut visited = Vec::new(); + tree.for_each_leaf(&mut |leaf: &u32| visited.push(*leaf)); + assert_eq!(visited, vec![1, 2, 3, 4]); + + // try_map_leaves consumes leaves in the *same* order: tag each with its + // visit index and confirm the structure-preserving result lines up. + let mut idx = 0u32; + let mapped: CipherTree<(u32, u32)> = tree + .try_map_leaves(&mut |leaf| { + let pos = idx; + idx += 1; + Ok::<_, ()>((leaf, pos)) + }) + .expect("infallible"); + + let mut pairs = Vec::new(); + mapped.for_each_leaf(&mut |(leaf, pos): &(u32, u32)| pairs.push((*leaf, *pos))); + // (leaf value, position) — positions follow the same 0..4 leaf order. + assert_eq!(pairs, vec![(1, 0), (2, 1), (3, 2), (4, 3)]); + } +} diff --git a/packages/encrypt/src/cipher.rs b/packages/encrypt/src/cipher.rs index 79af310..3a7f98e 100644 --- a/packages/encrypt/src/cipher.rs +++ b/packages/encrypt/src/cipher.rs @@ -2,37 +2,21 @@ use crate::backend::{CipherKey, NONCE_LEN}; use crate::Key; use std::any::Any; use vitaminc_aead::{ - Aad, Cipher, CipherTextBuilder, Decipher, DecipherVisitor, Decrypt, Encrypt, IntoAad, - LocalCipherText, MapAccess, MapCipher, NonceGenerator, RandomNonceGenerator, SeqAccess, - SeqCipher, Unspecified, + Cipher, CipherTextBuilder, CipherTree, Decrypt, IntoAad, LeafOpener, LocalCipherText, + MapCipher, NonceGenerator, PendingLeaf, RandomNonceGenerator, SeqCipher, TreeCipher, + TreeDecipher, TreeMap, TreeSeq, Unspecified, }; use vitaminc_protected::Protected; /// The recursive ciphertext container produced by [`Aes256Cipher`]. /// -/// The shape mirrors the structure of the plaintext that was encrypted: a -/// single value yields [`Single`](AesCipherText::Single), a `Vec` yields -/// [`Sequence`](AesCipherText::Sequence), a `HashMap` yields -/// [`Map`](AesCipherText::Map). Nested structures are represented recursively. -#[derive(Debug)] -pub enum AesCipherText { - /// A single sealed value (nonce + ciphertext + tag). - Single(LocalCipherText), - /// A sequence of ciphertexts produced from a `Vec`-shaped plaintext. - Sequence(Vec), - /// A map of (cleartext key, ciphertext value) pairs produced from a - /// `HashMap`-shaped plaintext. Keys are not encrypted. - Map(Vec<(String, AesCipherText)>), - /// The authenticated absent marker produced by [`Cipher::encrypt_none`]. - /// Stores a sealed empty plaintext whose tag binds the supplied AAD. - None(LocalCipherText), - /// A typed value passed through unencrypted via [`Cipher::passthrough`]. - /// Not serializable to bytes — see - /// `packages/aead/src/cipher.rs` for the API-level type bound and the - /// runtime downcast performed by - /// [`Decipher::decrypt_passthrough`](vitaminc_aead::Decipher::decrypt_passthrough). - Passthrough(Box), -} +/// A [`CipherTree`] whose leaves are sealed [`LocalCipherText`]s. Its shape +/// mirrors the encrypted plaintext: a single value yields +/// [`Single`](CipherTree::Single), a `Vec` yields +/// [`Sequence`](CipherTree::Sequence), a `HashMap` yields +/// [`Map`](CipherTree::Map). The structural drive/walk is shared with any other +/// cipher built on [`CipherTree`]; only the per-leaf sealing differs. +pub type AesCipherText = CipherTree; /// Implements AES-256-GCM. Backend is selected at compile time: /// `aws-lc-rs` on native targets, `aes-gcm` (RustCrypto) on `wasm32`. @@ -53,13 +37,71 @@ impl Aes256Cipher { key: key.cipher_key()?, }) } + + /// Seal a single [`PendingLeaf`] under a fresh random nonce. This is the AES + /// per-leaf step the generic [`TreeCipher`] codec leaves to the cipher. + fn seal_leaf(&self, leaf: PendingLeaf) -> Result { + let nonce = self.nonce_generator.generate()?; + // Copy the nonce bytes out without consuming the nonce (it is still + // appended to the ciphertext below) and without a fallible slice + // conversion — `Nonce` already wraps `[u8; NONCE_LEN]`. + let nonce_bytes = *nonce.as_array(); + + CipherTextBuilder::new() + .append_nonce(nonce) + .append_target_plaintext(leaf.plaintext) + .accepts_ciphertext_and_tag_ok(|mut buf| { + self.key + .seal(&nonce_bytes, &leaf.aad, &mut buf) + .map(|()| buf) + }) + .build() + } + + /// Seal every leaf of a pending tree, preserving structure. AES seals + /// eagerly leaf-by-leaf; a batching cipher would instead size a key request + /// with [`CipherTree::leaf_count`] and pull keys here. Leaves are sealed in + /// [`CipherTree::for_each_leaf`] order (see the `vitaminc_aead::tree` + /// module's "Leaf ordering contract"). + fn seal_tree(&self, tree: CipherTree) -> Result { + tree.try_map_leaves(&mut |leaf| self.seal_leaf(leaf)) + } +} + +/// The per-leaf opening step for [`Aes256Cipher`]: decrypt a leaf with the +/// cipher's key and the nonce stored in the leaf. A `Copy` lookup wrapper so the +/// generic [`TreeDecipher`] can thread it through every leaf of the structure. +#[derive(Clone, Copy)] +pub struct AesOpener<'c>(&'c Aes256Cipher); + +impl LeafOpener for AesOpener<'_> { + type Leaf = LocalCipherText; + + fn open(self, leaf: LocalCipherText, aad: &[u8]) -> Result>, Unspecified> { + let (nonce, reader) = leaf.into_reader().read_nonce::()?; + let nonce_bytes = nonce.into_inner(); + reader + .accepts_plaintext_ok(|data| self.0.key.open(&nonce_bytes, aad, data)) + .read() + } } +/// A [`Decipher`](vitaminc_aead::Decipher) over an [`AesCipherText`], produced by +/// [`Aes256Cipher::decipher`]. The structural walk is shared via [`TreeDecipher`]; +/// [`AesOpener`] supplies the per-leaf decryption. +pub type AesDecipher<'c> = TreeDecipher>; + +/// `&Aes256Cipher` is the [`Cipher`] so values can be encrypted ergonomically as +/// `value.encrypt(&cipher)`. The structural traversal is delegated to the generic +/// [`TreeCipher`] codec (which only *collects* leaves); this impl adds the AES +/// seal step, sealing the collected tree before handing it back. A coherence +/// blanket impl is impossible (it would conflict with every direct `Cipher` +/// impl), so a future backend repeats this thin collect-then-seal shape. impl<'c> Cipher for &'c Aes256Cipher { type Ok = AesCipherText; type Error = Unspecified; - type SeqCipher = AesSeqCipher<'c>; - type MapCipher = AesMapCipher<'c>; + type SeqCipher = AesSeq<'c>; + type MapCipher = AesMap<'c>; fn encrypt_bytes_vec<'a, A>( self, @@ -69,34 +111,20 @@ impl<'c> Cipher for &'c Aes256Cipher { where A: IntoAad<'a>, { - let nonce = self.nonce_generator.generate()?; - let nonce_bytes: [u8; NONCE_LEN] = nonce.as_ref().try_into().map_err(|_| Unspecified)?; - let aad = aad.into_aad(); - - CipherTextBuilder::new() - .append_nonce(nonce) - .append_target_plaintext(data) - .accepts_ciphertext_and_tag_ok(|mut buf| { - self.key - .seal(&nonce_bytes, aad.as_bytes(), &mut buf) - .map(|()| buf) - }) - .build() - .map(AesCipherText::Single) + self.seal_tree(TreeCipher.encrypt_bytes_vec(data, aad)?) } fn encrypt_seq(self, size_hint: Option) -> Self::SeqCipher { - AesSeqCipher { + AesSeq { + inner: TreeCipher.encrypt_seq(size_hint), cipher: self, - items: Vec::with_capacity(size_hint.unwrap_or(0)), } } fn encrypt_map(self) -> Self::MapCipher { - AesMapCipher { + AesMap { + inner: TreeCipher.encrypt_map(), cipher: self, - entries: Vec::new(), - current_key: None, } } @@ -104,49 +132,34 @@ impl<'c> Cipher for &'c Aes256Cipher { where A: IntoAad<'a>, { - let nonce = self.nonce_generator.generate()?; - let nonce_bytes: [u8; NONCE_LEN] = nonce.as_ref().try_into().map_err(|_| Unspecified)?; - let aad = aad.into_aad(); - - CipherTextBuilder::new() - .append_nonce(nonce) - .append_target_plaintext(Vec::::new()) - .accepts_ciphertext_and_tag_ok(|mut buf| { - self.key - .seal(&nonce_bytes, aad.as_bytes(), &mut buf) - .map(|()| buf) - }) - .build() - .map(AesCipherText::None) + self.seal_tree(TreeCipher.encrypt_none(aad)?) } fn passthrough(self, value: T) -> Result where T: Any + Send + 'static, { - Ok(AesCipherText::Passthrough(Box::new(value))) + self.seal_tree(TreeCipher.passthrough(value)?) } } -/// [`SeqCipher`] driver for [`Aes256Cipher`]. Encrypts each element under its -/// own fresh nonce and accumulates the results into an -/// [`AesCipherText::Sequence`]. -pub struct AesSeqCipher<'c> { +/// [`SeqCipher`] for `&Aes256Cipher`: accumulates a pending tree via the generic +/// [`TreeSeq`], then seals the whole sequence in one pass at [`end`](SeqCipher::end). +pub struct AesSeq<'c> { + inner: TreeSeq, cipher: &'c Aes256Cipher, - items: Vec, } -impl<'c> SeqCipher for AesSeqCipher<'c> { +impl SeqCipher for AesSeq<'_> { type Ok = AesCipherText; type Error = Unspecified; fn encrypt_next<'a, T, A>(mut self, data: T, aad: A) -> Result where - T: Encrypt, + T: vitaminc_aead::Encrypt, A: IntoAad<'a>, { - let encrypted = data.encrypt_with_aad(self.cipher, aad)?; - self.items.push(encrypted); + self.inner = self.inner.encrypt_next(data, aad)?; Ok(self) } @@ -154,59 +167,38 @@ impl<'c> SeqCipher for AesSeqCipher<'c> { where T: Any + Send + 'static, { - self.items.push(AesCipherText::Passthrough(Box::new(value))); + self.inner = self.inner.passthrough_next(value)?; Ok(self) } fn end(self) -> Result { - Ok(AesCipherText::Sequence(self.items)) + self.cipher.seal_tree(self.inner.end()?) } } -/// [`MapCipher`] driver for [`Aes256Cipher`]. Keys are stored in the clear; -/// values are encrypted under their own fresh nonce and accumulated into an -/// [`AesCipherText::Map`]. -/// -/// This driver is intended for encrypting sources that already enforce key -/// uniqueness themselves — `HashMap`s and structs (whose field names are -/// unique by construction). It therefore stores `entries` as a *positional -/// list* and performs **no duplicate-key checks** of its own. -/// -/// Implementors should be aware: if the same key is pushed by two completed -/// `encrypt_value` / `passthrough_entry` calls, both are kept, and on decrypt -/// the `HashMap`-shaped visitor is **last-wins**. The built-in -/// `Encrypt for HashMap` impl never does this (the source map already dedups), -/// so it is unreachable today; a custom encoder driving this trait directly is -/// responsible for not emitting duplicate keys. -pub struct AesMapCipher<'c> { +/// [`MapCipher`] for `&Aes256Cipher`: accumulates a pending tree via the generic +/// [`TreeMap`] (which enforces the key→value contract), then seals at +/// [`end`](MapCipher::end). +pub struct AesMap<'c> { + inner: TreeMap, cipher: &'c Aes256Cipher, - entries: Vec<(String, AesCipherText)>, - current_key: Option<&'static str>, } -impl<'c> MapCipher for AesMapCipher<'c> { +impl MapCipher for AesMap<'_> { type Ok = AesCipherText; type Error = Unspecified; fn encrypt_key(mut self, key: &'static str) -> Result { - // A key already pending means `encrypt_key` was called twice with no - // intervening `encrypt_value` — a trait-contract violation. Fail rather - // than silently drop the first key. - if self.current_key.is_some() { - return Err(Unspecified); - } - self.current_key = Some(key); + self.inner = self.inner.encrypt_key(key)?; Ok(self) } - fn encrypt_value<'a, U, A>(mut self, value: U, aad: A) -> Result + fn encrypt_value<'a, T, A>(mut self, value: T, aad: A) -> Result where - U: Encrypt, + T: vitaminc_aead::Encrypt, A: IntoAad<'a>, { - let key = self.current_key.take().ok_or(Unspecified)?; - let encrypted = value.encrypt_with_aad(self.cipher, aad)?; - self.entries.push((key.to_string(), encrypted)); + self.inner = self.inner.encrypt_value(value, aad)?; Ok(self) } @@ -214,24 +206,12 @@ impl<'c> MapCipher for AesMapCipher<'c> { where T: Any + Send + 'static, { - // A key already pending means `encrypt_key` ran without a matching - // `encrypt_value` — adopting it here would silently drop the pending - // key, which is the same trait-contract violation `encrypt_key` - // rejects. - if self.current_key.is_some() { - return Err(Unspecified); - } - self.entries - .push((key.to_string(), AesCipherText::Passthrough(Box::new(value)))); + self.inner = self.inner.passthrough_entry(key, value)?; Ok(self) } fn end(self) -> Result { - // Finalising with a pending key would silently drop the entry. - if self.current_key.is_some() { - return Err(Unspecified); - } - Ok(AesCipherText::Map(self.entries)) + self.cipher.seal_tree(self.inner.end()?) } } @@ -264,202 +244,19 @@ impl Aes256Cipher { T::decrypt_with_aad(self.decipher(ciphertext), aad) } - /// Construct a [`Decipher`] over `ciphertext` bound to this cipher. + /// Construct a [`Decipher`](vitaminc_aead::Decipher) over `ciphertext` bound + /// to this cipher. /// /// This is the decrypt-side counterpart to passing `&cipher` (a [`Cipher`]) - /// on the encrypt side: it hands callers a concrete [`Decipher`] they can - /// drive directly via [`Decrypt::decrypt_with_aad`], which is what generic + /// on the encrypt side: it hands callers a concrete decipher they can drive + /// directly via [`Decrypt::decrypt_with_aad`], which is what generic /// decrypt-side helpers (e.g. `ContextTag`) build on. The ergonomic /// [`decrypt`](Aes256Cipher::decrypt) / /// [`decrypt_with_aad`](Aes256Cipher::decrypt_with_aad) methods are thin - /// wrappers around it. + /// wrappers around it. The structural walk lives in [`TreeDecipher`]; + /// [`AesOpener`] supplies the per-leaf decryption. pub fn decipher(&self, ciphertext: AesCipherText) -> AesDecipher<'_> { - AesDecipher { - cipher: self, - ciphertext, - } - } -} - -/// A [`Decipher`] over a single [`AesCipherText`], produced by -/// [`Aes256Cipher::decipher`]. Carries the cipher and ciphertext; the AAD is -/// supplied per call by [`Decrypt::decrypt_with_aad`]. -pub struct AesDecipher<'c> { - cipher: &'c Aes256Cipher, - ciphertext: AesCipherText, -} - -impl AesDecipher<'_> { - fn decrypt_local_ciphertext( - cipher: &Aes256Cipher, - ct: LocalCipherText, - aad: &[u8], - ) -> Result>, Unspecified> { - let (nonce, reader) = ct.into_reader().read_nonce::()?; - let nonce_bytes = nonce.into_inner(); - - reader - .accepts_plaintext_ok(|data| cipher.key.open(&nonce_bytes, aad, data)) - .read() - } -} - -impl<'c> Decipher<'c> for AesDecipher<'c> { - type Ok - = Result - where - T: Send + 'c; - - fn map_ok(ok: Self::Ok, f: F) -> Self::Ok - where - T: Send + 'c, - U: Send + 'c, - F: FnOnce(T) -> U, - { - ok.map(f) - } - - fn decrypt_bytes<'a, V, A>(self, visitor: V, aad: A) -> Self::Ok - where - V: DecipherVisitor<'c> + Send + 'c, - A: IntoAad<'a>, - { - match self.ciphertext { - AesCipherText::Single(ct) => { - let aad = aad.into_aad(); - let bytes = Self::decrypt_local_ciphertext(self.cipher, ct, aad.as_bytes())?; - visitor.visit_bytes_vec(bytes) - } - _ => Err(Unspecified), - } - } - - fn decrypt_seq<'a, V, A>(self, visitor: V, aad: A) -> Self::Ok - where - V: DecipherVisitor<'c> + Send + 'c, - A: IntoAad<'a>, - { - match self.ciphertext { - AesCipherText::Sequence(items) => { - let seq_access = AesSeqAccess { - cipher: self.cipher, - items: items.into_iter(), - aad: aad.into_aad(), - }; - visitor.visit_seq(seq_access) - } - _ => Err(Unspecified), - } - } - - fn decrypt_map<'a, V, A>(self, visitor: V, aad: A) -> Self::Ok - where - V: DecipherVisitor<'c> + Send + 'c, - A: IntoAad<'a>, - { - match self.ciphertext { - AesCipherText::Map(entries) => { - let map_access = AesMapAccess { - cipher: self.cipher, - entries: entries.into_iter(), - aad: aad.into_aad(), - }; - visitor.visit_map(map_access) - } - _ => Err(Unspecified), - } - } - - fn decrypt_passthrough(self) -> Self::Ok - where - T: Any + Send + 'static, - { - match self.ciphertext { - AesCipherText::Passthrough(boxed) => { - boxed.downcast::().map(|b| *b).map_err(|_| Unspecified) - } - _ => Err(Unspecified), - } - } - - fn decrypt_option<'a, T, A>(self, aad: A) -> Self::Ok> - where - T: Decrypt<'c> + 'c, - A: IntoAad<'a>, - { - match self.ciphertext { - AesCipherText::None(ct) => { - // Verify the AAD-bound tag over the empty plaintext. - let aad = aad.into_aad(); - Self::decrypt_local_ciphertext(self.cipher, ct, aad.as_bytes())?; - Ok(None) - } - // Passthrough must never be decoded as an Option payload. - AesCipherText::Passthrough(_) => Err(Unspecified), - // Any other variant is the `Some` payload: recurse into `T`. This is - // what lets `Option>`, `Option>`, - // `Option>` compose naturally. - // - // Note: there is no depth tag in the ciphertext, so the *shape* of - // nested options is decided by `T` at the call site, not by the - // bytes — `Some(Some(x))` and `Some(x)` seal to identical - // `Single(_)` ciphertexts. Decoding the same ciphertext as - // `Option` yields `Some("x")` and as `Option>` - // yields `Some(Some("x"))`; both succeed. This mirrors serde's - // treatment of `Option` and is intentional — every caller fixes a - // concrete type at the call site. See `nested_option_shape_is_caller_decided`. - other => { - let inner = self.cipher.decipher(other); - T::decrypt_with_aad(inner, aad).map(Some) - } - } - } -} - -struct AesSeqAccess<'c, 'a> { - cipher: &'c Aes256Cipher, - items: std::vec::IntoIter, - // Held as `Aad` (copy-on-write) rather than an owned `Vec` so a borrowed - // AAD stays borrowed. `next_element` re-supplies it per element by *borrowing* - // these bytes, so there is no per-element allocation in either the borrowed - // (`&str`/`&[u8]`) or the owned (`Vec`/PAE) case. - aad: Aad<'a>, -} - -impl<'c, 'a> SeqAccess<'c> for AesSeqAccess<'c, 'a> { - type Error = Unspecified; - - fn next_element + 'c>(&mut self) -> Result, Self::Error> { - let ct = match self.items.next() { - Some(ct) => ct, - None => return Ok(None), - }; - let decipher = self.cipher.decipher(ct); - // Each element was sealed with the same AAD; re-supply it per element by - // *borrowing* the stored bytes — no per-element allocation, even when the - // AAD is owned. Mirrors `SeqCipher::encrypt_next` binding AAD per element. - T::decrypt_with_aad(decipher, self.aad.as_bytes()).map(Some) - } -} - -struct AesMapAccess<'c, 'a> { - cipher: &'c Aes256Cipher, - entries: std::vec::IntoIter<(String, AesCipherText)>, - aad: Aad<'a>, -} - -impl<'c, 'a> MapAccess<'c> for AesMapAccess<'c, 'a> { - type Error = Unspecified; - - fn next_entry + 'c>(&mut self) -> Result, Self::Error> { - let (key, ct) = match self.entries.next() { - Some(entry) => entry, - None => return Ok(None), - }; - let decipher = self.cipher.decipher(ct); - // Borrow the stored AAD bytes per entry — see `AesSeqAccess::next_element`. - let value = T::decrypt_with_aad(decipher, self.aad.as_bytes())?; - Ok(Some((key, value))) + TreeDecipher::new(AesOpener(self), ciphertext) } } @@ -483,7 +280,7 @@ mod test { use crate::key::tests::DifferingKeyPair; use quickcheck_macros::quickcheck; use std::collections::HashMap; - use vitaminc_aead::Encrypt; + use vitaminc_aead::{Decipher, Encrypt, MapCipher, SeqCipher}; #[quickcheck] fn roundtrip_byte_array(key: Key, plaintext: [u8; 16]) -> bool { @@ -563,9 +360,9 @@ mod test { fn decrypt_seq_fails_with_wrong_aad(key: Key, plaintext: Vec) -> bool { // The Sequence path re-supplies the same AAD to every element, so a wrong // AAD must fail the per-element authentication rather than silently - // decrypting (see `AesSeqAccess::next_element`). An empty sequence has no - // element tags to reject — the container shape itself is not - // AEAD-authenticated on either side — so skip it. + // decrypting. An empty sequence has no element tags to reject — the + // container shape itself is not AEAD-authenticated on either side — so + // skip it. if plaintext.is_empty() { return true; } @@ -901,8 +698,7 @@ mod test { where T: Any + Send + 'static, { - let decipher = cipher.decipher(ciphertext); - decipher.decrypt_passthrough::() + cipher.decipher(ciphertext).decrypt_passthrough::() } #[quickcheck]