Merge 2e669e76f83fb27e9d0ebbb62b0d9a8799aea1fd into abb18568dc06bba7f2fc26fda81134f917dd9035

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);
 

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);

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"] }

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);
+        }
+    }
+}

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<AccountWithMetadata>,
     pub public_post_states: Vec<Account>,

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 {

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]

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)]

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);

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());

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 _};

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<Nullifier> {
         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<usize> {
+        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 {

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();
 }