From 8339690964c6cac15c61522b844e271dafc23bf0 Mon Sep 17 00:00:00 2001 From: Marvin Jones Date: Mon, 15 Jun 2026 15:10:56 -0400 Subject: [PATCH] Clean up key protocol --- .../tests/auth_transfer/private.rs | 4 +- .../key_management/ephemeral_key_holder.rs | 11 --- .../src/key_management/group_key_holder.rs | 13 ++- .../key_management/key_tree/chain_index.rs | 25 +++--- .../key_management/key_tree/keys_private.rs | 80 +++++-------------- .../key_management/key_tree/keys_public.rs | 69 +++++----------- .../src/key_management/key_tree/mod.rs | 61 +++++++------- lee/key_protocol/src/key_management/mod.rs | 39 ++++----- .../src/key_management/secret_holders.rs | 39 ++------- lee/state_machine/core/src/encoding.rs | 4 +- lee/state_machine/core/src/encryption/mod.rs | 5 +- .../src/encryption/shared_key_derivation.rs | 9 ++- lee/state_machine/core/src/lib.rs | 3 +- lee/state_machine/src/state.rs | 4 +- 14 files changed, 130 insertions(+), 236 deletions(-) diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index 30f0cfdd..9094460e 100644 --- a/integration_tests/tests/auth_transfer/private.rs +++ b/integration_tests/tests/auth_transfer/private.rs @@ -11,7 +11,7 @@ use lee::{ privacy_preserving_transaction::circuit::ProgramWithDependencies, program::Program, }; use lee_core::{ - EncryptedAccountData, InputAccountIdentity, NullifierPublicKey, + EncryptedAccountData, InputAccountIdentity, ML_KEM_768_CIPHERTEXT_LEN, NullifierPublicKey, account::AccountWithMetadata, encryption::{EphemeralPublicKey, ViewingPublicKey}, }; @@ -664,7 +664,7 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> { let npk = NullifierPublicKey::from(&nsk); let vpk = ViewingPublicKey::from_bytes(vec![4_u8; 1184]).unwrap(); let ssk = SharedSecretKey([55_u8; 32]); - let epk = EphemeralPublicKey(vec![55_u8; 1088]); + let epk = EphemeralPublicKey(vec![55_u8; ML_KEM_768_CIPHERTEXT_LEN]); let attacker_vault_id = { let seed = vault_core::compute_vault_seed(attacker_id); AccountId::for_private_pda(&vault_program_id, &seed, &npk, 1337) diff --git a/lee/key_protocol/src/key_management/ephemeral_key_holder.rs b/lee/key_protocol/src/key_management/ephemeral_key_holder.rs index a53ae47c..9bc81391 100644 --- a/lee/key_protocol/src/key_management/ephemeral_key_holder.rs +++ b/lee/key_protocol/src/key_management/ephemeral_key_holder.rs @@ -48,14 +48,3 @@ impl EphemeralKeyHolder { self.shared_secret } } - -/// Encapsulates a fresh shared secret toward `vpk` and returns `(shared_secret, ciphertext)`. -/// -/// Used when the local side is acting as an "ephemeral receiver" — i.e. generating a -/// one-sided encryption that only the holder of the VSK can decrypt. -#[must_use] -pub fn produce_one_sided_shared_secret_receiver( - vpk: &ViewingPublicKey, -) -> (SharedSecretKey, EphemeralPublicKey) { - SharedSecretKey::encapsulate(vpk) -} diff --git a/lee/key_protocol/src/key_management/group_key_holder.rs b/lee/key_protocol/src/key_management/group_key_holder.rs index 7fb24713..70d503ac 100644 --- a/lee/key_protocol/src/key_management/group_key_holder.rs +++ b/lee/key_protocol/src/key_management/group_key_holder.rs @@ -1,7 +1,7 @@ use aes_gcm::{Aes256Gcm, KeyInit as _, aead::Aead as _}; use lee_core::{ SharedSecretKey, - encryption::{EphemeralPublicKey, ViewingPublicKey}, + encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, ViewingPublicKey}, program::{PdaSeed, ProgramId}, }; use rand::{RngCore as _, rngs::OsRng}; @@ -170,7 +170,7 @@ impl GroupKeyHolder { .encrypt(&nonce, self.gms.as_ref()) .expect("AES-GCM encryption should not fail with valid key/nonce"); - let capacity = 1088_usize + let capacity = ML_KEM_768_CIPHERTEXT_LEN .checked_add(12) .and_then(|n| n.checked_add(ciphertext.len())) .expect("seal capacity overflow"); @@ -187,20 +187,19 @@ impl GroupKeyHolder { /// doesn't verify (wrong key or tampered data). pub fn unseal(sealed: &[u8], own_key: &SealingSecretKey) -> Result { // kem_ciphertext (1088) + nonce (12) = header, then AES-GCM tag (16) minimum. - const KEM_CT_LEN: usize = 1088; - const HEADER_LEN: usize = KEM_CT_LEN + 12; + const HEADER_LEN: usize = ML_KEM_768_CIPHERTEXT_LEN + 12; const MIN_LEN: usize = HEADER_LEN + 16; if sealed.len() < MIN_LEN { return Err(SealError::TooShort); } - let kem_ct = EphemeralPublicKey(sealed[..KEM_CT_LEN].to_vec()); - let nonce = aes_gcm::Nonce::from_slice(&sealed[KEM_CT_LEN..HEADER_LEN]); + let kem_ct = EphemeralPublicKey(sealed[..ML_KEM_768_CIPHERTEXT_LEN].to_vec()); + let nonce = aes_gcm::Nonce::from_slice(&sealed[ML_KEM_768_CIPHERTEXT_LEN..HEADER_LEN]); let ciphertext = &sealed[HEADER_LEN..]; let shared = SharedSecretKey::decapsulate(&kem_ct, &own_key.d, &own_key.z) - .expect("key_protocol::group_key_holder::GroupKeyHolder::unseal: KEM_CT_LEN guarantees exactly 1088 bytes"); + .expect("key_protocol::group_key_holder::GroupKeyHolder::unseal: ML_KEM_768_CIPHERTEXT_LEN guarantees exactly 1088 bytes"); let aes_key = Self::seal_kdf(&shared); let cipher = Aes256Gcm::new(&aes_key.into()); diff --git a/lee/key_protocol/src/key_management/key_tree/chain_index.rs b/lee/key_protocol/src/key_management/key_tree/chain_index.rs index b22dc779..6ea2e8a1 100644 --- a/lee/key_protocol/src/key_management/key_tree/chain_index.rs +++ b/lee/key_protocol/src/key_management/key_tree/chain_index.rs @@ -139,7 +139,7 @@ impl ChainIndex { .map(|item| Self(item.into_iter().copied().collect())) } - pub fn chain_ids_at_depth(depth: usize) -> impl Iterator { + fn collect_chain_ids_at_depth(depth: usize) -> Vec { let mut stack = vec![Self(vec![0; depth])]; let mut cumulative_stack = vec![Self(vec![0; depth])]; @@ -152,23 +152,18 @@ impl ChainIndex { } } - cumulative_stack.into_iter().unique() + cumulative_stack + } + + pub fn chain_ids_at_depth(depth: usize) -> impl Iterator { + Self::collect_chain_ids_at_depth(depth).into_iter().unique() } pub fn chain_ids_at_depth_rev(depth: usize) -> impl Iterator { - let mut stack = vec![Self(vec![0; depth])]; - let mut cumulative_stack = vec![Self(vec![0; depth])]; - - while let Some(top_id) = stack.pop() { - if let Some(collapsed_id) = top_id.collapse_back() { - for id in collapsed_id.shuffle_iter() { - stack.push(id.clone()); - cumulative_stack.push(id); - } - } - } - - cumulative_stack.into_iter().rev().unique() + Self::collect_chain_ids_at_depth(depth) + .into_iter() + .rev() + .unique() } } diff --git a/lee/key_protocol/src/key_management/key_tree/keys_private.rs b/lee/key_protocol/src/key_management/key_tree/keys_private.rs index 5a27be79..e2b982cc 100644 --- a/lee/key_protocol/src/key_management/key_tree/keys_private.rs +++ b/lee/key_protocol/src/key_management/key_tree/keys_private.rs @@ -6,7 +6,7 @@ use sha2::Digest as _; use crate::key_management::{ KeyChain, - key_tree::traits::KeyTreeNode, + key_tree::{split_hash, traits::KeyTreeNode}, secret_holders::{PrivateKeyHolder, SecretSpendingKey}, }; @@ -23,38 +23,11 @@ impl ChildKeysPrivate { #[must_use] pub fn root(seed: [u8; 64]) -> Self { let hash_value = hmac_sha512::HMAC::mac(seed, b"LEE_master_priv"); + let (first, ccc) = split_hash(&hash_value); - let ssk = SecretSpendingKey( - *hash_value - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32"), - ); - let ccc = *hash_value - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"); + let ssk = SecretSpendingKey(first); - let nsk = ssk.generate_nullifier_secret_key(None); - let vsk = ssk.generate_viewing_secret_seed_key(None); - - let npk = NullifierPublicKey::from(&nsk); - let vpk = ViewingPublicKey::from(&vsk); - - Self { - value: ( - KeyChain { - secret_spending_key: ssk, - nullifier_public_key: npk, - viewing_public_key: vpk, - private_key_holder: PrivateKeyHolder { - nullifier_secret_key: nsk, - viewing_secret_key: vsk, - }, - }, - BTreeMap::from_iter([(PrivateAccountKind::Regular(0), lee::Account::default())]), - ), - ccc, - cci: None, - } + Self::from_ssk_and_ccc(ssk, ccc, None) } #[must_use] @@ -77,18 +50,16 @@ impl ChildKeysPrivate { input.extend_from_slice(&cci.to_be_bytes()); let hash_value = hmac_sha512::HMAC::mac(input, self.ccc); + let (first, ccc) = split_hash(&hash_value); - let ssk = SecretSpendingKey( - *hash_value - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32"), - ); - let ccc = *hash_value - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"); + let ssk = SecretSpendingKey(first); - let nsk = ssk.generate_nullifier_secret_key(Some(cci)); - let vsk = ssk.generate_viewing_secret_seed_key(Some(cci)); + Self::from_ssk_and_ccc(ssk, ccc, Some(cci)) + } + + fn from_ssk_and_ccc(ssk: SecretSpendingKey, ccc: [u8; 32], cci: Option) -> Self { + let nsk = ssk.generate_nullifier_secret_key(cci); + let vsk = ssk.generate_viewing_secret_seed_key(cci); let npk = NullifierPublicKey::from(&nsk); let vpk = ViewingPublicKey::from(&vsk); @@ -107,7 +78,7 @@ impl ChildKeysPrivate { BTreeMap::from_iter([(PrivateAccountKind::Regular(0), lee::Account::default())]), ), ccc, - cci: Some(cci), + cci, } } } @@ -137,16 +108,16 @@ mod tests { use super::*; use crate::key_management::{self, secret_holders::ViewingSecretKey}; + const SEED: [u8; 64] = [ + 252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255, 49, + 43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18, 129, 223, + 176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56, 114, 39, 38, + 118, 197, 205, 225, + ]; + #[test] fn master_key_generation() { - let seed: [u8; 64] = [ - 252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255, - 49, 43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18, - 129, 223, 176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56, - 114, 39, 38, 118, 197, 205, 225, - ]; - - let keys = ChildKeysPrivate::root(seed); + let keys = ChildKeysPrivate::root(SEED); let expected_ssk = key_management::secret_holders::SecretSpendingKey([ 246, 79, 26, 124, 135, 95, 52, 51, 201, 27, 48, 194, 2, 144, 51, 219, 245, 128, 139, @@ -253,14 +224,7 @@ mod tests { #[test] fn child_keys_generation() { - let seed: [u8; 64] = [ - 252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255, - 49, 43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18, - 129, 223, 176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56, - 114, 39, 38, 118, 197, 205, 225, - ]; - - let root_node = ChildKeysPrivate::root(seed); + let root_node = ChildKeysPrivate::root(SEED); let child_node = ChildKeysPrivate::nth_child(&root_node, 42_u32); let expected_ssk = key_management::secret_holders::SecretSpendingKey([ diff --git a/lee/key_protocol/src/key_management/key_tree/keys_public.rs b/lee/key_protocol/src/key_management/key_tree/keys_public.rs index 947fb83c..4caad0e7 100644 --- a/lee/key_protocol/src/key_management/key_tree/keys_public.rs +++ b/lee/key_protocol/src/key_management/key_tree/keys_public.rs @@ -1,7 +1,7 @@ use k256::elliptic_curve::PrimeField as _; use serde::{Deserialize, Serialize}; -use crate::key_management::key_tree::traits::KeyTreeNode; +use crate::key_management::key_tree::{split_hash, traits::KeyTreeNode}; #[derive(Debug, Serialize, Deserialize, Clone)] #[cfg_attr(any(test, feature = "test_utils"), derive(PartialEq, Eq))] @@ -21,52 +21,32 @@ impl ChildKeysPublic { #[must_use] pub fn root(seed: [u8; 64]) -> Self { let hash_value = hmac_sha512::HMAC::mac(seed, "LEE_master_pub"); + let (first, cc) = split_hash(&hash_value); - let sk = lee::PrivateKey::try_new( - *hash_value - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32"), - ) - .expect("Expect a valid Private Key"); - let ssk = lee::PrivateKey::tweak(sk.value()).expect("`key_protocol::key_management::keys_public::root()`: Invalid private key produced from `tweak`"); + let sk = lee::PrivateKey::try_new(first).expect("Expect a valid Private Key"); - let cc = *hash_value - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"); - let pk = lee::PublicKey::new_from_private_key(&ssk); - - Self { - sk, - ssk, - pk, - cc, - cci: None, - } + Self::from_sk_and_cc(sk, cc, None) } #[must_use] pub fn nth_child(&self, cci: u32) -> Self { let hash_value = self.compute_hash_value(cci); + let (first, cc) = split_hash(&hash_value); - let lhs = k256::Scalar::from_repr( - (*hash_value - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32")) - .into(), - ) - .expect("Expect a valid k256 scalar"); + let lhs = k256::Scalar::from_repr(first.into()).expect("Expect a valid k256 scalar"); let rhs = k256::Scalar::from_repr((*self.sk.value()).into()).expect("Expect a valid k256 scalar"); let sk = lee::PrivateKey::try_new(lhs.add(&rhs).to_bytes().into()) .expect("Expect a valid private key"); - let ssk = lee::PrivateKey::tweak(sk.value()).expect("`key_protocol::key_management::keys_public::nth_child()`: Invalid private key produced from `tweak`"); - - let cc = *hash_value - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"); + Self::from_sk_and_cc(sk, cc, Some(cci)) + } + fn from_sk_and_cc(sk: lee::PrivateKey, cc: [u8; 32], cci: Option) -> Self { + let ssk = lee::PrivateKey::tweak(sk.value()).expect( + "`key_protocol::key_management::keys_public::ChildKeysPublic`: Invalid private key produced from `tweak`", + ); let pk = lee::PublicKey::new_from_private_key(&ssk); Self { @@ -74,7 +54,7 @@ impl ChildKeysPublic { ssk, pk, cc, - cci: Some(cci), + cci, } } @@ -128,15 +108,16 @@ mod tests { use super::*; + const SEED: [u8; 64] = [ + 88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173, + 134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87, 22, + 227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6, 187, + 148, 92, 44, 253, 210, 37, + ]; + #[test] fn master_keys_generation() { - let seed = [ - 88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173, - 134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87, - 22, 227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6, - 187, 148, 92, 44, 253, 210, 37, - ]; - let keys = ChildKeysPublic::root(seed); + let keys = ChildKeysPublic::root(SEED); let expected_cc = [ 238, 94, 84, 154, 56, 224, 80, 218, 133, 249, 179, 222, 9, 24, 17, 252, 120, 127, 222, @@ -169,13 +150,7 @@ mod tests { #[test] fn child_keys_generation() { - let seed = [ - 88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173, - 134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87, - 22, 227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6, - 187, 148, 92, 44, 253, 210, 37, - ]; - let root_keys = ChildKeysPublic::root(seed); + let root_keys = ChildKeysPublic::root(SEED); let cci = (2_u32).pow(31) + 13; let child_keys = ChildKeysPublic::nth_child(&root_keys, cci); diff --git a/lee/key_protocol/src/key_management/key_tree/mod.rs b/lee/key_protocol/src/key_management/key_tree/mod.rs index c15c09a5..ce1572b7 100644 --- a/lee/key_protocol/src/key_management/key_tree/mod.rs +++ b/lee/key_protocol/src/key_management/key_tree/mod.rs @@ -39,16 +39,7 @@ impl KeyTree { .try_into() .expect("SeedHolder seed is 64 bytes long"); - let root_keys = N::from_seed(seed_fit); - let account_id_map = root_keys - .account_ids() - .map(|id| (id, ChainIndex::root())) - .collect(); - - Self { - key_map: BTreeMap::from_iter([(ChainIndex::root(), root_keys)]), - account_id_map, - } + Self::new_from_root(N::from_seed(seed_fit)) } pub fn new_from_root(root: N) -> Self { @@ -63,6 +54,15 @@ impl KeyTree { } } + fn insert_child(&mut self, child_keys: N, chain_index: ChainIndex) -> ChainIndex { + for account_id in child_keys.account_ids() { + self.account_id_map.insert(account_id, chain_index.clone()); + } + self.key_map.insert(chain_index.clone(), child_keys); + + chain_index + } + pub fn generate_new_node(&mut self, parent_cci: &ChainIndex) -> Option { let parent_keys = self.key_map.get(parent_cci)?; let next_child_id = self @@ -71,14 +71,8 @@ impl KeyTree { let next_cci = parent_cci.nth_child(next_child_id); let child_keys = parent_keys.derive_child(next_child_id); - let account_ids = child_keys.account_ids(); - for account_id in account_ids { - self.account_id_map.insert(account_id, next_cci.clone()); - } - self.key_map.insert(next_cci.clone(), child_keys); - - Some(next_cci) + Some(self.insert_child(child_keys, next_cci)) } pub fn fill_node(&mut self, chain_index: &ChainIndex) -> Option { @@ -86,14 +80,8 @@ impl KeyTree { let child_id = *chain_index.chain().last()?; let child_keys = parent_keys.derive_child(child_id); - let account_ids = child_keys.account_ids(); - for account_id in account_ids { - self.account_id_map.insert(account_id, chain_index.clone()); - } - self.key_map.insert(chain_index.clone(), child_keys); - - Some(chain_index.clone()) + Some(self.insert_child(child_keys, chain_index.clone())) } #[must_use] @@ -200,24 +188,27 @@ impl KeyTree { } impl KeyTree { + /// Pairs `cci` with the account ID of the node stored at it. + fn account_id_for_cci(&self, cci: ChainIndex) -> Option<(lee::AccountId, ChainIndex)> { + let node = self.key_map.get(&cci)?; + let account_id = node.account_ids().next()?; + Some((account_id, cci)) + } + /// Generate a new public key node, returning the account ID and chain index. pub fn generate_new_public_node( &mut self, parent_cci: &ChainIndex, ) -> Option<(lee::AccountId, ChainIndex)> { let cci = self.generate_new_node(parent_cci)?; - let node = self.key_map.get(&cci)?; - let account_id = node.account_ids().next()?; - Some((account_id, cci)) + self.account_id_for_cci(cci) } /// Generate a new public key node using layered placement, returning the account ID and chain /// index. pub fn generate_new_public_node_layered(&mut self) -> Option<(lee::AccountId, ChainIndex)> { let cci = self.generate_new_node_layered()?; - let node = self.key_map.get(&cci)?; - let account_id = node.account_ids().next()?; - Some((account_id, cci)) + self.account_id_for_cci(cci) } /// Cleanup of non-initialized accounts in a public tree. @@ -322,6 +313,16 @@ impl KeyTree { } } +const fn split_hash(hash_value: &[u8; 64]) -> ([u8; 32], [u8; 32]) { + let first = *hash_value + .first_chunk::<32>() + .expect("hash_value is 64 bytes, must be safe to get first 32"); + let last = *hash_value + .last_chunk::<32>() + .expect("hash_value is 64 bytes, must be safe to get last 32"); + (first, last) +} + #[cfg(test)] mod tests { #![expect(clippy::shadow_unrelated, reason = "We don't care about this in tests")] diff --git a/lee/key_protocol/src/key_management/mod.rs b/lee/key_protocol/src/key_management/mod.rs index 459badf0..3a066fd9 100644 --- a/lee/key_protocol/src/key_management/mod.rs +++ b/lee/key_protocol/src/key_management/mod.rs @@ -25,8 +25,22 @@ impl KeyChain { #[must_use] pub fn new_os_random() -> Self { // Currently dropping SeedHolder at the end of initialization. - // Now entirely sure if we need it in the future. + // Not entirely sure if we need it in the future. let seed_holder = SeedHolder::new_os_random(); + + Self::from_seed_holder(&seed_holder) + } + + #[must_use] + pub fn new_mnemonic(passphrase: &str) -> (Self, bip39::Mnemonic) { + // Currently dropping SeedHolder at the end of initialization. + // Not entirely sure if we need it in the future. + let (seed_holder, mnemonic) = SeedHolder::new_mnemonic(passphrase); + + (Self::from_seed_holder(&seed_holder), mnemonic) + } + + fn from_seed_holder(seed_holder: &SeedHolder) -> Self { let secret_spending_key = seed_holder.produce_top_secret_key_holder(); let private_key_holder = secret_spending_key.produce_private_key_holder(None); @@ -42,29 +56,6 @@ impl KeyChain { } } - #[must_use] - pub fn new_mnemonic(passphrase: &str) -> (Self, bip39::Mnemonic) { - // Currently dropping SeedHolder at the end of initialization. - // Not entirely sure if we need it in the future. - let (seed_holder, mnemonic) = SeedHolder::new_mnemonic(passphrase); - let secret_spending_key = seed_holder.produce_top_secret_key_holder(); - - let private_key_holder = secret_spending_key.produce_private_key_holder(None); - - let nullifier_public_key = private_key_holder.generate_nullifier_public_key(); - let viewing_public_key = private_key_holder.generate_viewing_public_key(); - - ( - Self { - secret_spending_key, - private_key_holder, - nullifier_public_key, - viewing_public_key, - }, - mnemonic, - ) - } - #[must_use] pub fn calculate_shared_secret_receiver( &self, diff --git a/lee/key_protocol/src/key_management/secret_holders.rs b/lee/key_protocol/src/key_management/secret_holders.rs index 7bda4ffb..4035080d 100644 --- a/lee/key_protocol/src/key_management/secret_holders.rs +++ b/lee/key_protocol/src/key_management/secret_holders.rs @@ -43,16 +43,7 @@ pub struct PrivateKeyHolder { impl SeedHolder { #[must_use] pub fn new_os_random() -> Self { - let mut enthopy_bytes: [u8; 32] = [0; 32]; - OsRng.fill_bytes(&mut enthopy_bytes); - - let mnemonic = Mnemonic::from_entropy(&enthopy_bytes) - .expect("Enthropy must be a multiple of 32 bytes"); - let seed_wide = mnemonic.to_seed("mnemonic"); - - Self { - seed: seed_wide.to_vec(), - } + Self::new_mnemonic("mnemonic").0 } #[must_use] @@ -62,14 +53,8 @@ impl SeedHolder { let mnemonic = Mnemonic::from_entropy(&entropy_bytes).expect("Entropy must be a multiple of 32 bytes"); - let seed_wide = mnemonic.to_seed(passphrase); - ( - Self { - seed: seed_wide.to_vec(), - }, - mnemonic, - ) + (Self::from_mnemonic(&mnemonic, passphrase), mnemonic) } #[must_use] @@ -107,10 +92,7 @@ impl SecretSpendingKey { const SUFFIX_1: &[u8; 1] = &[1]; const SUFFIX_2: &[u8; 19] = &[0; 19]; - let index = match index { - None => 0_u32, - _ => index.expect("Expect a valid u32"), - }; + let index = index.unwrap_or(0); let mut hasher = sha2::Sha256::new(); hasher.update(PREFIX); @@ -129,10 +111,7 @@ impl SecretSpendingKey { const SUFFIX_1: &[u8; 1] = &[2]; const SUFFIX_2: &[u8; 19] = &[0; 19]; - let index = match index { - None => 0_u32, - _ => index.expect("Expect a valid u32"), - }; + let index = index.unwrap_or(0); let mut bytes: Vec = Vec::with_capacity(64); bytes.extend_from_slice(PREFIX); @@ -146,14 +125,7 @@ impl SecretSpendingKey { let full_seed = hmac_sha512::HMAC::mac(bytes, b"LEE_viewing_seed"); - ViewingSecretKey::new( - *full_seed - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32"), - *full_seed - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"), - ) + Self::generate_viewing_secret_key(full_seed) } #[must_use] @@ -201,7 +173,6 @@ impl PrivateKeyHolder { mod tests { use super::*; - // TODO? are these necessary? #[test] fn seed_generation_test() { let seed_holder = SeedHolder::new_os_random(); diff --git a/lee/state_machine/core/src/encoding.rs b/lee/state_machine/core/src/encoding.rs index 59df4b06..0fe3e08f 100644 --- a/lee/state_machine/core/src/encoding.rs +++ b/lee/state_machine/core/src/encoding.rs @@ -7,7 +7,7 @@ use std::io::Read as _; #[cfg(feature = "host")] use crate::Nullifier; #[cfg(feature = "host")] -use crate::encryption::EphemeralPublicKey; +use crate::encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN}; #[cfg(feature = "host")] use crate::error::LeeCoreError; use crate::{ @@ -168,7 +168,7 @@ impl EphemeralPublicKey { /// Deserializes an ML-KEM-768 ciphertext from a cursor. /// Reads exactly 1088 bytes — the fixed ciphertext size for ML-KEM-768. pub fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result { - let mut value = vec![0_u8; 1088]; + let mut value = vec![0_u8; ML_KEM_768_CIPHERTEXT_LEN]; cursor.read_exact(&mut value)?; Ok(Self(value)) } diff --git a/lee/state_machine/core/src/encryption/mod.rs b/lee/state_machine/core/src/encryption/mod.rs index 5fa80b60..47e0995d 100644 --- a/lee/state_machine/core/src/encryption/mod.rs +++ b/lee/state_machine/core/src/encryption/mod.rs @@ -12,13 +12,16 @@ use crate::{Commitment, account::Account, program::PrivateAccountKind}; #[cfg(feature = "host")] pub mod shared_key_derivation; +/// Length in bytes of an ML-KEM-768 ciphertext (the `EphemeralPublicKey` payload). +pub const ML_KEM_768_CIPHERTEXT_LEN: usize = 1088; + pub type Scalar = [u8; 32]; #[derive(Serialize, Deserialize, Clone, Copy)] pub struct SharedSecretKey(pub [u8; 32]); /// The ML-KEM-768 ciphertext produced during encapsulation; transmitted on-wire in place of the -/// former ECDH ephemeral public key. Always 1088 bytes for ML-KEM-768. +/// former ECDH ephemeral public key. Always `ML_KEM_768_CIPHERTEXT_LEN` (1088) bytes. #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct EphemeralPublicKey(pub Vec); diff --git a/lee/state_machine/core/src/encryption/shared_key_derivation.rs b/lee/state_machine/core/src/encryption/shared_key_derivation.rs index 5c982c6f..71cd8ab5 100644 --- a/lee/state_machine/core/src/encryption/shared_key_derivation.rs +++ b/lee/state_machine/core/src/encryption/shared_key_derivation.rs @@ -146,6 +146,7 @@ mod tests { use ml_kem::KeyExport as _; use super::*; + use crate::ML_KEM_768_CIPHERTEXT_LEN; #[test] fn encapsulate_decapsulate_round_trip() { @@ -164,7 +165,11 @@ mod tests { let receiver_ss = SharedSecretKey::decapsulate(&epk, &d, &z).unwrap(); assert_eq!(sender_ss.0, receiver_ss.0, "shared secrets must match"); - assert_eq!(epk.0.len(), 1088, "ML-KEM-768 ciphertext is 1088 bytes"); + assert_eq!( + epk.0.len(), + ML_KEM_768_CIPHERTEXT_LEN, + "ML-KEM-768 ciphertext is 1088 bytes" + ); assert_eq!( ek.0.len(), 1184, @@ -185,7 +190,7 @@ mod tests { ); // Too long — 1089 bytes instead of 1088. - let long_epk = EphemeralPublicKey(vec![42_u8; 1089]); + let long_epk = EphemeralPublicKey(vec![42_u8; ML_KEM_768_CIPHERTEXT_LEN + 1]); assert!( SharedSecretKey::decapsulate(&long_epk, &d, &z).is_none(), "long EphemeralPublicKey must return None" diff --git a/lee/state_machine/core/src/lib.rs b/lee/state_machine/core/src/lib.rs index 9ad2858e..900de667 100644 --- a/lee/state_machine/core/src/lib.rs +++ b/lee/state_machine/core/src/lib.rs @@ -11,7 +11,8 @@ pub use commitment::{ compute_digest_for_path, }; pub use encryption::{ - EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, SharedSecretKey, ViewTag, + EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, + SharedSecretKey, ViewTag, }; pub use nullifier::{Identifier, Nullifier, NullifierPublicKey, NullifierSecretKey}; pub use program::PrivateAccountKind; diff --git a/lee/state_machine/src/state.rs b/lee/state_machine/src/state.rs index c399cea1..57cd0bf6 100644 --- a/lee/state_machine/src/state.rs +++ b/lee/state_machine/src/state.rs @@ -330,7 +330,7 @@ pub mod tests { BlockId, Commitment, EncryptedAccountData, InputAccountIdentity, Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, Timestamp, account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data}, - encryption::{EphemeralPublicKey, ViewingPublicKey}, + encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, ViewingPublicKey}, program::{ BlockValidityWindow, ExecutionValidationError, MAX_NUMBER_CHAINED_CALLS, PdaSeed, ProgramId, TimestampValidityWindow, WrappedBalanceSum, @@ -4379,7 +4379,7 @@ pub mod tests { vec![ InputAccountIdentity::Public, InputAccountIdentity::PrivatePdaUpdate { - epk: EphemeralPublicKey(vec![12_u8; 1088]), + epk: EphemeralPublicKey(vec![12_u8; ML_KEM_768_CIPHERTEXT_LEN]), view_tag: EncryptedAccountData::compute_view_tag( &alice_npk, &alice_keys.vpk(),