From 7abe4d8c0ec3d618090bc4b903a86c4191106f44 Mon Sep 17 00:00:00 2001 From: Dan Draper Date: Sun, 7 Jun 2026 21:11:23 +0800 Subject: [PATCH 1/2] feat(aead): extract a reusable CipherTree codec for structured ciphers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The structural drive and visitor walk for a recursive ciphertext (single value, sequence, map, option, passthrough) were identical between Aes256Cipher and any other cipher producing the same shape — only the per-leaf seal/open differs. Factor that out into vitaminc-aead: - CipherTree: the recursive container, generic over the leaf payload. - LeafSealer / LeafOpener: the two small per-cipher operations. - TreeCipher / TreeDecipher: generic Cipher/Decipher (plus TreeSeq/TreeMap and the SeqAccess/MapAccess drivers) that do the whole structural traversal. These are generic impls on the concrete local types TreeCipher/TreeDecipher, not a blanket impl over &C — a conditional blanket would conflict (E0119) with the direct Cipher/Decipher impls those traits exist to allow. Aes256Cipher is refactored onto it: AesCipherText is now CipherTree, AesDecipher is TreeDecipher, and the hand-written Seq/Map ciphers, Decipher impl, and access structs are gone (~290 fewer lines). Behaviour is unchanged — the full encrypt suite passes on both the aws-lc-rs and RustCrypto backends, including the map-contract, option-shape, passthrough, AAD, and nonce-uniqueness tests. This is the shared foundation a ZeroKMS-backed cipher (cipherstash-suite stack-encrypt) can build on, sealing each leaf under its own data key while reusing all of the structural machinery here. --- packages/aead/src/lib.rs | 6 + packages/aead/src/tree.rs | 448 +++++++++++++++++++++++++++++++++ packages/encrypt/src/cipher.rs | 434 ++++++-------------------------- 3 files changed, 526 insertions(+), 362 deletions(-) create mode 100644 packages/aead/src/tree.rs diff --git a/packages/aead/src/lib.rs b/packages/aead/src/lib.rs index 2718bb7..01f9c28 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, LeafSealer, TreeCipher, TreeDecipher, TreeMap, TreeMapAccess, TreeSeq, + TreeSeqAccess, +}; diff --git a/packages/aead/src/tree.rs b/packages/aead/src/tree.rs new file mode 100644 index 0000000..07db028 --- /dev/null +++ b/packages/aead/src/tree.rs @@ -0,0 +1,448 @@ +//! A reusable recursive ciphertext tree and the generic [`Cipher`]/[`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`] for any [`LeafSealer`] `S`; it +//! handles the sequence/map/option/passthrough structure and calls +//! [`LeafSealer::seal`] at each byte leaf. +//! - [`TreeDecipher`] implements [`Decipher`] for any [`LeafOpener`] `O`; it +//! drives the [`DecipherVisitor`] walk and calls [`LeafOpener::open`] at each +//! leaf. +//! +//! A concrete cipher therefore only has to define its leaf type and the two +//! leaf operations. `vitaminc_encrypt::Aes256Cipher` seals each leaf eagerly +//! under a random nonce; a ZeroKMS-backed cipher can stash plaintext leaves and +//! batch-seal them afterwards — both reuse everything here. +//! +//! ## Coherence +//! +//! The generic impls are on the concrete local types [`TreeCipher`] / +//! [`TreeDecipher`], **not** a blanket `impl Cipher for &C`. A conditional +//! blanket would conflict (E0119) with any direct `Cipher`/`Decipher` impl, +//! and those traits are public extension points meant to be implemented +//! directly. Concrete ciphers wrap themselves in a [`LeafSealer`]/[`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 sealed `LocalCipherText`, or a +/// pending plaintext awaiting a data key). +#[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`]. + 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 + /// keys in the same order, so collectors and sealers line up. + 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) => { + let mut out = Vec::with_capacity(items.len()); + for item in items { + out.push(item.try_map_leaves(f)?); + } + CipherTree::Sequence(out) + } + CipherTree::Map(entries) => { + let mut out = Vec::with_capacity(entries.len()); + for (k, v) in entries { + out.push((k, v.try_map_leaves(f)?)); + } + CipherTree::Map(out) + } + CipherTree::Passthrough(value) => CipherTree::Passthrough(value), + }) + } +} + +/// Seals one plaintext leaf. The single per-cipher operation on the encrypt side. +/// +/// Implementors are typically a `Copy` wrapper around a borrowed cipher (so the +/// sealer can be threaded through nested structures without ownership churn). +pub trait LeafSealer: Copy { + /// The sealed (or pending) per-leaf payload this sealer produces. + type Leaf; + + /// Seal `plaintext` bound to `aad`, producing the leaf payload. + fn seal(self, plaintext: Protected>, aad: &[u8]) -> Result; +} + +/// Opens one leaf back to plaintext. The single per-cipher operation on the +/// decrypt side. `Copy` for the same threading reason as [`LeafSealer`]. +pub trait LeafOpener: Copy { + /// The per-leaf payload this opener consumes (matches a [`LeafSealer::Leaf`]). + type Leaf; + + /// Open `leaf`, authenticating against `aad`, returning the plaintext bytes. + fn open(self, leaf: Self::Leaf, aad: &[u8]) -> Result>, Unspecified>; +} + +// ============================================================================= +// Encrypt side +// ============================================================================= + +/// A [`Cipher`] that builds a [`CipherTree`] by delegating each leaf to a +/// [`LeafSealer`]. Wrap a cipher's sealer in this to get the full structural +/// drive for free. +pub struct TreeCipher(pub S); + +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>, + { + let aad = aad.into_aad(); + Ok(CipherTree::Single(self.0.seal(data, aad.as_bytes())?)) + } + + fn encrypt_seq(self, size_hint: Option) -> Self::SeqCipher { + TreeSeq { + sealer: self.0, + items: Vec::with_capacity(size_hint.unwrap_or(0)), + } + } + + fn encrypt_map(self) -> Self::MapCipher { + TreeMap { + sealer: self.0, + entries: Vec::new(), + current_key: None, + } + } + + fn encrypt_none<'a, A>(self, aad: A) -> Result + where + A: IntoAad<'a>, + { + // Seal an empty plaintext so the leaf's tag binds the AAD. + let aad = aad.into_aad(); + Ok(CipherTree::None( + self.0.seal(Protected::new(Vec::new()), aad.as_bytes())?, + )) + } + + fn passthrough(self, value: T) -> Result + where + T: Any + Send + 'static, + { + Ok(CipherTree::Passthrough(Box::new(value))) + } +} + +/// [`SeqCipher`] for [`TreeCipher`]: accumulates a [`CipherTree`] per element. +pub struct TreeSeq { + sealer: S, + 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 `TreeCipher` over the (Copy) sealer. + self.items + .push(data.encrypt_with_aad(TreeCipher(self.sealer), 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 { + sealer: S, + 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(self.sealer), 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 +// ============================================================================= + +/// 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 { + tree: other, + opener: self.opener, + }, + 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 { + tree, + opener: self.opener, + }, + 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 { + tree, + opener: self.opener, + }, + self.aad.as_bytes(), + )?; + Ok(Some((key, value))) + } + None => Ok(None), + } + } +} diff --git a/packages/encrypt/src/cipher.rs b/packages/encrypt/src/cipher.rs index 79af310..5268735 100644 --- a/packages/encrypt/src/cipher.rs +++ b/packages/encrypt/src/cipher.rs @@ -2,37 +2,22 @@ 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, LeafSealer, + LocalCipherText, NonceGenerator, RandomNonceGenerator, 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 (see +/// [`AesSealer`] / [`AesOpener`]). +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`. @@ -55,183 +40,92 @@ impl Aes256Cipher { } } -impl<'c> Cipher for &'c Aes256Cipher { - type Ok = AesCipherText; - type Error = Unspecified; - type SeqCipher = AesSeqCipher<'c>; - type MapCipher = AesMapCipher<'c>; +/// The per-leaf sealing step for [`Aes256Cipher`]: seal a leaf under a fresh +/// random nonce. A `Copy` wrapper over the cipher so the generic [`TreeCipher`] +/// machinery can thread it through nested structures. +#[derive(Clone, Copy)] +pub struct AesSealer<'c>(&'c Aes256Cipher); - fn encrypt_bytes_vec<'a, A>( +impl LeafSealer for AesSealer<'_> { + type Leaf = LocalCipherText; + + fn seal( self, - data: Protected>, - aad: A, - ) -> Result - where - A: IntoAad<'a>, - { - let nonce = self.nonce_generator.generate()?; + plaintext: Protected>, + aad: &[u8], + ) -> Result { + let nonce = self.0.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) + .append_target_plaintext(plaintext) .accepts_ciphertext_and_tag_ok(|mut buf| { - self.key - .seal(&nonce_bytes, aad.as_bytes(), &mut buf) - .map(|()| buf) + self.0.key.seal(&nonce_bytes, aad, &mut buf).map(|()| buf) }) .build() - .map(AesCipherText::Single) - } - - fn encrypt_seq(self, size_hint: Option) -> Self::SeqCipher { - AesSeqCipher { - cipher: self, - items: Vec::with_capacity(size_hint.unwrap_or(0)), - } - } - - fn encrypt_map(self) -> Self::MapCipher { - AesMapCipher { - cipher: self, - entries: Vec::new(), - current_key: None, - } } +} - fn encrypt_none<'a, A>(self, aad: A) -> Result - 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(); +/// The per-leaf opening step for [`Aes256Cipher`]: decrypt a leaf with the +/// cipher's key and the nonce stored in the leaf. +#[derive(Clone, Copy)] +pub struct AesOpener<'c>(&'c Aes256Cipher); - 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) - } +impl LeafOpener for AesOpener<'_> { + type Leaf = LocalCipherText; - fn passthrough(self, value: T) -> Result - where - T: Any + Send + 'static, - { - Ok(AesCipherText::Passthrough(Box::new(value))) + 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() } } -/// [`SeqCipher`] driver for [`Aes256Cipher`]. Encrypts each element under its -/// own fresh nonce and accumulates the results into an -/// [`AesCipherText::Sequence`]. -pub struct AesSeqCipher<'c> { - cipher: &'c Aes256Cipher, - items: Vec, -} +/// 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>; -impl<'c> SeqCipher for AesSeqCipher<'c> { +impl<'c> Cipher for &'c Aes256Cipher { type Ok = AesCipherText; type Error = Unspecified; + type SeqCipher = TreeSeq>; + type MapCipher = TreeMap>; - fn encrypt_next<'a, T, A>(mut self, data: T, aad: A) -> Result + fn encrypt_bytes_vec<'a, A>( + self, + data: Protected>, + aad: A, + ) -> Result where - T: Encrypt, A: IntoAad<'a>, { - let encrypted = data.encrypt_with_aad(self.cipher, aad)?; - self.items.push(encrypted); - Ok(self) + TreeCipher(AesSealer(self)).encrypt_bytes_vec(data, aad) } - fn passthrough_next(mut self, value: T) -> Result - where - T: Any + Send + 'static, - { - self.items.push(AesCipherText::Passthrough(Box::new(value))); - Ok(self) - } - - fn end(self) -> Result { - Ok(AesCipherText::Sequence(self.items)) + fn encrypt_seq(self, size_hint: Option) -> Self::SeqCipher { + TreeCipher(AesSealer(self)).encrypt_seq(size_hint) } -} -/// [`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> { - cipher: &'c Aes256Cipher, - entries: Vec<(String, AesCipherText)>, - current_key: Option<&'static str>, -} - -impl<'c> MapCipher for AesMapCipher<'c> { - 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); - Ok(self) + fn encrypt_map(self) -> Self::MapCipher { + TreeCipher(AesSealer(self)).encrypt_map() } - fn encrypt_value<'a, U, A>(mut self, value: U, aad: A) -> Result + fn encrypt_none<'a, A>(self, aad: A) -> Result where - U: 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)); - Ok(self) + TreeCipher(AesSealer(self)).encrypt_none(aad) } - fn passthrough_entry(mut self, key: &'static str, value: T) -> Result + fn passthrough(self, value: T) -> Result 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)))); - 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)) + TreeCipher(AesSealer(self)).passthrough(value) } } @@ -264,202 +158,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 +194,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 +274,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 +612,7 @@ mod test { where T: Any + Send + 'static, { - let decipher = cipher.decipher(ciphertext); - decipher.decrypt_passthrough::() + cipher.decipher(ciphertext).decrypt_passthrough::() } #[quickcheck] From c93f2f101fa71deb10a482371957e2bee23cd9f8 Mon Sep 17 00:00:00 2001 From: Dan Draper Date: Mon, 8 Jun 2026 10:51:10 +0800 Subject: [PATCH 2/2] refactor(aead): split structural codec from sealing for batch-cipher reuse MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `LeafSealer` seam sealed each leaf eagerly during the encrypt traversal, behind a `Copy` bound. That fits an eager single-key cipher but not a batching one (e.g. ZeroKMS), which must collect all leaves, fetch data keys in one request, then seal — it could only satisfy the old seam with a no-op `seal` and a second pass anyway. Split structure from sealing so one codec serves both: - `TreeCipher`/`TreeSeq`/`TreeMap` are now a non-generic structural codec that collects each byte leaf (plaintext + bound AAD) into a `CipherTree` without sealing. `LeafSealer` is removed. - Sealing is a separate per-cipher step. `Aes256Cipher` seals the collected tree via `try_map_leaves` (eager, fresh nonce per leaf); a batch cipher can size a key request with `leaf_count` and seal the same way. `value.encrypt(&cipher)` is unchanged for callers. - `LeafOpener` is kept (decrypt openers are stateless per-leaf lookups, so `Copy` is correct there) and its three recursion sites now use `TreeDecipher::new`. Also: - Add `Nonce::as_array` so `seal` copies the nonce bytes without the previously-dead fallible slice conversion (it needs the nonce twice: for the AEAD call and to append to the ciphertext). - Pin the leaf-ordering contract (`leaf_count`/`for_each_leaf`/ `try_map_leaves` agree on set and order) with a test. - Rewrite `try_map_leaves` seq/map arms as iterator combinators. - Document `CipherTree::None`'s per-leaf key cost and `Passthrough`'s non-serializable type erasure for the future batch cipher. --- packages/aead/src/lib.rs | 2 +- packages/aead/src/nonce.rs | 9 + packages/aead/src/tree.rs | 289 +++++++++++++++++++++------------ packages/encrypt/src/cipher.rs | 150 +++++++++++++---- 4 files changed, 312 insertions(+), 138 deletions(-) diff --git a/packages/aead/src/lib.rs b/packages/aead/src/lib.rs index 01f9c28..f5909a6 100644 --- a/packages/aead/src/lib.rs +++ b/packages/aead/src/lib.rs @@ -27,6 +27,6 @@ pub use encrypt::Encrypt; pub use nonce::{Nonce, NonceGenerator, RandomNonceGenerator}; #[doc(inline)] pub use tree::{ - CipherTree, LeafOpener, LeafSealer, TreeCipher, TreeDecipher, TreeMap, TreeMapAccess, TreeSeq, + 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 index 07db028..9f55ba6 100644 --- a/packages/aead/src/tree.rs +++ b/packages/aead/src/tree.rs @@ -1,5 +1,5 @@ -//! A reusable recursive ciphertext tree and the generic [`Cipher`]/[`Decipher`] -//! machinery that drives it. +//! 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 @@ -7,26 +7,53 @@ //! //! - [`CipherTree`] is the recursive container, generic over the per-leaf //! payload `L`. -//! - [`TreeCipher`] implements [`Cipher`] for any [`LeafSealer`] `S`; it -//! handles the sequence/map/option/passthrough structure and calls -//! [`LeafSealer::seal`] at each byte leaf. +//! - [`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. //! -//! A concrete cipher therefore only has to define its leaf type and the two -//! leaf operations. `vitaminc_encrypt::Aes256Cipher` seals each leaf eagerly -//! under a random nonce; a ZeroKMS-backed cipher can stash plaintext leaves and -//! batch-seal them afterwards — both reuse everything here. +//! ## 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 //! -//! The generic impls are on the concrete local types [`TreeCipher`] / -//! [`TreeDecipher`], **not** a blanket `impl Cipher for &C`. A conditional -//! blanket would conflict (E0119) with any direct `Cipher`/`Decipher` impl, -//! and those traits are public extension points meant to be implemented -//! directly. Concrete ciphers wrap themselves in a [`LeafSealer`]/[`LeafOpener`] -//! instead. +//! [`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; @@ -39,8 +66,8 @@ use crate::{ /// A recursive ciphertext container whose shape mirrors the encrypted plaintext. /// -/// Generic over the per-leaf payload `L` (e.g. a sealed `LocalCipherText`, or a -/// pending plaintext awaiting a data key). +/// 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. @@ -52,6 +79,13 @@ pub enum 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), } @@ -69,7 +103,8 @@ impl CipherTree { /// 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 - /// keys in the same order, so collectors and sealers line up. + /// 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), @@ -90,61 +125,60 @@ impl CipherTree { Ok(match self { CipherTree::Single(leaf) => CipherTree::Single(f(leaf)?), CipherTree::None(leaf) => CipherTree::None(f(leaf)?), - CipherTree::Sequence(items) => { - let mut out = Vec::with_capacity(items.len()); - for item in items { - out.push(item.try_map_leaves(f)?); - } - CipherTree::Sequence(out) - } - CipherTree::Map(entries) => { - let mut out = Vec::with_capacity(entries.len()); - for (k, v) in entries { - out.push((k, v.try_map_leaves(f)?)); - } - CipherTree::Map(out) - } + 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), }) } } -/// Seals one plaintext leaf. The single per-cipher operation on the encrypt side. +/// A byte leaf collected by [`TreeCipher`], awaiting sealing. /// -/// Implementors are typically a `Copy` wrapper around a borrowed cipher (so the -/// sealer can be threaded through nested structures without ownership churn). -pub trait LeafSealer: Copy { - /// The sealed (or pending) per-leaf payload this sealer produces. - type Leaf; - - /// Seal `plaintext` bound to `aad`, producing the leaf payload. - fn seal(self, plaintext: Protected>, aad: &[u8]) -> Result; +/// 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, } -/// Opens one leaf back to plaintext. The single per-cipher operation on the -/// decrypt side. `Copy` for the same threading reason as [`LeafSealer`]. -pub trait LeafOpener: Copy { - /// The per-leaf payload this opener consumes (matches a [`LeafSealer::Leaf`]). - type Leaf; - - /// Open `leaf`, authenticating against `aad`, returning the plaintext bytes. - fn open(self, leaf: Self::Leaf, aad: &[u8]) -> Result>, Unspecified>; +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 +// Encrypt side — a structural codec that collects, but does not seal. // ============================================================================= -/// A [`Cipher`] that builds a [`CipherTree`] by delegating each leaf to a -/// [`LeafSealer`]. Wrap a cipher's sealer in this to get the full structural -/// drive for free. -pub struct TreeCipher(pub S); +/// 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; +impl Cipher for TreeCipher { + type Ok = CipherTree; type Error = Unspecified; - type SeqCipher = TreeSeq; - type MapCipher = TreeMap; + type SeqCipher = TreeSeq; + type MapCipher = TreeMap; fn encrypt_bytes_vec<'a, A>( self, @@ -154,20 +188,17 @@ impl Cipher for TreeCipher { where A: IntoAad<'a>, { - let aad = aad.into_aad(); - Ok(CipherTree::Single(self.0.seal(data, aad.as_bytes())?)) + Ok(CipherTree::Single(PendingLeaf::new(data, aad))) } fn encrypt_seq(self, size_hint: Option) -> Self::SeqCipher { TreeSeq { - sealer: self.0, items: Vec::with_capacity(size_hint.unwrap_or(0)), } } fn encrypt_map(self) -> Self::MapCipher { TreeMap { - sealer: self.0, entries: Vec::new(), current_key: None, } @@ -177,11 +208,11 @@ impl Cipher for TreeCipher { where A: IntoAad<'a>, { - // Seal an empty plaintext so the leaf's tag binds the AAD. - let aad = aad.into_aad(); - Ok(CipherTree::None( - self.0.seal(Protected::new(Vec::new()), aad.as_bytes())?, - )) + // 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 @@ -192,14 +223,13 @@ impl Cipher for TreeCipher { } } -/// [`SeqCipher`] for [`TreeCipher`]: accumulates a [`CipherTree`] per element. -pub struct TreeSeq { - sealer: S, - items: Vec>, +/// [`SeqCipher`] for [`TreeCipher`]: accumulates a pending [`CipherTree`] per element. +pub struct TreeSeq { + items: Vec>, } -impl SeqCipher for TreeSeq { - type Ok = CipherTree; +impl SeqCipher for TreeSeq { + type Ok = CipherTree; type Error = Unspecified; fn encrypt_next<'a, T, A>(mut self, data: T, aad: A) -> Result @@ -207,9 +237,8 @@ impl SeqCipher for TreeSeq { T: Encrypt, A: IntoAad<'a>, { - // Recurse with a fresh `TreeCipher` over the (Copy) sealer. - self.items - .push(data.encrypt_with_aad(TreeCipher(self.sealer), aad)?); + // Recurse with a fresh structural codec — the element is collected, not sealed. + self.items.push(data.encrypt_with_aad(TreeCipher, aad)?); Ok(self) } @@ -228,14 +257,13 @@ impl SeqCipher for TreeSeq { /// [`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 { - sealer: S, - entries: Vec<(String, CipherTree)>, +pub struct TreeMap { + entries: Vec<(String, CipherTree)>, current_key: Option<&'static str>, } -impl MapCipher for TreeMap { - type Ok = CipherTree; +impl MapCipher for TreeMap { + type Ok = CipherTree; type Error = Unspecified; fn encrypt_key(mut self, key: &'static str) -> Result { @@ -253,7 +281,7 @@ impl MapCipher for TreeMap { A: IntoAad<'a>, { let key = self.current_key.take().ok_or(Unspecified)?; - let value = value.encrypt_with_aad(TreeCipher(self.sealer), aad)?; + let value = value.encrypt_with_aad(TreeCipher, aad)?; self.entries.push((key.to_string(), value)); Ok(self) } @@ -282,6 +310,21 @@ impl MapCipher for TreeMap { // 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 { @@ -381,14 +424,7 @@ impl<'c, O: LeafOpener> Decipher<'c> for TreeDecipher { } CipherTree::Passthrough(_) => Err(Unspecified), // Any other shape is the `Some` payload — recurse into `T`. - other => T::decrypt_with_aad( - TreeDecipher { - tree: other, - opener: self.opener, - }, - aad, - ) - .map(Some), + other => T::decrypt_with_aad(TreeDecipher::new(self.opener, other), aad).map(Some), } } } @@ -406,14 +442,10 @@ impl<'c, 'a, O: LeafOpener> SeqAccess<'c> for TreeSeqAccess<'a, O> { fn next_element + 'c>(&mut self) -> Result, Self::Error> { match self.items.next() { - Some(tree) => T::decrypt_with_aad( - TreeDecipher { - tree, - opener: self.opener, - }, - self.aad.as_bytes(), - ) - .map(Some), + Some(tree) => { + T::decrypt_with_aad(TreeDecipher::new(self.opener, tree), self.aad.as_bytes()) + .map(Some) + } None => Ok(None), } } @@ -433,16 +465,63 @@ impl<'c, 'a, O: LeafOpener> MapAccess<'c> for TreeMapAccess<'a, O> { fn next_entry + 'c>(&mut self) -> Result, Self::Error> { match self.entries.next() { Some((key, tree)) => { - let value = T::decrypt_with_aad( - TreeDecipher { - tree, - opener: self.opener, - }, - self.aad.as_bytes(), - )?; + 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 5268735..3a7f98e 100644 --- a/packages/encrypt/src/cipher.rs +++ b/packages/encrypt/src/cipher.rs @@ -2,9 +2,9 @@ use crate::backend::{CipherKey, NONCE_LEN}; use crate::Key; use std::any::Any; use vitaminc_aead::{ - Cipher, CipherTextBuilder, CipherTree, Decrypt, IntoAad, LeafOpener, LeafSealer, - LocalCipherText, NonceGenerator, RandomNonceGenerator, TreeCipher, TreeDecipher, TreeMap, - TreeSeq, Unspecified, + Cipher, CipherTextBuilder, CipherTree, Decrypt, IntoAad, LeafOpener, LocalCipherText, + MapCipher, NonceGenerator, PendingLeaf, RandomNonceGenerator, SeqCipher, TreeCipher, + TreeDecipher, TreeMap, TreeSeq, Unspecified, }; use vitaminc_protected::Protected; @@ -15,8 +15,7 @@ use vitaminc_protected::Protected; /// [`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 (see -/// [`AesSealer`] / [`AesOpener`]). +/// cipher built on [`CipherTree`]; only the per-leaf sealing differs. pub type AesCipherText = CipherTree; /// Implements AES-256-GCM. Backend is selected at compile time: @@ -38,37 +37,40 @@ impl Aes256Cipher { key: key.cipher_key()?, }) } -} - -/// The per-leaf sealing step for [`Aes256Cipher`]: seal a leaf under a fresh -/// random nonce. A `Copy` wrapper over the cipher so the generic [`TreeCipher`] -/// machinery can thread it through nested structures. -#[derive(Clone, Copy)] -pub struct AesSealer<'c>(&'c Aes256Cipher); -impl LeafSealer for AesSealer<'_> { - type Leaf = LocalCipherText; - - fn seal( - self, - plaintext: Protected>, - aad: &[u8], - ) -> Result { - let nonce = self.0.nonce_generator.generate()?; - let nonce_bytes: [u8; NONCE_LEN] = nonce.as_ref().try_into().map_err(|_| Unspecified)?; + /// 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(plaintext) + .append_target_plaintext(leaf.plaintext) .accepts_ciphertext_and_tag_ok(|mut buf| { - self.0.key.seal(&nonce_bytes, aad, &mut buf).map(|()| 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. +/// 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); @@ -89,11 +91,17 @@ impl LeafOpener for AesOpener<'_> { /// [`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 = TreeSeq>; - type MapCipher = TreeMap>; + type SeqCipher = AesSeq<'c>; + type MapCipher = AesMap<'c>; fn encrypt_bytes_vec<'a, A>( self, @@ -103,29 +111,107 @@ impl<'c> Cipher for &'c Aes256Cipher { where A: IntoAad<'a>, { - TreeCipher(AesSealer(self)).encrypt_bytes_vec(data, aad) + self.seal_tree(TreeCipher.encrypt_bytes_vec(data, aad)?) } fn encrypt_seq(self, size_hint: Option) -> Self::SeqCipher { - TreeCipher(AesSealer(self)).encrypt_seq(size_hint) + AesSeq { + inner: TreeCipher.encrypt_seq(size_hint), + cipher: self, + } } fn encrypt_map(self) -> Self::MapCipher { - TreeCipher(AesSealer(self)).encrypt_map() + AesMap { + inner: TreeCipher.encrypt_map(), + cipher: self, + } } fn encrypt_none<'a, A>(self, aad: A) -> Result where A: IntoAad<'a>, { - TreeCipher(AesSealer(self)).encrypt_none(aad) + self.seal_tree(TreeCipher.encrypt_none(aad)?) } fn passthrough(self, value: T) -> Result where T: Any + Send + 'static, { - TreeCipher(AesSealer(self)).passthrough(value) + self.seal_tree(TreeCipher.passthrough(value)?) + } +} + +/// [`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, +} + +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: vitaminc_aead::Encrypt, + A: IntoAad<'a>, + { + self.inner = self.inner.encrypt_next(data, aad)?; + Ok(self) + } + + fn passthrough_next(mut self, value: T) -> Result + where + T: Any + Send + 'static, + { + self.inner = self.inner.passthrough_next(value)?; + Ok(self) + } + + fn end(self) -> Result { + self.cipher.seal_tree(self.inner.end()?) + } +} + +/// [`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, +} + +impl MapCipher for AesMap<'_> { + type Ok = AesCipherText; + type Error = Unspecified; + + fn encrypt_key(mut self, key: &'static str) -> Result { + self.inner = self.inner.encrypt_key(key)?; + Ok(self) + } + + fn encrypt_value<'a, T, A>(mut self, value: T, aad: A) -> Result + where + T: vitaminc_aead::Encrypt, + A: IntoAad<'a>, + { + self.inner = self.inner.encrypt_value(value, aad)?; + Ok(self) + } + + fn passthrough_entry(mut self, key: &'static str, value: T) -> Result + where + T: Any + Send + 'static, + { + self.inner = self.inner.passthrough_entry(key, value)?; + Ok(self) + } + + fn end(self) -> Result { + self.cipher.seal_tree(self.inner.end()?) } }