diff --git a/lee/key_protocol/src/key_management/group_key_holder.rs b/lee/key_protocol/src/key_management/group_key_holder.rs index db6d138d..b18960cf 100644 --- a/lee/key_protocol/src/key_management/group_key_holder.rs +++ b/lee/key_protocol/src/key_management/group_key_holder.rs @@ -149,9 +149,9 @@ impl GroupKeyHolder { /// Derive keys for a shared regular account from its `identifier`. /// /// Computes the derivation seed via the `SharedAccountTag` domain separator, then delegates - /// to [`Self::derive_keys_for_shared_account`] so controllers sharing the GMS agree on keys. + /// to [`Self::derive_keys_for_shared_account`]. #[must_use] - pub fn derive_keys_for_regular_shared_account( + pub fn derive_regular_shared_account_keys_from_identifier( &self, identifier: Identifier, ) -> PrivateKeyHolder { diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index c4250afd..17ffa610 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -402,7 +402,7 @@ impl WalletCore { .group_key_holder(&group_name) .context(format!("Group '{group_name}' not found"))?; - let keys = holder.derive_keys_for_regular_shared_account(identifier); + let keys = holder.derive_regular_shared_account_keys_from_identifier(identifier); let npk = keys.generate_nullifier_public_key(); let vpk = keys.generate_viewing_public_key(); let account_id = AccountId::from((&npk, identifier)); @@ -691,18 +691,20 @@ impl WalletCore { let mut blocks = std::pin::pin!(poller.poll_block_range(last_synced_block.saturating_add(1)..=block_id)); - let mut watch = self.storage.key_chain().build_nullifier_watch(); + // Get the latest nullifiers for all owned accounts. + let mut index = self.storage.key_chain().build_latest_nullifier_index(); let bar = indicatif::ProgressBar::new(num_of_blocks); while let Some(block) = blocks.try_next().await? { for tx in block.body.transactions { + // Eagerly decrypt note updates using expected nullifiers. let handled = if let LeeTransaction::PrivacyPreserving(pp_tx) = &tx { self.storage .key_chain_mut() - .sync_updates_via_nullifiers(&pp_tx.message, &mut watch) + .sync_updates_via_nullifiers(&pp_tx.message, &mut index) } else { HashSet::new() }; - self.sync_private_accounts_with_tx(tx, &mut watch, &handled); + self.sync_private_accounts_with_tx(tx, &mut index, &handled); } self.storage.set_last_synced_block(block.header.block_id); @@ -722,7 +724,7 @@ impl WalletCore { fn sync_private_accounts_with_tx( &mut self, tx: LeeTransaction, - watch: &mut HashMap, + index: &mut HashMap, handled: &HashSet, ) { let LeeTransaction::PrivacyPreserving(tx) = tx else { @@ -745,6 +747,9 @@ impl WalletCore { .iter() .enumerate() .filter(move |(ciph_id, encrypted_data)| { + // If we have not decrypted the update using the nullifiers, + // the note may be an initialized one, for which we should + // scan. !handled.contains(ciph_id) && encrypted_data.view_tag == view_tag }) .filter_map(move |(ciph_id, encrypted_data)| { @@ -779,19 +784,21 @@ impl WalletCore { .key_chain_mut() .insert_private_account(affected_account_id, kind, new_acc) .expect("Account Id should exist"); + // Insert the initial nullifier awaited for the new ID. + // This allows to skip tag matching for noticing updates to it. self.storage .key_chain() - .refresh_watch_entry(watch, affected_account_id); + .refresh_next_nullifier_entry(index, affected_account_id); } // Scan for updates to shared accounts (GMS-derived). - self.sync_shared_private_accounts_with_tx(&tx, watch, handled); + self.sync_shared_private_accounts_with_tx(&tx, index, handled); } fn sync_shared_private_accounts_with_tx( &mut self, tx: &PrivacyPreservingTransaction, - watch: &mut HashMap, + index: &mut HashMap, handled: &HashSet, ) { let shared_keys: Vec<_> = self @@ -816,6 +823,7 @@ impl WalletCore { .iter() .enumerate() { + // If already decrypted or the tag does nto match, skip. if handled.contains(&ciph_id) || encrypted_data.view_tag != view_tag { continue; } @@ -841,7 +849,7 @@ impl WalletCore { .update_shared_private_account_state(&account_id, new_acc); self.storage .key_chain() - .refresh_watch_entry(watch, account_id); + .refresh_next_nullifier_entry(index, account_id); } } } diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs index cbdba58b..74c209a9 100644 --- a/lez/wallet/src/storage/key_chain.rs +++ b/lez/wallet/src/storage/key_chain.rs @@ -297,64 +297,73 @@ impl UserKeyChain { 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, - _ => holder.derive_keys_for_regular_shared_account(entry.identifier), + _ => holder.derive_regular_shared_account_keys_from_identifier(entry.identifier), }) } /// 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_nullifier_watch(&self) -> HashMap { - let mut watch = HashMap::new(); + pub fn build_latest_nullifier_index(&self) -> HashMap { + let mut index = HashMap::new(); + // 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. 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); + index.insert(Nullifier::for_account_update(&commitment, &nsk), account_id); } + // Same for the shared accounts. 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); + index.insert(Nullifier::for_account_update(&commitment, &nsk), account_id); } - watch + index } /// Applies every watched nullifier the `message` publishes: decrypts the position-aligned - /// note, stores the new state, and rolls the watch to the account's next nullifier. Returns + /// note, stores the new state, and rolls the index to the account's next nullifier. Returns /// the output slots handled, so the view-tag pass can skip them. pub fn sync_updates_via_nullifiers( &mut self, message: &Message, - watch: &mut HashMap, + index: &mut HashMap, ) -> HashSet { + // Get the nullifier information if awaiting the nullifier. let hits: Vec<(usize, Nullifier, AccountId)> = message .new_nullifiers .iter() .enumerate() - .filter_map(|(i, (nullifier, _))| watch.get(nullifier).map(|&id| (i, *nullifier, id))) + .filter_map(|(i, (nullifier, _))| index.get(nullifier).map(|&id| (i, *nullifier, id))) .collect(); let mut handled = HashSet::new(); for (i, old_nullifier, account_id) in hits { + // 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) { - watch.remove(&old_nullifier); - watch.insert(new_nullifier, account_id); + 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); + // Record that this nullifier's position can be skipped for scanning. handled.insert(i); } } handled } - /// Decrypts the note at slot `i` for `account_id` (shared or sole-owned), stores the new - /// state, and returns the account's next update nullifier. `None` if keys or decryption fail. + /// Decrypts the note at slot `i` for `account_id`, stores the new state, and returns the + /// account's next update nullifier. `None` if keys or decryption fail. fn apply_nullifier_update( &mut self, account_id: AccountId, @@ -400,11 +409,10 @@ impl UserKeyChain { Some(new_nullifier) } - /// Recomputes `account_id`'s current update nullifier and inserts it into the watch -- used - /// after the view-tag pass stores a freshly discovered or updated account mid-sync. - pub fn refresh_watch_entry( + /// Recomputes `account_id`'s current update nullifier and inserts it into the index. + pub fn refresh_next_nullifier_entry( &self, - watch: &mut HashMap, + index: &mut HashMap, account_id: AccountId, ) { let (account, nsk) = if let Some(entry) = self.shared_private_account(account_id) { @@ -421,7 +429,7 @@ impl UserKeyChain { return; }; let nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, account), &nsk); - watch.insert(nullifier, account_id); + index.insert(nullifier, account_id); } pub fn add_imported_public_account(&mut self, private_key: lee::PrivateKey) { @@ -846,8 +854,8 @@ mod tests { let old_nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, &old_account), &nsk); - let mut watch = kc.build_nullifier_watch(); - assert_eq!(watch.get(&old_nullifier), Some(&account_id)); + let mut index = kc.build_latest_nullifier_index(); + assert_eq!(index.get(&old_nullifier), Some(&account_id)); let new_account = Account { balance: 150, @@ -876,7 +884,7 @@ mod tests { ..Default::default() }; - let handled = kc.sync_updates_via_nullifiers(&message, &mut watch); + let handled = kc.sync_updates_via_nullifiers(&message, &mut index); assert_eq!(handled, HashSet::from([0])); assert_eq!( @@ -885,8 +893,8 @@ mod tests { ); let new_nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk); - assert_eq!(watch.get(&new_nullifier), Some(&account_id)); - assert!(!watch.contains_key(&old_nullifier)); + assert_eq!(index.get(&new_nullifier), Some(&account_id)); + assert!(!index.contains_key(&old_nullifier)); } #[test] @@ -896,7 +904,7 @@ mod tests { let label = Label::new("group"); let holder = GroupKeyHolder::new(); let identifier = 0; - let keys = holder.derive_keys_for_regular_shared_account(identifier); + let keys = holder.derive_regular_shared_account_keys_from_identifier(identifier); let npk = keys.generate_nullifier_public_key(); let vpk = keys.generate_viewing_public_key(); let nsk = keys.nullifier_secret_key; @@ -917,8 +925,8 @@ mod tests { let old_nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, &old_account), &nsk); - let mut watch = kc.build_nullifier_watch(); - assert_eq!(watch.get(&old_nullifier), Some(&account_id)); + let mut index = kc.build_latest_nullifier_index(); + assert_eq!(index.get(&old_nullifier), Some(&account_id)); let new_account = Account { balance: 250, @@ -941,7 +949,7 @@ mod tests { ..Default::default() }; - let handled = kc.sync_updates_via_nullifiers(&message, &mut watch); + let handled = kc.sync_updates_via_nullifiers(&message, &mut index); assert_eq!(handled, HashSet::from([0])); assert_eq!( @@ -950,12 +958,12 @@ mod tests { ); let new_nullifier = Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk); - assert_eq!(watch.get(&new_nullifier), Some(&account_id)); - assert!(!watch.contains_key(&old_nullifier)); + assert_eq!(index.get(&new_nullifier), Some(&account_id)); + assert!(!index.contains_key(&old_nullifier)); } #[test] - fn nullifier_sync_ignores_unwatched_nullifier() { + fn nullifier_sync_ignores_unindexed_nullifier() { let mut kc = UserKeyChain::default(); let key_chain = KeyChain::new_os_random(); @@ -967,17 +975,17 @@ mod tests { let account = Account::default(); kc.add_imported_private_account(key_chain, None, identifier, account.clone()); - let mut watch = kc.build_nullifier_watch(); - let unwatched = Nullifier::for_account_update( + let mut index = kc.build_latest_nullifier_index(); + let unindexed = Nullifier::for_account_update( &Commitment::new(&AccountId::new([9; 32]), &Account::default()), &[9; 32], ); let message = Message { - new_nullifiers: vec![(unwatched, [0; 32])], + new_nullifiers: vec![(unindexed, [0; 32])], ..Default::default() }; - let handled = kc.sync_updates_via_nullifiers(&message, &mut watch); + let handled = kc.sync_updates_via_nullifiers(&message, &mut index); assert!(handled.is_empty()); assert_eq!(kc.private_account(account_id).unwrap().account, &account);