diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index 80b42e17..60dcc228 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -7,7 +7,7 @@ reason = "Most of the shadows come from args parsing which is ok" )] -use std::path::PathBuf; +use std::{collections::HashMap, path::PathBuf}; pub use account_manager::AccountIdentity; use anyhow::{Context as _, Result}; @@ -22,7 +22,8 @@ use lee::{ }, }; use lee_core::{ - Commitment, MembershipProof, SharedSecretKey, account::Nonce, program::InstructionData, + Commitment, MembershipProof, Nullifier, SharedSecretKey, account::Nonce, + program::InstructionData, }; use log::info; use sequencer_service_rpc::{RpcClient as _, SequencerClient, SequencerClientBuilder}; @@ -789,28 +790,7 @@ impl WalletCore { .key_chain() .shared_private_accounts_iter() .filter_map(|(&account_id, entry)| { - let holder = self - .storage - .key_chain() - .group_key_holder(&entry.group_label)?; - - let keys = match (&entry.pda_seed, &entry.authority_program_id) { - (Some(pda_seed), Some(program_id)) => { - holder.derive_keys_for_pda(program_id, pda_seed) - } - (Some(_), None) => return None, // PDA without program_id, skip - _ => { - let derivation_seed = { - use sha2::Digest as _; - let mut hasher = sha2::Sha256::new(); - hasher.update(b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00"); - hasher.update(entry.identifier.to_le_bytes()); - let result: [u8; 32] = hasher.finalize().into(); - result - }; - holder.derive_keys_for_shared_account(&derivation_seed) - } - }; + let keys = self.storage.key_chain().derive_shared_account_keys(entry)?; let npk = keys.generate_nullifier_public_key(); let vpk = keys.generate_viewing_public_key(); let vsk = keys.viewing_secret_key; @@ -855,6 +835,31 @@ impl WalletCore { } } + #[expect(dead_code, reason = "wired into the sync loop in the update-scan unit")] + 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 + } + #[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 f73833eb..bf4d8ab2 100644 --- a/lez/wallet/src/storage/key_chain.rs +++ b/lez/wallet/src/storage/key_chain.rs @@ -6,7 +6,7 @@ use key_protocol::key_management::{ KeyChain, group_key_holder::GroupKeyHolder, key_tree::{KeyTreePrivate, KeyTreePublic, chain_index::ChainIndex, traits::KeyTreeNode as _}, - secret_holders::{SeedHolder, ViewingSecretKey}, + secret_holders::{PrivateKeyHolder, SeedHolder, ViewingSecretKey}, }; use lee::{Account, AccountId}; use lee_core::{Identifier, PrivateAccountKind}; @@ -209,41 +209,11 @@ impl UserKeyChain { /// those. #[must_use] pub fn private_account(&self, account_id: AccountId) -> Option> { - self.imported_private_accounts - .iter() - .flat_map(|(key, data)| { - data.accounts - .iter() - .map(|(kind, account)| FoundPrivateAccount { - account, - key_chain: &key.key_chain, - kind, - chain_index: key.chain_index.clone(), - }) - }) - .chain( - self.private_key_tree - .key_map - .iter() - .flat_map(|(chain_index, data)| { - data.value - .1 - .iter() - .map(|(kind, account)| FoundPrivateAccount { - account, - key_chain: &data.value.0, - kind, - chain_index: Some(chain_index.clone()), - }) - }), - ) - .find_map(|found| { - let expected_id = AccountId::for_private_account( - &found.key_chain.nullifier_public_key, - found.kind, - ); - (expected_id == account_id).then_some(found) - }) + self.private_accounts().find_map(|found| { + let expected_id = + AccountId::for_private_account(&found.key_chain.nullifier_public_key, found.kind); + (expected_id == account_id).then_some(found) + }) } #[must_use] @@ -281,6 +251,60 @@ impl UserKeyChain { ) } + pub fn private_accounts(&self) -> impl Iterator> { + self.imported_private_accounts + .iter() + .flat_map(|(key, data)| { + data.accounts + .iter() + .map(|(kind, account)| FoundPrivateAccount { + account, + key_chain: &key.key_chain, + kind, + chain_index: key.chain_index.clone(), + }) + }) + .chain( + self.private_key_tree + .key_map + .iter() + .flat_map(|(chain_index, data)| { + data.value + .1 + .iter() + .map(|(kind, account)| FoundPrivateAccount { + account, + key_chain: &data.value.0, + kind, + chain_index: Some(chain_index.clone()), + }) + }), + ) + } + + #[must_use] + pub fn derive_shared_account_keys( + &self, + entry: &SharedAccountEntry, + ) -> Option { + let holder = self.group_key_holder(&entry.group_label)?; + Some(match (&entry.pda_seed, &entry.authority_program_id) { + (Some(pda_seed), Some(program_id)) => holder.derive_keys_for_pda(program_id, pda_seed), + (Some(_), None) => return None, + _ => { + let derivation_seed = { + use sha2::Digest as _; + let mut hasher = sha2::Sha256::new(); + hasher.update(b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00"); + hasher.update(entry.identifier.to_le_bytes()); + let result: [u8; 32] = hasher.finalize().into(); + result + }; + holder.derive_keys_for_shared_account(&derivation_seed) + } + }) + } + 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));