feat(wallet): nullifier-watch for updated private accounts

This commit is contained in:
Artem Gureev 2026-07-01 11:37:52 +00:00 committed by agureev
parent a58fbce2ff
commit eea069c663
2 changed files with 89 additions and 60 deletions

View File

@ -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<Nullifier, AccountId> {
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

View File

@ -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<FoundPrivateAccount<'_>> {
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<Item = FoundPrivateAccount<'_>> {
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<PrivateKeyHolder> {
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));