From d2633af413fb2ad18c11996a59a36fa8f7044c5d Mon Sep 17 00:00:00 2001 From: Artem Gureev Date: Wed, 1 Jul 2026 18:23:19 +0000 Subject: [PATCH] refactor(wallet): introduce a NullifierIndex struct --- lez/wallet/src/lib.rs | 33 +++++-------- lez/wallet/src/storage/key_chain.rs | 72 +++++++++++++++++++++-------- 2 files changed, 64 insertions(+), 41 deletions(-) diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index 412ca34c..3fe103c3 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -7,10 +7,7 @@ reason = "Most of the shadows come from args parsing which is ok" )] -use std::{ - collections::{HashMap, HashSet}, - path::PathBuf, -}; +use std::{collections::HashSet, path::PathBuf}; pub use account_manager::AccountIdentity; use anyhow::{Context as _, Result}; @@ -25,7 +22,7 @@ use lee::{ }, }; use lee_core::{ - Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, Nullifier, SharedSecretKey, + Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, SharedSecretKey, account::Nonce, compute_digest_for_path, program::InstructionData, }; use log::info; @@ -37,7 +34,7 @@ use crate::{ account::{AccountIdWithPrivacy, Label}, config::WalletConfigOverrides, poller::TxPoller, - storage::key_chain::SharedAccountEntry, + storage::key_chain::{NullifierIndex, SharedAccountEntry}, }; pub mod account; @@ -758,7 +755,7 @@ impl WalletCore { fn sync_private_accounts_with_tx( &mut self, tx: LeeTransaction, - index: &mut HashMap, + index: &mut NullifierIndex, handled: &HashSet, ) { let LeeTransaction::PrivacyPreserving(tx) = tx else { @@ -804,28 +801,24 @@ impl WalletCore { let npk = &key_chain.nullifier_public_key; let account_id = lee::AccountId::for_private_account(npk, &kind); let nsk = key_chain.private_key_holder.nullifier_secret_key; - let nullifier = Nullifier::for_account_update( - &Commitment::new(&account_id, &res_acc), - &nsk, - ); - (account_id, kind, res_acc, nullifier) + (account_id, kind, res_acc, nsk) }) }) .collect::>() }) .collect::>(); - for (affected_account_id, kind, new_acc, nullifier) in affected_accounts { + for (affected_account_id, kind, new_acc, nsk) in affected_accounts { info!( "Received new account for account_id {affected_account_id:#?} with account object {new_acc:#?}" ); + // Await the account's next update by its nullifier, so later updates + // to it are caught without tag matching. + index.track(affected_account_id, &new_acc, &nsk); self.storage .key_chain_mut() .insert_private_account(affected_account_id, kind, new_acc) .expect("Account Id should exist"); - // Await the account's next update by its nullifier, so later updates - // to it are caught without tag matching. - index.insert(nullifier, affected_account_id); } // Scan for updates to shared accounts (GMS-derived). @@ -835,7 +828,7 @@ impl WalletCore { fn sync_shared_private_accounts_with_tx( &mut self, tx: &PrivacyPreservingTransaction, - index: &mut HashMap, + index: &mut NullifierIndex, handled: &HashSet, ) { let shared_keys: Vec<_> = self @@ -882,14 +875,10 @@ impl WalletCore { .expect("Ciphertext ID is expected to fit in u32"), ) { info!("Synced shared account {account_id:#?} with new state {new_acc:#?}"); - let nullifier = Nullifier::for_account_update( - &Commitment::new(&account_id, &new_acc), - &nsk, - ); + index.track(account_id, &new_acc, &nsk); self.storage .key_chain_mut() .update_shared_private_account_state(&account_id, new_acc); - index.insert(nullifier, account_id); } } } diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs index 50c6e68d..948b9d1d 100644 --- a/lez/wallet/src/storage/key_chain.rs +++ b/lez/wallet/src/storage/key_chain.rs @@ -10,7 +10,8 @@ use key_protocol::key_management::{ }; use lee::{Account, AccountId, privacy_preserving_transaction::message::Message}; use lee_core::{ - Commitment, EncryptionScheme, Identifier, Nullifier, PrivateAccountKind, SharedSecretKey, + Commitment, EncryptionScheme, Identifier, Nullifier, NullifierSecretKey, PrivateAccountKind, + SharedSecretKey, }; use log::{debug, warn}; use serde::{Deserialize, Serialize}; @@ -61,6 +62,41 @@ pub struct SharedAccountEntry { pub account: Account, } +/// Maps each owned or shared private account to the nullifier its next update will publish, +/// so sync can spot updates by nullifier rather than view tag. +#[derive(Default)] +pub struct NullifierIndex(HashMap); + +impl NullifierIndex { + fn next_update_nullifier( + account_id: AccountId, + account: &Account, + nsk: &NullifierSecretKey, + ) -> Nullifier { + Nullifier::for_account_update(&Commitment::new(&account_id, account), nsk) + } + + /// Returns the account whose next update would publish `nullifier`. + #[must_use] + pub fn account_for(&self, nullifier: &Nullifier) -> Option { + self.0.get(nullifier).copied() + } + + /// Indexes `account_id` by the nullifier its next update will publish. + pub fn track(&mut self, account_id: AccountId, account: &Account, nsk: &NullifierSecretKey) { + self.0.insert( + Self::next_update_nullifier(account_id, account, nsk), + account_id, + ); + } + + /// Replaces a spent nullifier with the account's `next` one. + pub fn update(&mut self, spent: &Nullifier, next: Nullifier, account_id: AccountId) { + self.0.remove(spent); + self.0.insert(next, account_id); + } +} + #[derive(Clone, Debug)] #[cfg_attr(test, derive(PartialEq, Eq))] pub struct UserKeyChain { @@ -304,8 +340,8 @@ impl UserKeyChain { /// Maps each owned and shared account's current-state update nullifier to its `account_id`, /// so co-owner updates are found during sync by nullifier rather than view tag. #[must_use] - pub fn build_latest_nullifier_index(&self) -> HashMap { - let mut index = HashMap::new(); + pub fn build_latest_nullifier_index(&self) -> NullifierIndex { + let mut index = NullifierIndex::default(); // For each (regular) found account the user owns, compute its nullifier and put // into the map. This is the next nullifier it will look for. @@ -313,8 +349,7 @@ impl UserKeyChain { let account_id = AccountId::for_private_account(&found.key_chain.nullifier_public_key, found.kind); let nsk = found.key_chain.private_key_holder.nullifier_secret_key; - let commitment = Commitment::new(&account_id, found.account); - index.insert(Nullifier::for_account_update(&commitment, &nsk), account_id); + index.track(account_id, found.account, &nsk); } // Same for the shared accounts. @@ -323,8 +358,7 @@ impl UserKeyChain { continue; }; let nsk = keys.nullifier_secret_key; - let commitment = Commitment::new(&account_id, &entry.account); - index.insert(Nullifier::for_account_update(&commitment, &nsk), account_id); + index.track(account_id, &entry.account, &nsk); } index @@ -336,14 +370,16 @@ impl UserKeyChain { pub fn sync_updates_via_nullifiers( &mut self, message: &Message, - index: &mut HashMap, + index: &mut NullifierIndex, ) -> HashSet { // Get the nullifier information if awaiting the nullifier. let hits: Vec<(usize, Nullifier, AccountId)> = message .new_nullifiers .iter() .enumerate() - .filter_map(|(i, (nullifier, _))| index.get(nullifier).map(|&id| (i, *nullifier, id))) + .filter_map(|(i, (nullifier, _))| { + index.account_for(nullifier).map(|id| (i, *nullifier, id)) + }) .collect(); let mut handled = HashSet::new(); @@ -351,10 +387,9 @@ impl UserKeyChain { // Try decrypting the commitment connectted to the nullifier and get the next // nullifier to await. if let Some(new_nullifier) = self.apply_nullifier_update(account_id, message, i) { - index.remove(&old_nullifier); // Update the index to await for the new state of the account, i.e. // the new nullifier. - index.insert(new_nullifier, account_id); + index.update(&old_nullifier, new_nullifier, account_id); // Record that this nullifier's position can be skipped for scanning. handled.insert(i); } @@ -397,8 +432,7 @@ impl UserKeyChain { let (kind, new_account) = EncryptionScheme::decrypt(&encrypted.ciphertext, &secret, commitment, ciph_id)?; - let new_nullifier = - Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk); + let new_nullifier = NullifierIndex::next_update_nullifier(account_id, &new_account, &nsk); if is_shared { self.update_shared_private_account_state(&account_id, new_account); @@ -832,7 +866,7 @@ mod tests { let old_nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, &old_account), &nsk); let mut index = kc.build_latest_nullifier_index(); - assert_eq!(index.get(&old_nullifier), Some(&account_id)); + assert_eq!(index.account_for(&old_nullifier), Some(account_id)); let new_account = Account { balance: 150, @@ -870,8 +904,8 @@ mod tests { ); let new_nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk); - assert_eq!(index.get(&new_nullifier), Some(&account_id)); - assert!(!index.contains_key(&old_nullifier)); + assert_eq!(index.account_for(&new_nullifier), Some(account_id)); + assert!(index.account_for(&old_nullifier).is_none()); } #[test] @@ -903,7 +937,7 @@ mod tests { let old_nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, &old_account), &nsk); let mut index = kc.build_latest_nullifier_index(); - assert_eq!(index.get(&old_nullifier), Some(&account_id)); + assert_eq!(index.account_for(&old_nullifier), Some(account_id)); let new_account = Account { balance: 250, @@ -935,8 +969,8 @@ mod tests { ); let new_nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk); - assert_eq!(index.get(&new_nullifier), Some(&account_id)); - assert!(!index.contains_key(&old_nullifier)); + assert_eq!(index.account_for(&new_nullifier), Some(account_id)); + assert!(index.account_for(&old_nullifier).is_none()); } #[test]