From 6bb2d7ffc35f424cf30ade467a5cd130399b938a Mon Sep 17 00:00:00 2001 From: Artem Gureev Date: Wed, 1 Jul 2026 13:34:50 +0000 Subject: [PATCH] refactor(wallet): move nullifier-watch sync onto UserKeyChain --- lez/wallet/src/lib.rs | 139 ++-------------------------- lez/wallet/src/storage/key_chain.rs | 122 +++++++++++++++++++++++- 2 files changed, 129 insertions(+), 132 deletions(-) diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index e635f062..acea846d 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -699,12 +699,14 @@ impl WalletCore { let mut blocks = std::pin::pin!(poller.poll_block_range(last_synced_block.saturating_add(1)..=block_id)); - let mut watch = self.build_nullifier_watch(); + let mut watch = self.storage.key_chain().build_nullifier_watch(); let bar = indicatif::ProgressBar::new(num_of_blocks); while let Some(block) = blocks.try_next().await? { for tx in block.body.transactions { let handled = if let LeeTransaction::PrivacyPreserving(pp_tx) = &tx { - self.sync_updates_via_nullifiers(pp_tx, &mut watch) + self.storage + .key_chain_mut() + .sync_updates_via_nullifiers(&pp_tx.message, &mut watch) } else { HashSet::new() }; @@ -785,7 +787,9 @@ impl WalletCore { .key_chain_mut() .insert_private_account(affected_account_id, kind, new_acc) .expect("Account Id should exist"); - self.refresh_watch_entry(watch, affected_account_id); + self.storage + .key_chain() + .refresh_watch_entry(watch, affected_account_id); } // Scan for updates to shared accounts (GMS-derived). @@ -843,137 +847,14 @@ impl WalletCore { self.storage .key_chain_mut() .update_shared_private_account_state(&account_id, new_acc); - self.refresh_watch_entry(watch, account_id); + self.storage + .key_chain() + .refresh_watch_entry(watch, account_id); } } } } - fn build_nullifier_watch(&self) -> HashMap { - let key_chain = self.storage.key_chain(); - let mut watch = HashMap::new(); - - for found in key_chain.private_accounts() { - 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); - watch.insert(Nullifier::for_account_update(&commitment, &nsk), account_id); - } - - for (&account_id, entry) in key_chain.shared_private_accounts_iter() { - let Some(keys) = key_chain.derive_shared_account_keys(entry) else { - continue; - }; - let nsk = keys.nullifier_secret_key; - let commitment = Commitment::new(&account_id, &entry.account); - watch.insert(Nullifier::for_account_update(&commitment, &nsk), account_id); - } - - watch - } - - fn sync_updates_via_nullifiers( - &mut self, - tx: &PrivacyPreservingTransaction, - watch: &mut HashMap, - ) -> HashSet { - let hits: Vec<(usize, Nullifier, AccountId)> = tx - .message - .new_nullifiers - .iter() - .enumerate() - .filter_map(|(i, (nullifier, _))| watch.get(nullifier).map(|&id| (i, *nullifier, id))) - .collect(); - - let mut handled = HashSet::new(); - for (i, old_nullifier, account_id) in hits { - if let Some(new_nullifier) = self.apply_nullifier_update(account_id, tx, i) { - watch.remove(&old_nullifier); - watch.insert(new_nullifier, account_id); - handled.insert(i); - } - } - handled - } - - fn apply_nullifier_update( - &mut self, - account_id: AccountId, - tx: &PrivacyPreservingTransaction, - i: usize, - ) -> Option { - let encrypted = &tx.message.encrypted_private_post_states[i]; - let commitment = &tx.message.new_commitments[i]; - let ciph_id = u32::try_from(i).ok()?; - - let key_chain = self.storage.key_chain(); - let (nsk, secret, is_shared) = - if let Some(entry) = key_chain.shared_private_account(account_id) { - let keys = key_chain.derive_shared_account_keys(entry)?; - let secret = SharedSecretKey::decapsulate( - &encrypted.epk, - &keys.viewing_secret_key.d, - &keys.viewing_secret_key.z, - )?; - (keys.nullifier_secret_key, secret, true) - } else { - let found = key_chain.private_account(account_id)?; - let secret = found - .key_chain - .calculate_shared_secret_receiver(&encrypted.epk)?; - ( - found.key_chain.private_key_holder.nullifier_secret_key, - secret, - false, - ) - }; - - let (kind, new_account) = lee_core::EncryptionScheme::decrypt( - &encrypted.ciphertext, - &secret, - commitment, - ciph_id, - )?; - let new_nullifier = - Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk); - - if is_shared { - self.storage - .key_chain_mut() - .update_shared_private_account_state(&account_id, new_account); - } else { - self.storage - .key_chain_mut() - .insert_private_account(account_id, kind, new_account) - .ok()?; - } - Some(new_nullifier) - } - - fn refresh_watch_entry( - &self, - watch: &mut HashMap, - account_id: AccountId, - ) { - let key_chain = self.storage.key_chain(); - let (account, nsk) = if let Some(entry) = key_chain.shared_private_account(account_id) { - let Some(keys) = key_chain.derive_shared_account_keys(entry) else { - return; - }; - (&entry.account, keys.nullifier_secret_key) - } else if let Some(found) = key_chain.private_account(account_id) { - ( - found.account, - found.key_chain.private_key_holder.nullifier_secret_key, - ) - } else { - return; - }; - let nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, account), &nsk); - watch.insert(nullifier, account_id); - } - #[must_use] pub const fn config_path(&self) -> &PathBuf { &self.config_path diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs index bf4d8ab2..2694f574 100644 --- a/lez/wallet/src/storage/key_chain.rs +++ b/lez/wallet/src/storage/key_chain.rs @@ -1,5 +1,5 @@ use core::panic; -use std::collections::{BTreeMap, btree_map::Entry}; +use std::collections::{BTreeMap, HashMap, HashSet, btree_map::Entry}; use anyhow::{Context as _, Result, anyhow}; use key_protocol::key_management::{ @@ -8,8 +8,10 @@ use key_protocol::key_management::{ key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex, traits::KeyTreeNode as _}, secret_holders::{PrivateKeyHolder, SeedHolder, ViewingSecretKey}, }; -use lee::{Account, AccountId}; -use lee_core::{Identifier, PrivateAccountKind}; +use lee::{Account, AccountId, privacy_preserving_transaction::message::Message}; +use lee_core::{ + Commitment, EncryptionScheme, Identifier, Nullifier, PrivateAccountKind, SharedSecretKey, +}; use log::{debug, warn}; use serde::{Deserialize, Serialize}; use testnet_initial_state::{PrivateAccountPrivateInitialData, PublicAccountPrivateInitialData}; @@ -305,6 +307,120 @@ impl UserKeyChain { }) } + #[must_use] + pub fn build_nullifier_watch(&self) -> HashMap { + let mut watch = HashMap::new(); + + for found in self.private_accounts() { + 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); + watch.insert(Nullifier::for_account_update(&commitment, &nsk), account_id); + } + + for (&account_id, entry) in self.shared_private_accounts_iter() { + let Some(keys) = self.derive_shared_account_keys(entry) else { + continue; + }; + let nsk = keys.nullifier_secret_key; + let commitment = Commitment::new(&account_id, &entry.account); + watch.insert(Nullifier::for_account_update(&commitment, &nsk), account_id); + } + + watch + } + + pub fn sync_updates_via_nullifiers( + &mut self, + message: &Message, + watch: &mut HashMap, + ) -> HashSet { + let hits: Vec<(usize, Nullifier, AccountId)> = message + .new_nullifiers + .iter() + .enumerate() + .filter_map(|(i, (nullifier, _))| watch.get(nullifier).map(|&id| (i, *nullifier, id))) + .collect(); + + let mut handled = HashSet::new(); + for (i, old_nullifier, account_id) in hits { + if let Some(new_nullifier) = self.apply_nullifier_update(account_id, message, i) { + watch.remove(&old_nullifier); + watch.insert(new_nullifier, account_id); + handled.insert(i); + } + } + handled + } + + fn apply_nullifier_update( + &mut self, + account_id: AccountId, + message: &Message, + 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) + { + let keys = self.derive_shared_account_keys(entry)?; + let secret = SharedSecretKey::decapsulate( + &encrypted.epk, + &keys.viewing_secret_key.d, + &keys.viewing_secret_key.z, + )?; + (keys.nullifier_secret_key, secret, true) + } else { + let found = self.private_account(account_id)?; + let secret = found + .key_chain + .calculate_shared_secret_receiver(&encrypted.epk)?; + ( + found.key_chain.private_key_holder.nullifier_secret_key, + secret, + false, + ) + }; + + 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); + + if is_shared { + self.update_shared_private_account_state(&account_id, new_account); + } else { + self.insert_private_account(account_id, kind, new_account) + .ok()?; + } + Some(new_nullifier) + } + + pub fn refresh_watch_entry( + &self, + watch: &mut HashMap, + account_id: AccountId, + ) { + let (account, nsk) = if let Some(entry) = self.shared_private_account(account_id) { + let Some(keys) = self.derive_shared_account_keys(entry) else { + return; + }; + (&entry.account, keys.nullifier_secret_key) + } else if let Some(found) = self.private_account(account_id) { + ( + found.account, + found.key_chain.private_key_holder.nullifier_secret_key, + ) + } else { + return; + }; + let nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, account), &nsk); + watch.insert(nullifier, account_id); + } + 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));