From ebcb3dde6b2ea5d330eb1821f04b61770063a6c1 Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Thu, 22 Feb 2024 18:26:30 -0700 Subject: [PATCH] zcash_client_backend: Factor out common note decryption from `scan_block_with_runner` --- zcash_client_backend/src/scan.rs | 23 ++- zcash_client_backend/src/scanning.rs | 278 +++++++++++++++++---------- 2 files changed, 197 insertions(+), 104 deletions(-) diff --git a/zcash_client_backend/src/scan.rs b/zcash_client_backend/src/scan.rs index 98a768622e..6b792f29f1 100644 --- a/zcash_client_backend/src/scan.rs +++ b/zcash_client_backend/src/scan.rs @@ -506,6 +506,14 @@ where } } +pub(crate) trait BatchResults { + fn collect_results( + &mut self, + block_tag: BlockHash, + txid: TxId, + ) -> HashMap<(TxId, usize), DecryptedOutput>; +} + impl BatchRunner where A: Clone + Send + 'static, @@ -554,13 +562,26 @@ where self.running_tasks.run_task(batch); } } +} +impl BatchResults for BatchRunner +where + A: Clone + Send + 'static, + D: BatchDomain + Send + 'static, + //D::IncomingViewingKey: Clone + Send, + D::Memo: Send, + D::Note: Send, + D::Recipient: Send, + Output: Clone + Send + 'static, + Dec: Decryptor, + T: Tasks>, +{ /// Collects the pending decryption results for the given transaction. /// /// `block_tag` is the hash of the block that triggered this txid being added to the /// batch, or the all-zeros hash to indicate that no block triggered it (i.e. it was a /// mempool change). - pub(crate) fn collect_results( + fn collect_results( &mut self, block_tag: BlockHash, txid: TxId, diff --git a/zcash_client_backend/src/scanning.rs b/zcash_client_backend/src/scanning.rs index d0cf5bdb0b..329254285c 100644 --- a/zcash_client_backend/src/scanning.rs +++ b/zcash_client_backend/src/scanning.rs @@ -11,11 +11,14 @@ use sapling::{ SaplingIvk, }; use subtle::{ConditionallySelectable, ConstantTimeEq, CtOption}; -use zcash_note_encryption::batch; +use zcash_note_encryption::{batch, BatchDomain, ShieldedOutput}; +use zcash_primitives::block::BlockHash; use zcash_primitives::consensus::{self, BlockHeight, NetworkUpgrade}; +use zcash_primitives::transaction::TxId; use zip32::{AccountId, Scope}; use crate::data_api::{BlockMetadata, ScannedBlock, ScannedBundles}; +use crate::scan::BatchResults; use crate::{ proto::compact_formats::CompactBlock, scan::{Batch, BatchRunner, CompactDecryptor, Tasks}, @@ -513,11 +516,16 @@ pub(crate) fn scan_block_with_runner< let tx_actions_len = u32::try_from(tx.actions.len()).expect("Orchard action count cannot exceed a u32"); - // Check for incoming notes while incrementing tree and witnesses - let mut shielded_outputs: Vec> = vec![]; - { - let decoded = &tx - .outputs + let (sapling_outputs, mut sapling_nc) = find_notes( + cur_height, + cur_hash, + compact_block_tx_count, + txid, + tx_idx, + sapling_commitment_tree_size, + sapling_keys, + &spent_from_accounts, + &tx.outputs .into_iter() .map(|output| { ( @@ -526,108 +534,32 @@ pub(crate) fn scan_block_with_runner< .expect("Invalid output found in compact block decoding."), ) }) - .collect::>(); - - let decrypted: Vec<_> = if let Some(runner) = sapling_batch_runner.as_mut() { - let sapling_keys = sapling_keys - .iter() - .flat_map(|(a, k)| { - k.to_ivks() - .into_iter() - .map(move |(scope, _, nk)| ((**a, scope), nk)) - }) - .collect::>(); - - let mut decrypted = runner.collect_results(cur_hash, txid); - (0..decoded.len()) - .map(|i| { - decrypted.remove(&(txid, i)).map(|d_note| { - let a = d_note.ivk_tag.0; - let nk = sapling_keys.get(&d_note.ivk_tag).expect( - "The batch runner and scan_block must use the same set of IVKs.", - ); - - (d_note.note, a, d_note.ivk_tag.1, (*nk).clone()) - }) - }) - .collect() - } else { - let sapling_keys = sapling_keys - .iter() - .flat_map(|(a, k)| { - k.to_ivks() - .into_iter() - .map(move |(scope, ivk, nk)| (**a, scope, ivk, nk)) - }) - .collect::>(); - - let ivks = sapling_keys - .iter() - .map(|(_, _, ivk, _)| PreparedIncomingViewingKey::new(ivk)) - .collect::>(); - - batch::try_compact_note_decryption(&ivks, &decoded[..]) - .into_iter() - .map(|v| { - v.map(|((note, _), ivk_idx)| { - let (account, scope, _, nk) = &sapling_keys[ivk_idx]; - (note, *account, scope.clone(), (*nk).clone()) - }) - }) - .collect() - }; - - for (output_idx, ((_, output), dec_output)) in decoded.iter().zip(decrypted).enumerate() - { - // Collect block note commitments - let node = sapling::Node::from_cmu(&output.cmu); - let is_checkpoint = - output_idx + 1 == decoded.len() && tx_idx + 1 == compact_block_tx_count; - let retention = match (dec_output.is_some(), is_checkpoint) { - (is_marked, true) => Retention::Checkpoint { - id: cur_height, - is_marked, - }, - (true, false) => Retention::Marked, - (false, false) => Retention::Ephemeral, - }; - - if let Some((note, account, scope, nk)) = dec_output { - // A note is marked as "change" if the account that received it - // also spent notes in the same transaction. This will catch, - // for instance: - // - Change created by spending fractions of notes. - // - Notes created by consolidation transactions. - // - Notes sent from one account to itself. - let is_change = spent_from_accounts.contains(&account); - let note_commitment_tree_position = Position::from(u64::from( - sapling_commitment_tree_size + u32::try_from(output_idx).unwrap(), - )); - let nf = SK::nf(&nk, ¬e, note_commitment_tree_position); - - shielded_outputs.push(WalletSaplingOutput::from_parts( - output_idx, - output.cmu, - output.ephemeral_key.clone(), - account, - note, - is_change, - note_commitment_tree_position, - nf, - scope, - )); - } - - sapling_note_commitments.push((node, retention)); - } - } + .collect::>(), + &mut sapling_batch_runner, + PreparedIncomingViewingKey::new, + |output| sapling::Node::from_cmu(&output.cmu), + |output_idx, output, account, note, is_change, position, nf, scope| { + WalletSaplingOutput::from_parts( + output_idx, + output.cmu, + output.ephemeral_key.clone(), + account, + note, + is_change, + position, + nf, + scope, + ) + }, + ); + sapling_note_commitments.append(&mut sapling_nc); - if !(sapling_spends.is_empty() && shielded_outputs.is_empty()) { + if !(sapling_spends.is_empty() && sapling_outputs.is_empty()) { wtxs.push(WalletTx { txid, index: tx_index as usize, sapling_spends, - sapling_outputs: shielded_outputs, + sapling_outputs, }); } @@ -714,6 +646,146 @@ fn check_nullifiers( (found_spent, unlinked_nullifiers) } +fn find_notes< + M, + D: BatchDomain, + SK: ScanningKey, + Output: ShieldedOutput, + B: BatchResults<(AccountId, SK::Scope), D, M>, + WalletOutput, + NoteCommitment, +>( + block_height: BlockHeight, + block_tag: BlockHash, + block_tx_count: usize, + txid: TxId, + tx_idx: usize, + commitment_tree_size: u32, + keys: &[(&AccountId, SK)], + spent_from_accounts: &HashSet, + decoded: &[(D, Output)], + batch_runner: &mut Option<&mut B>, + prepare_key: impl Fn(&SK::IncomingViewingKey) -> D::IncomingViewingKey, + extract_note_commitment: impl Fn(&Output) -> NoteCommitment, + new_wallet_output: impl Fn( + usize, + &Output, + AccountId, + SK::Note, + bool, + Position, + SK::Nf, + SK::Scope, + ) -> WalletOutput, +) -> ( + Vec, + Vec<(NoteCommitment, Retention)>, +) { + // Check for incoming notes while incrementing tree and witnesses + let (decrypted_opts, decrypted_len) = if let Some(runner) = batch_runner.as_mut() { + let tagged_keys = keys + .iter() + .flat_map(|(a, k)| { + k.to_ivks() + .into_iter() + .map(move |(scope, _, nk)| ((**a, scope), nk)) + }) + .collect::>(); + + let mut decrypted = runner.collect_results(block_tag, txid); + let decrypted_len = decrypted.len(); + ( + (0..decoded.len()) + .map(|i| { + decrypted.remove(&(txid, i)).map(|d_note| { + let a = d_note.ivk_tag.0; + let nk = tagged_keys.get(&d_note.ivk_tag).expect( + "The batch runner and scan_block must use the same set of IVKs.", + ); + + (d_note.note, a, d_note.ivk_tag.1, (*nk).clone()) + }) + }) + .collect::>(), + decrypted_len, + ) + } else { + let tagged_keys = keys + .iter() + .flat_map(|(a, k)| { + k.to_ivks() + .into_iter() + .map(move |(scope, ivk, nk)| (**a, scope, ivk, nk)) + }) + .collect::>(); + + let ivks = tagged_keys + .iter() + .map(|(_, _, ivk, _)| prepare_key(ivk)) + .collect::>(); + + let mut decrypted_len = 0; + ( + batch::try_compact_note_decryption(&ivks, &decoded[..]) + .into_iter() + .map(|v| { + v.map(|((note, _), ivk_idx)| { + decrypted_len += 1; + let (account, scope, _, nk) = &tagged_keys[ivk_idx]; + (note, *account, scope.clone(), (*nk).clone()) + }) + }) + .collect::>(), + decrypted_len, + ) + }; + + let mut shielded_outputs = Vec::with_capacity(decrypted_len); + let mut note_commitments = Vec::with_capacity(decoded.len()); + for (output_idx, ((_, output), dec_output)) in decoded.iter().zip(decrypted_opts).enumerate() { + // Collect block note commitments + let node = extract_note_commitment(output); + let is_checkpoint = output_idx + 1 == decoded.len() && tx_idx + 1 == block_tx_count; + let retention = match (dec_output.is_some(), is_checkpoint) { + (is_marked, true) => Retention::Checkpoint { + id: block_height, + is_marked, + }, + (true, false) => Retention::Marked, + (false, false) => Retention::Ephemeral, + }; + + if let Some((note, account, scope, nk)) = dec_output { + // A note is marked as "change" if the account that received it + // also spent notes in the same transaction. This will catch, + // for instance: + // - Change created by spending fractions of notes. + // - Notes created by consolidation transactions. + // - Notes sent from one account to itself. + let is_change = spent_from_accounts.contains(&account); + let note_commitment_tree_position = Position::from(u64::from( + commitment_tree_size + u32::try_from(output_idx).unwrap(), + )); + let nf = SK::nf(&nk, ¬e, note_commitment_tree_position); + + shielded_outputs.push(new_wallet_output( + output_idx, + output, + account, + note, + is_change, + note_commitment_tree_position, + nf, + scope, + )); + } + + note_commitments.push((node, retention)) + } + + (shielded_outputs, note_commitments) +} + #[cfg(test)] mod tests { use group::{