From 46c9cee8c58d923fd4c52330a4fd2ac6854976e7 Mon Sep 17 00:00:00 2001 From: Artem Gureev Date: Tue, 7 Jul 2026 16:27:34 +0000 Subject: [PATCH] feat!: key note encryption on the nullifier, drop the output index --- lee/privacy_preserving_circuit/src/output.rs | 16 +------ lee/state_machine/core/src/encryption/mod.rs | 47 +++++++------------ .../circuit/tests.rs | 12 ++--- .../privacy_preserving_transaction/message.rs | 5 +- lez/wallet/src/lib.rs | 26 ++-------- lez/wallet/src/storage/key_chain.rs | 10 ++-- .../benches/primitives.rs | 11 ++--- 7 files changed, 39 insertions(+), 88 deletions(-) diff --git a/lee/privacy_preserving_circuit/src/output.rs b/lee/privacy_preserving_circuit/src/output.rs index 9ae530b2..5105005c 100644 --- a/lee/privacy_preserving_circuit/src/output.rs +++ b/lee/privacy_preserving_circuit/src/output.rs @@ -32,7 +32,6 @@ pub fn compute_circuit_output( "Invalid account_identities length" ); - let mut output_index = 0; for (pos, (account_identity, (pre_state, post_state))) in account_identities.iter().zip(states_iter).enumerate() { @@ -71,7 +70,6 @@ pub fn compute_circuit_output( emit_private_output( &mut output, - &mut output_index, post_state, &account_id, &PrivateAccountKind::Regular(*identifier), @@ -109,7 +107,6 @@ pub fn compute_circuit_output( emit_private_output( &mut output, - &mut output_index, post_state, &account_id, &PrivateAccountKind::Regular(*identifier), @@ -149,7 +146,6 @@ pub fn compute_circuit_output( emit_private_output( &mut output, - &mut output_index, post_state, &account_id, &PrivateAccountKind::Regular(*identifier), @@ -197,7 +193,6 @@ pub fn compute_circuit_output( let view_tag = EncryptedAccountData::compute_view_tag(npk, vpk); emit_private_output( &mut output, - &mut output_index, post_state, &account_id, &PrivateAccountKind::Pda { @@ -245,7 +240,6 @@ pub fn compute_circuit_output( .expect("PrivatePdaUpdate position must be in pda_seed_by_position"); emit_private_output( &mut output, - &mut output_index, post_state, &account_id, &PrivateAccountKind::Pda { @@ -294,7 +288,6 @@ fn emit_dummy_output(output: &mut PrivacyPreservingCircuitOutput, dummy: DummyIn )] fn emit_private_output( output: &mut PrivacyPreservingCircuitOutput, - output_index: &mut u32, post_state: Account, account_id: &AccountId, kind: &PrivateAccountKind, @@ -304,8 +297,6 @@ fn emit_private_output( new_nullifier: (Nullifier, CommitmentSetDigest), new_nonce: Nonce, ) { - output.new_nullifiers.push(new_nullifier); - let mut post_with_updated_nonce = post_state; post_with_updated_nonce.nonce = new_nonce; @@ -318,10 +309,10 @@ fn emit_private_output( &post_with_updated_nonce, kind, &shared_secret, - &commitment_post, - *output_index, + &new_nullifier.0, ); + output.new_nullifiers.push(new_nullifier); output.new_commitments.push(commitment_post); output .encrypted_private_post_states @@ -330,9 +321,6 @@ fn emit_private_output( epk, view_tag, }); - *output_index = output_index - .checked_add(1) - .unwrap_or_else(|| panic!("Too many private accounts, output index overflow")); } fn compute_update_nullifier_and_set_digest( diff --git a/lee/state_machine/core/src/encryption/mod.rs b/lee/state_machine/core/src/encryption/mod.rs index f503279d..3b244b5c 100644 --- a/lee/state_machine/core/src/encryption/mod.rs +++ b/lee/state_machine/core/src/encryption/mod.rs @@ -7,7 +7,7 @@ use risc0_zkvm::sha::{Impl, Sha256 as _}; use serde::{Deserialize, Serialize}; pub use shared_key_derivation::{MlKem768EncapsulationKey, ViewingPublicKey}; -use crate::{Commitment, account::Account, program::PrivateAccountKind}; +use crate::{Nullifier, account::Account, program::PrivateAccountKind}; pub mod shared_key_derivation; /// Length in bytes of an ML-KEM-768 ciphertext (the `EphemeralPublicKey` payload). @@ -114,39 +114,32 @@ impl EncryptionScheme { account: &Account, kind: &PrivateAccountKind, shared_secret: &SharedSecretKey, - commitment: &Commitment, - output_index: u32, + nullifier: &Nullifier, ) -> Ciphertext { // Plaintext: PrivateAccountKind::HEADER_LEN bytes header || account bytes. // Both variants produce the same header length — see PrivateAccountKind::to_header_bytes. let mut buffer = kind.to_header_bytes().to_vec(); buffer.extend_from_slice(&account.to_bytes()); - Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index); + Self::symmetric_transform(&mut buffer, shared_secret, nullifier); Ciphertext(buffer) } fn symmetric_transform( buffer: &mut [u8], shared_secret: &SharedSecretKey, - commitment: &Commitment, - output_index: u32, + nullifier: &Nullifier, ) { - let key = Self::kdf(shared_secret, commitment, output_index); + let key = Self::kdf(shared_secret, nullifier); let mut cipher = ChaCha20::new(&key.into(), &[0; 12].into()); cipher.apply_keystream(buffer); } - fn kdf( - shared_secret: &SharedSecretKey, - commitment: &Commitment, - output_index: u32, - ) -> [u8; 32] { + fn kdf(shared_secret: &SharedSecretKey, nullifier: &Nullifier) -> [u8; 32] { const PREFIX: &[u8; 20] = b"LEE/v0.2/KDF-SHA256/"; - let mut bytes = [0_u8; 20 + 32 + 32 + 4]; + let mut bytes = [0_u8; 20 + 32 + 32]; bytes[0..20].copy_from_slice(PREFIX); bytes[20..52].copy_from_slice(&shared_secret.0); - bytes[52..84].copy_from_slice(&commitment.to_byte_array()); - bytes[84..88].copy_from_slice(&output_index.to_le_bytes()); + bytes[52..84].copy_from_slice(&nullifier.to_byte_array()); Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap() } @@ -160,12 +153,11 @@ impl EncryptionScheme { pub fn decrypt( ciphertext: &Ciphertext, shared_secret: &SharedSecretKey, - commitment: &Commitment, - output_index: u32, + nullifier: &Nullifier, ) -> Option<(PrivateAccountKind, Account)> { use std::io::Cursor; let mut buffer = ciphertext.0.clone(); - Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index); + Self::symmetric_transform(&mut buffer, shared_secret, nullifier); if buffer.len() < PrivateAccountKind::HEADER_LEN { return None; @@ -180,8 +172,7 @@ impl EncryptionScheme { println!( "Failed to decode {ciphertext:?} \n with secret {:?} ,\n - commitment {commitment:?} ,\n - and output_index {output_index} ,\n + nullifier {nullifier:?} ,\n with error {err:?}", shared_secret.0 ); @@ -203,14 +194,13 @@ mod tests { fn encrypt_same_length_for_account_and_pda() { let account = Account::default(); let secret = SharedSecretKey([0_u8; 32]); - let commitment = crate::Commitment::new(&AccountId::new([0_u8; 32]), &Account::default()); + let nullifier = Nullifier::for_account_initialization(&AccountId::new([0_u8; 32])); let account_ct = EncryptionScheme::encrypt( &account, &PrivateAccountKind::Regular(42), &secret, - &commitment, - 0, + &nullifier, ); let pda_ct = EncryptionScheme::encrypt( &account, @@ -220,8 +210,7 @@ mod tests { identifier: 42, }, &secret, - &commitment, - 0, + &nullifier, ); assert_eq!(account_ct.0.len(), pda_ct.0.len()); @@ -245,11 +234,11 @@ mod tests { ..Account::default() }; let kind = PrivateAccountKind::Regular(0); - let commitment = crate::Commitment::new(&AccountId::new([7_u8; 32]), &account); + let nullifier = Nullifier::for_account_initialization(&AccountId::new([7_u8; 32])); - let ct = EncryptionScheme::encrypt(&account, &kind, &sender_ss, &commitment, 0); + let ct = EncryptionScheme::encrypt(&account, &kind, &sender_ss, &nullifier); let (decoded_kind, decoded_account) = - EncryptionScheme::decrypt(&ct, &receiver_ss, &commitment, 0) + EncryptionScheme::decrypt(&ct, &receiver_ss, &nullifier) .expect("decryption must succeed with correct shared secret"); assert_eq!(decoded_account, account); @@ -257,7 +246,7 @@ mod tests { // Wrong shared secret must not decrypt correctly. let wrong_ss = SharedSecretKey([0_u8; 32]); - let bad = EncryptionScheme::decrypt(&ct, &wrong_ss, &commitment, 0); + let bad = EncryptionScheme::decrypt(&ct, &wrong_ss, &nullifier); assert!( bad.is_none() || bad.is_some_and(|(_, a)| a.balance != 999), "wrong shared secret must not produce the correct plaintext" diff --git a/lee/state_machine/src/privacy_preserving_transaction/circuit/tests.rs b/lee/state_machine/src/privacy_preserving_transaction/circuit/tests.rs index 225d542a..c94675a9 100644 --- a/lee/state_machine/src/privacy_preserving_transaction/circuit/tests.rs +++ b/lee/state_machine/src/privacy_preserving_transaction/circuit/tests.rs @@ -26,8 +26,7 @@ fn decrypt_kind( let (kind, _) = EncryptionScheme::decrypt( &output.encrypted_private_post_states[idx].ciphertext, ssk, - &output.new_commitments[idx], - u32::try_from(idx).expect("idx fits in u32"), + &output.new_nullifiers[idx].0, ) .unwrap(); kind @@ -114,8 +113,7 @@ fn prove_privacy_preserving_execution_circuit_public_and_private_pre_accounts() let (_identifier, recipient_post) = EncryptionScheme::decrypt( &output.encrypted_private_post_states[0].ciphertext, &shared_secret, - &output.new_commitments[0], - 0, + &output.new_nullifiers[0].0, ) .unwrap(); assert_eq!(recipient_post, expected_recipient_post); @@ -227,8 +225,7 @@ fn prove_privacy_preserving_execution_circuit_fully_private() { let (_identifier, sender_post) = EncryptionScheme::decrypt( &output.encrypted_private_post_states[0].ciphertext, &shared_secret_1, - &expected_new_commitments[0], - 0, + &output.new_nullifiers[0].0, ) .unwrap(); assert_eq!(sender_post, expected_private_account_1); @@ -236,8 +233,7 @@ fn prove_privacy_preserving_execution_circuit_fully_private() { let (_identifier, recipient_post) = EncryptionScheme::decrypt( &output.encrypted_private_post_states[1].ciphertext, &shared_secret_2, - &expected_new_commitments[1], - 1, + &output.new_nullifiers[1].0, ) .unwrap(); assert_eq!(recipient_post, expected_private_account_2); diff --git a/lee/state_machine/src/privacy_preserving_transaction/message.rs b/lee/state_machine/src/privacy_preserving_transaction/message.rs index 648fcb6a..a7b75caf 100644 --- a/lee/state_machine/src/privacy_preserving_transaction/message.rs +++ b/lee/state_machine/src/privacy_preserving_transaction/message.rs @@ -199,15 +199,14 @@ pub mod tests { let vpk = ViewingPublicKey::from_seed(&[2_u8; 32], &[3_u8; 32]); let account = Account::default(); let account_id = lee_core::account::AccountId::for_regular_private_account(&npk, &vpk, 0); - let commitment = Commitment::new(&account_id, &account); + let nullifier = Nullifier::for_account_initialization(&account_id); let (shared_secret, epk) = SharedSecretKey::encapsulate_deterministic(&vpk, &EphemeralSecretKey([0_u8; 32])); let ciphertext = EncryptionScheme::encrypt( &account, &PrivateAccountKind::Regular(0), &shared_secret, - &commitment, - 2, + &nullifier, ); let encrypted_account_data = EncryptedAccountData::new(ciphertext.clone(), &npk, &vpk, epk.clone()); diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index 27839eb4..75c56fec 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -498,15 +498,11 @@ impl WalletCore { match acc_decode_data { AccDecodeData::Decode(secret, acc_account_id) => { let acc_ead = tx.message.encrypted_private_post_states[output_index].clone(); - let acc_comm = tx.message.new_commitments[output_index].clone(); let (kind, res_acc) = lee_core::EncryptionScheme::decrypt( &acc_ead.ciphertext, secret, - &acc_comm, - output_index - .try_into() - .expect("Output index is expected to fit in u32"), + &tx.message.new_nullifiers[output_index].0, ) .unwrap(); @@ -758,7 +754,7 @@ impl WalletCore { &key_chain.nullifier_public_key, &key_chain.viewing_public_key, ); - let new_commitments = &message.new_commitments; + let new_nullifiers = &message.new_nullifiers; message .encrypted_private_post_states @@ -772,18 +768,11 @@ impl WalletCore { }) .filter_map(move |(ciph_id, encrypted_data)| { let ciphertext = &encrypted_data.ciphertext; - let commitment = &new_commitments[ciph_id]; + let nullifier = &new_nullifiers[ciph_id].0; let shared_secret = key_chain.calculate_shared_secret_receiver(&encrypted_data.epk)?; - lee_core::EncryptionScheme::decrypt( - ciphertext, - &shared_secret, - commitment, - ciph_id - .try_into() - .expect("Ciphertext ID is expected to fit in u32"), - ) + lee_core::EncryptionScheme::decrypt(ciphertext, &shared_secret, nullifier) .map(|(kind, res_acc)| { let npk = &key_chain.nullifier_public_key; let account_id = lee::AccountId::for_private_account( @@ -852,15 +841,10 @@ impl WalletCore { else { continue; }; - let commitment = &message.new_commitments[ciph_id]; - if let Some((_kind, new_acc)) = lee_core::EncryptionScheme::decrypt( &encrypted_data.ciphertext, &shared_secret, - commitment, - ciph_id - .try_into() - .expect("Ciphertext ID is expected to fit in u32"), + &message.new_nullifiers[ciph_id].0, ) { info!("Synced shared account {account_id:#?} with new state {new_acc:#?}"); index.track(account_id, &new_acc, &nsk); diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs index 6e41eb25..f03b035d 100644 --- a/lez/wallet/src/storage/key_chain.rs +++ b/lez/wallet/src/storage/key_chain.rs @@ -410,8 +410,6 @@ impl UserKeyChain { i: usize, ) -> Option { let encrypted = &message.encrypted_private_post_states[i]; - let commitment = &message.new_commitments[i]; - let ciph_id = u32::try_from(i).ok()?; let (nsk, secret, is_shared) = if let Some(entry) = self.shared_private_account(account_id) { @@ -435,7 +433,7 @@ impl UserKeyChain { }; let (kind, new_account) = - EncryptionScheme::decrypt(&encrypted.ciphertext, &secret, commitment, ciph_id)?; + EncryptionScheme::decrypt(&encrypted.ciphertext, &secret, &message.new_nullifiers[i].0)?; let new_nullifier = NullifierIndex::next_update_nullifier(account_id, &new_account, &nsk); if is_shared { @@ -892,8 +890,7 @@ mod tests { &new_account, &PrivateAccountKind::Regular(identifier), &sender_ss, - &new_commitment, - 0, + &old_nullifier, ); let note = EncryptedAccountData::new( ciphertext, @@ -963,8 +960,7 @@ mod tests { &new_account, &PrivateAccountKind::Regular(identifier), &sender_ss, - &new_commitment, - 0, + &old_nullifier, ); let note = EncryptedAccountData::new(ciphertext, &npk, &vpk, epk); let message = Message { diff --git a/tools/crypto_primitives_bench/benches/primitives.rs b/tools/crypto_primitives_bench/benches/primitives.rs index 9a42305c..85b799cb 100644 --- a/tools/crypto_primitives_bench/benches/primitives.rs +++ b/tools/crypto_primitives_bench/benches/primitives.rs @@ -11,7 +11,7 @@ use std::time::Duration; use criterion::{Criterion, criterion_group, criterion_main}; use key_protocol::key_management::KeyChain; use lee_core::{ - Commitment, EncryptionScheme, SharedSecretKey, + EncryptionScheme, Nullifier, SharedSecretKey, account::{Account, AccountId}, program::PrivateAccountKind, }; @@ -50,19 +50,18 @@ fn bench_encryption(c: &mut Criterion) { let account = Account::default(); let account_id = AccountId::for_regular_private_account(&npk, &recipient_kc.viewing_public_key, 0); - let commitment = Commitment::new(&account_id, &account); + let nullifier = Nullifier::for_account_initialization(&account_id); let (shared, _epk) = SharedSecretKey::encapsulate(&recipient_kc.viewing_public_key); let kind = PrivateAccountKind::Regular(0_u128); - let output_index: u32 = 0; let mut g = c.benchmark_group("encryption"); g.sample_size(50).noise_threshold(0.05); g.bench_function("encrypt", |b| { - b.iter(|| EncryptionScheme::encrypt(&account, &kind, &shared, &commitment, output_index)); + b.iter(|| EncryptionScheme::encrypt(&account, &kind, &shared, &nullifier)); }); - let ct = EncryptionScheme::encrypt(&account, &kind, &shared, &commitment, output_index); + let ct = EncryptionScheme::encrypt(&account, &kind, &shared, &nullifier); g.bench_function("decrypt", |b| { - b.iter(|| EncryptionScheme::decrypt(&ct, &shared, &commitment, output_index)); + b.iter(|| EncryptionScheme::decrypt(&ct, &shared, &nullifier)); }); g.finish(); }