diff --git a/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin b/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin index ebc8d365..c0645c49 100644 Binary files a/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin and b/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin differ diff --git a/artifacts/lez/programs/amm.bin b/artifacts/lez/programs/amm.bin index 3940bbd2..a23b5517 100644 Binary files a/artifacts/lez/programs/amm.bin and b/artifacts/lez/programs/amm.bin differ diff --git a/artifacts/lez/programs/associated_token_account.bin b/artifacts/lez/programs/associated_token_account.bin index 5313badf..99076825 100644 Binary files a/artifacts/lez/programs/associated_token_account.bin and b/artifacts/lez/programs/associated_token_account.bin differ diff --git a/artifacts/lez/programs/authenticated_transfer.bin b/artifacts/lez/programs/authenticated_transfer.bin index 5b22fd1a..ee8929dd 100644 Binary files a/artifacts/lez/programs/authenticated_transfer.bin and b/artifacts/lez/programs/authenticated_transfer.bin differ diff --git a/artifacts/lez/programs/bridge.bin b/artifacts/lez/programs/bridge.bin index 061cb7c2..b887af29 100644 Binary files a/artifacts/lez/programs/bridge.bin and b/artifacts/lez/programs/bridge.bin differ diff --git a/artifacts/lez/programs/bridge_lock.bin b/artifacts/lez/programs/bridge_lock.bin index fcec36f0..87da00db 100644 Binary files a/artifacts/lez/programs/bridge_lock.bin and b/artifacts/lez/programs/bridge_lock.bin differ diff --git a/artifacts/lez/programs/clock.bin b/artifacts/lez/programs/clock.bin index b402bc20..e94c08bf 100644 Binary files a/artifacts/lez/programs/clock.bin and b/artifacts/lez/programs/clock.bin differ diff --git a/artifacts/lez/programs/cross_zone_inbox.bin b/artifacts/lez/programs/cross_zone_inbox.bin index 3c7de003..2cbd007d 100644 Binary files a/artifacts/lez/programs/cross_zone_inbox.bin and b/artifacts/lez/programs/cross_zone_inbox.bin differ diff --git a/artifacts/lez/programs/cross_zone_outbox.bin b/artifacts/lez/programs/cross_zone_outbox.bin index 8df676a7..19dfce4b 100644 Binary files a/artifacts/lez/programs/cross_zone_outbox.bin and b/artifacts/lez/programs/cross_zone_outbox.bin differ diff --git a/artifacts/lez/programs/faucet.bin b/artifacts/lez/programs/faucet.bin index ded3feb0..f0c1d99e 100644 Binary files a/artifacts/lez/programs/faucet.bin and b/artifacts/lez/programs/faucet.bin differ diff --git a/artifacts/lez/programs/pinata.bin b/artifacts/lez/programs/pinata.bin index 24383422..df02776e 100644 Binary files a/artifacts/lez/programs/pinata.bin and b/artifacts/lez/programs/pinata.bin differ diff --git a/artifacts/lez/programs/pinata_token.bin b/artifacts/lez/programs/pinata_token.bin index aa88de51..e787a117 100644 Binary files a/artifacts/lez/programs/pinata_token.bin and b/artifacts/lez/programs/pinata_token.bin differ diff --git a/artifacts/lez/programs/ping_receiver.bin b/artifacts/lez/programs/ping_receiver.bin index 85de2df9..7454d1ff 100644 Binary files a/artifacts/lez/programs/ping_receiver.bin and b/artifacts/lez/programs/ping_receiver.bin differ diff --git a/artifacts/lez/programs/ping_sender.bin b/artifacts/lez/programs/ping_sender.bin index fd93cd74..09240097 100644 Binary files a/artifacts/lez/programs/ping_sender.bin and b/artifacts/lez/programs/ping_sender.bin differ diff --git a/artifacts/lez/programs/token.bin b/artifacts/lez/programs/token.bin index fa57f5f6..109cd9c4 100644 Binary files a/artifacts/lez/programs/token.bin and b/artifacts/lez/programs/token.bin differ diff --git a/artifacts/lez/programs/vault.bin b/artifacts/lez/programs/vault.bin index 1e8bf6b7..45b89776 100644 Binary files a/artifacts/lez/programs/vault.bin and b/artifacts/lez/programs/vault.bin differ diff --git a/artifacts/lez/programs/wrapped_token.bin b/artifacts/lez/programs/wrapped_token.bin index d0362994..777d20fc 100644 Binary files a/artifacts/lez/programs/wrapped_token.bin and b/artifacts/lez/programs/wrapped_token.bin differ diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index 8136f9df..a472c260 100644 --- a/integration_tests/tests/auth_transfer/private.rs +++ b/integration_tests/tests/auth_transfer/private.rs @@ -81,7 +81,7 @@ async fn private_transfer_to_foreign_account() -> Result<()> { .context("Failed to get private account commitment for sender")?; let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash).await; - assert_eq!(tx.message.new_commitments[0], new_commitment1); + assert!(tx.message.new_commitments.contains(&new_commitment1)); for commitment in tx.message.new_commitments { assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await); @@ -169,7 +169,7 @@ async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> { .wallet() .get_private_account_commitment(from) .context("Failed to get private account commitment for sender")?; - assert_eq!(tx.message.new_commitments[0], sender_commitment); + assert!(tx.message.new_commitments.contains(&sender_commitment)); for commitment in tx.message.new_commitments { assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await); @@ -245,13 +245,9 @@ async fn shielded_transfer_to_foreign_account() -> Result<()> { let acc_1_balance = account_balance(&ctx, from).await?; - assert!( - verify_commitment_is_in_state( - tx.message.new_commitments[0].clone(), - ctx.sequencer_client() - ) - .await - ); + for commitment in tx.message.new_commitments { + assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await); + } assert_eq!(acc_1_balance, 9900); diff --git a/integration_tests/tests/keys.rs b/integration_tests/tests/keys.rs index f0a66e0e..fb958ec0 100644 --- a/integration_tests/tests/keys.rs +++ b/integration_tests/tests/keys.rs @@ -71,7 +71,7 @@ async fn sync_private_account_with_non_zero_chain_index() -> Result<()> { .wallet() .get_private_account_commitment(from) .context("Failed to get private account commitment for sender")?; - assert_eq!(tx.message.new_commitments[0], new_commitment1); + assert!(tx.message.new_commitments.contains(&new_commitment1)); for commitment in tx.message.new_commitments { assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await); diff --git a/lee/privacy_preserving_circuit/Cargo.toml b/lee/privacy_preserving_circuit/Cargo.toml index 26dfa132..a4f74644 100644 --- a/lee/privacy_preserving_circuit/Cargo.toml +++ b/lee/privacy_preserving_circuit/Cargo.toml @@ -15,3 +15,6 @@ workspace = true [dependencies] lee_core.workspace = true risc0-zkvm.workspace = true + +[dev-dependencies] +lee_core = { workspace = true, features = ["host"] } diff --git a/lee/privacy_preserving_circuit/src/output.rs b/lee/privacy_preserving_circuit/src/output.rs index 9ae530b2..ab00f8c3 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 { @@ -267,9 +261,24 @@ pub fn compute_circuit_output( emit_dummy_output(&mut output, dummy); } + obfuscate_output_ordering(&mut output); + output } +fn obfuscate_output_ordering(output: &mut PrivacyPreservingCircuitOutput) { + output + .new_commitments + .sort_unstable_by_key(Commitment::to_byte_array); + + let mut notes: Vec<_> = core::mem::take(&mut output.new_nullifiers) + .into_iter() + .zip(core::mem::take(&mut output.encrypted_private_post_states)) + .collect(); + notes.sort_unstable_by_key(|((nullifier, _), _)| nullifier.to_byte_array()); + (output.new_nullifiers, output.encrypted_private_post_states) = notes.into_iter().unzip(); +} + fn emit_dummy_output(output: &mut PrivacyPreservingCircuitOutput, dummy: DummyInput) { // Note: the nullifiers and commitments are generated from seeds. // The prover is responsible for their randomness. @@ -294,7 +303,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 +312,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 +324,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 +336,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( @@ -346,3 +349,83 @@ fn compute_update_nullifier_and_set_digest( let nullifier = Nullifier::for_account_update(&commitment_pre, nsk); (nullifier, set_digest) } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use lee_core::{DUMMY_COMMITMENT_HASH, EphemeralPublicKey}; + + use super::*; + + fn note(tag: u8) -> (Nullifier, Commitment, EncryptedAccountData) { + let nullifier = Nullifier::for_dummy(&[tag; 32]); + let commitment = Commitment::for_dummy(&nullifier, &[tag; 32]); + let ciphertext = EncryptionScheme::encrypt( + &Account::default(), + &PrivateAccountKind::Regular(0), + &SharedSecretKey([0; 32]), + &nullifier, + ); + let encrypted = EncryptedAccountData { + ciphertext, + epk: EphemeralPublicKey(vec![tag]), + view_tag: 0, + }; + (nullifier, commitment, encrypted) + } + + #[test] + fn obfuscate_byte_sorts_commitments_and_nullifiers() { + let mut output = PrivacyPreservingCircuitOutput::default(); + for tag in 0..3 { + let (nullifier, commitment, encrypted) = note(tag); + output + .new_nullifiers + .push((nullifier, DUMMY_COMMITMENT_HASH)); + output.new_commitments.push(commitment); + output.encrypted_private_post_states.push(encrypted); + } + + obfuscate_output_ordering(&mut output); + + assert!( + output + .new_commitments + .is_sorted_by_key(Commitment::to_byte_array) + ); + assert!( + output + .new_nullifiers + .is_sorted_by_key(|(nullifier, _)| nullifier.to_byte_array()) + ); + } + + #[test] + fn obfuscate_keeps_each_nullifier_with_its_ciphertext() { + let mut output = PrivacyPreservingCircuitOutput::default(); + for tag in 0..3 { + let (nullifier, _, encrypted) = note(tag); + output + .new_nullifiers + .push((nullifier, DUMMY_COMMITMENT_HASH)); + output.encrypted_private_post_states.push(encrypted); + } + let paired: HashMap<[u8; 32], EphemeralPublicKey> = output + .new_nullifiers + .iter() + .zip(&output.encrypted_private_post_states) + .map(|((nullifier, _), note)| (nullifier.to_byte_array(), note.epk.clone())) + .collect(); + + obfuscate_output_ordering(&mut output); + + for ((nullifier, _), note) in output + .new_nullifiers + .iter() + .zip(&output.encrypted_private_post_states) + { + assert_eq!(paired[&nullifier.to_byte_array()], note.epk); + } + } +} diff --git a/lee/state_machine/core/src/circuit_io.rs b/lee/state_machine/core/src/circuit_io.rs index 9b283eae..b5e0986f 100644 --- a/lee/state_machine/core/src/circuit_io.rs +++ b/lee/state_machine/core/src/circuit_io.rs @@ -148,7 +148,7 @@ impl InputAccountIdentity { } #[derive(Serialize, Deserialize)] -#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] +#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq, Default))] pub struct PrivacyPreservingCircuitOutput { pub public_pre_states: Vec, pub public_post_states: Vec, diff --git a/lee/state_machine/core/src/encoding.rs b/lee/state_machine/core/src/encoding.rs index 0fe3e08f..e9b4a845 100644 --- a/lee/state_machine/core/src/encoding.rs +++ b/lee/state_machine/core/src/encoding.rs @@ -97,11 +97,6 @@ impl NullifierPublicKey { #[cfg(feature = "host")] impl Nullifier { - #[must_use] - pub const fn to_byte_array(&self) -> [u8; 32] { - self.0 - } - #[cfg(feature = "host")] #[must_use] pub const fn from_byte_array(bytes: [u8; 32]) -> Self { diff --git a/lee/state_machine/core/src/encryption/mod.rs b/lee/state_machine/core/src/encryption/mod.rs index f503279d..13ad7f5e 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,11 +246,19 @@ 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_via_ss = EncryptionScheme::decrypt(&ct, &wrong_ss, &nullifier); assert!( - bad.is_none() || bad.is_some_and(|(_, a)| a.balance != 999), + bad_via_ss.is_none() || bad_via_ss.is_some_and(|(_, a)| a.balance != 999), "wrong shared secret must not produce the correct plaintext" ); + + // Wrong nullifier must not decrypt correctly. + let wrong_nullifier = Nullifier::for_account_initialization(&AccountId::new([9; 32])); + let bad_via_nlf = EncryptionScheme::decrypt(&ct, &receiver_ss, &wrong_nullifier); + assert!( + bad_via_nlf.is_none() || bad_via_nlf.is_some_and(|(_, a)| a.balance != 999), + "wrong nullifier must not produce the correct plaintext" + ); } #[test] diff --git a/lee/state_machine/core/src/nullifier.rs b/lee/state_machine/core/src/nullifier.rs index 336f275f..59e7b5c2 100644 --- a/lee/state_machine/core/src/nullifier.rs +++ b/lee/state_machine/core/src/nullifier.rs @@ -117,6 +117,11 @@ impl Nullifier { bytes.extend_from_slice(nullifier_seed); Self(Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap()) } + + #[must_use] + pub const fn to_byte_array(&self) -> [u8; 32] { + self.0 + } } #[cfg(test)] 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 fb66186b..d5663b6d 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); @@ -220,24 +218,41 @@ fn prove_privacy_preserving_execution_circuit_fully_private() { assert!(proof.is_valid_for(&output)); assert!(output.public_pre_states.is_empty()); assert!(output.public_post_states.is_empty()); + let sender_nullifier = expected_new_nullifiers[0].0; + let recipient_nullifier = expected_new_nullifiers[1].0; + + let mut expected_new_commitments = expected_new_commitments; + expected_new_commitments.sort_unstable_by_key(Commitment::to_byte_array); assert_eq!(output.new_commitments, expected_new_commitments); + + let mut expected_new_nullifiers = expected_new_nullifiers; + expected_new_nullifiers.sort_unstable_by_key(|(nullifier, _)| nullifier.to_byte_array()); assert_eq!(output.new_nullifiers, expected_new_nullifiers); + assert_eq!(output.encrypted_private_post_states.len(), 2); + let sender_slot = output + .new_nullifiers + .iter() + .position(|(nullifier, _)| *nullifier == sender_nullifier) + .unwrap(); let (_identifier, sender_post) = EncryptionScheme::decrypt( - &output.encrypted_private_post_states[0].ciphertext, + &output.encrypted_private_post_states[sender_slot].ciphertext, &shared_secret_1, - &expected_new_commitments[0], - 0, + &output.new_nullifiers[sender_slot].0, ) .unwrap(); assert_eq!(sender_post, expected_private_account_1); + let recipient_slot = output + .new_nullifiers + .iter() + .position(|(nullifier, _)| *nullifier == recipient_nullifier) + .unwrap(); let (_identifier, recipient_post) = EncryptionScheme::decrypt( - &output.encrypted_private_post_states[1].ciphertext, + &output.encrypted_private_post_states[recipient_slot].ciphertext, &shared_secret_2, - &expected_new_commitments[1], - 1, + &output.new_nullifiers[recipient_slot].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 c1f2615c..3b6704ff 100644 --- a/lee/state_machine/src/privacy_preserving_transaction/message.rs +++ b/lee/state_machine/src/privacy_preserving_transaction/message.rs @@ -227,15 +227,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 595d1a15..86f456d6 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -556,21 +556,18 @@ impl WalletCore { "Decode mask has {} entries but the transaction has {note_count} notes", acc_decode_mask.len(), ); - for (output_index, acc_decode_data) in acc_decode_mask.iter().enumerate() { + for acc_decode_data in acc_decode_mask { 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"), - ) - .unwrap(); + let Some(output_index) = self + .storage + .key_chain() + .locate_spend(*acc_account_id, &tx.message) + else { + continue; + }; + let (kind, res_acc) = + decrypt_note_at(&tx.message, output_index, secret).unwrap(); println!("Received new acc {res_acc:#?}"); @@ -826,8 +823,6 @@ impl WalletCore { &key_chain.nullifier_public_key, &key_chain.viewing_public_key, ); - let new_commitments = &message.new_commitments; - message .encrypted_private_post_states .iter() @@ -839,20 +834,10 @@ impl WalletCore { !handled.contains(ciph_id) && encrypted_data.view_tag == view_tag }) .filter_map(move |(ciph_id, encrypted_data)| { - let ciphertext = &encrypted_data.ciphertext; - let commitment = &new_commitments[ciph_id]; 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"), - ) - .map(|(kind, res_acc)| { + decrypt_note_at(message, ciph_id, &shared_secret).map(|(kind, res_acc)| { let npk = &key_chain.nullifier_public_key; let account_id = lee::AccountId::for_private_account( npk, @@ -920,16 +905,7 @@ 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"), - ) { + if let Some((_kind, new_acc)) = decrypt_note_at(message, ciph_id, &shared_secret) { info!("Synced shared account {account_id:#?} with new state {new_acc:#?}"); index.track(account_id, &new_acc, &nsk); self.storage @@ -956,6 +932,18 @@ impl WalletCore { } } +fn decrypt_note_at( + message: &Message, + i: usize, + secret: &SharedSecretKey, +) -> Option<(lee_core::PrivateAccountKind, Account)> { + lee_core::EncryptionScheme::decrypt( + &message.encrypted_private_post_states[i].ciphertext, + secret, + &message.new_nullifiers[i].0, + ) +} + #[cfg(test)] mod tests { use std::{ffi::CString, str::FromStr as _}; diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs index bb1971d7..54493386 100644 --- a/lez/wallet/src/storage/key_chain.rs +++ b/lez/wallet/src/storage/key_chain.rs @@ -10,8 +10,7 @@ use key_protocol::key_management::{ }; use lee::{Account, AccountId, privacy_preserving_transaction::message::Message}; use lee_core::{ - Commitment, EncryptionScheme, Identifier, Nullifier, NullifierSecretKey, PrivateAccountKind, - SharedSecretKey, + Commitment, Identifier, Nullifier, NullifierSecretKey, PrivateAccountKind, SharedSecretKey, }; use log::{debug, warn}; use serde::{Deserialize, Serialize}; @@ -418,8 +417,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) { @@ -442,8 +439,7 @@ impl UserKeyChain { ) }; - let (kind, new_account) = - EncryptionScheme::decrypt(&encrypted.ciphertext, &secret, commitment, ciph_id)?; + let (kind, new_account) = crate::decrypt_note_at(message, i, &secret)?; let new_nullifier = NullifierIndex::next_update_nullifier(account_id, &new_account, &nsk); if is_shared { @@ -455,6 +451,21 @@ impl UserKeyChain { Some(new_nullifier) } + #[must_use] + pub fn locate_spend(&self, account_id: AccountId, message: &Message) -> Option { + let init = Nullifier::for_account_initialization(&account_id); + let update = self.private_account(account_id).map(|acc| { + Nullifier::for_account_update( + &Commitment::new(&account_id, acc.account), + &acc.key_chain.private_key_holder.nullifier_secret_key, + ) + }); + message + .new_nullifiers + .iter() + .position(|(nullifier, _)| *nullifier == init || Some(nullifier) == update.as_ref()) + } + pub fn add_imported_public_account(&mut self, private_key: lee::PrivateKey) { let account_id = AccountId::from(&lee::PublicKey::new_from_private_key(&private_key)); @@ -865,7 +876,7 @@ impl Default for UserKeyChain { #[cfg(test)] mod tests { - use lee_core::encryption::EncryptedAccountData; + use lee_core::{EncryptionScheme, encryption::EncryptedAccountData}; use super::*; @@ -900,8 +911,7 @@ mod tests { &new_account, &PrivateAccountKind::Regular(identifier), &sender_ss, - &new_commitment, - 0, + &old_nullifier, ); let note = EncryptedAccountData::new( ciphertext, @@ -971,8 +981,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 { @@ -1034,8 +1043,7 @@ mod tests { next, &PrivateAccountKind::Regular(identifier), &sender_ss, - &commitment, - 0, + &spent, ); let note = EncryptedAccountData::new(ciphertext, &npk, &vpk, epk); Message { diff --git a/test_fixtures/fixtures/prebuilt_sequencer_db.dump b/test_fixtures/fixtures/prebuilt_sequencer_db.dump index 26c582a1..771acc7a 100644 Binary files a/test_fixtures/fixtures/prebuilt_sequencer_db.dump and b/test_fixtures/fixtures/prebuilt_sequencer_db.dump differ 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(); }