diff --git a/zcash_client_backend/Cargo.toml b/zcash_client_backend/Cargo.toml index 8ff48f7eed..69cde1fc79 100644 --- a/zcash_client_backend/Cargo.toml +++ b/zcash_client_backend/Cargo.toml @@ -129,6 +129,7 @@ orchard = ["dep:orchard", "zcash_keys/orchard"] test-dependencies = [ "dep:proptest", "orchard?/test-dependencies", + "zcash_keys/test-dependencies", "zcash_primitives/test-dependencies", "incrementalmerkletree/test-dependencies", ] diff --git a/zcash_client_backend/src/data_api/chain.rs b/zcash_client_backend/src/data_api/chain.rs index 82e027af35..1555785f85 100644 --- a/zcash_client_backend/src/data_api/chain.rs +++ b/zcash_client_backend/src/data_api/chain.rs @@ -303,7 +303,7 @@ where .get_sapling_nullifiers(NullifierQuery::Unspent) .map_err(Error::Wallet)?; - let mut batch_runner = BatchRunner::<_, _, _, ()>::new( + let mut batch_runner = BatchRunner::<_, _, _, _, ()>::new( 100, dfvks .iter() diff --git a/zcash_client_backend/src/scan.rs b/zcash_client_backend/src/scan.rs index a568be4f0a..98a768622e 100644 --- a/zcash_client_backend/src/scan.rs +++ b/zcash_client_backend/src/scan.rs @@ -8,36 +8,107 @@ use std::sync::{ }; use memuse::DynamicUsage; -use zcash_note_encryption::{batch, BatchDomain, Domain, ShieldedOutput, COMPACT_NOTE_SIZE}; +use zcash_note_encryption::{ + batch, BatchDomain, Domain, ShieldedOutput, COMPACT_NOTE_SIZE, ENC_CIPHERTEXT_SIZE, +}; use zcash_primitives::{block::BlockHash, transaction::TxId}; -/// A decrypted note. -pub(crate) struct DecryptedNote { +/// A decrypted transaction output. +pub(crate) struct DecryptedOutput { /// The tag corresponding to the incoming viewing key used to decrypt the note. pub(crate) ivk_tag: A, /// The recipient of the note. pub(crate) recipient: D::Recipient, /// The note! pub(crate) note: D::Note, + /// The memo field, or `()` if this is a decrypted compact output. + pub(crate) memo: M, } -impl fmt::Debug for DecryptedNote +impl fmt::Debug for DecryptedOutput where A: fmt::Debug, D::IncomingViewingKey: fmt::Debug, D::Recipient: fmt::Debug, D::Note: fmt::Debug, - D::Memo: fmt::Debug, + M: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("DecryptedNote") + f.debug_struct("DecryptedOutput") .field("ivk_tag", &self.ivk_tag) .field("recipient", &self.recipient) .field("note", &self.note) + .field("memo", &self.memo) .finish() } } +/// A decryptor of transaction outputs. +pub(crate) trait Decryptor { + type Memo; + + // Once we reach MSRV 1.75.0, this can return `impl Iterator`. + fn batch_decrypt( + tags: &[A], + ivks: &[D::IncomingViewingKey], + outputs: &[(D, Output)], + ) -> Vec>>; +} + +/// A decryptor of outputs as encoded in transactions. +pub(crate) struct FullDecryptor; + +impl> Decryptor + for FullDecryptor +{ + type Memo = D::Memo; + + fn batch_decrypt( + tags: &[A], + ivks: &[D::IncomingViewingKey], + outputs: &[(D, Output)], + ) -> Vec>> { + batch::try_note_decryption(ivks, outputs) + .into_iter() + .map(|res| { + res.map(|((note, recipient, memo), ivk_idx)| DecryptedOutput { + ivk_tag: tags[ivk_idx].clone(), + recipient, + note, + memo, + }) + }) + .collect() + } +} + +/// A decryptor of outputs as encoded in compact blocks. +pub(crate) struct CompactDecryptor; + +impl> Decryptor + for CompactDecryptor +{ + type Memo = (); + + fn batch_decrypt( + tags: &[A], + ivks: &[D::IncomingViewingKey], + outputs: &[(D, Output)], + ) -> Vec>> { + batch::try_compact_note_decryption(ivks, outputs) + .into_iter() + .map(|res| { + res.map(|((note, recipient), ivk_idx)| DecryptedOutput { + ivk_tag: tags[ivk_idx].clone(), + recipient, + note, + memo: (), + }) + }) + .collect() + } +} + /// A value correlated with an output index. struct OutputIndex { /// The index of the output within the corresponding shielded bundle. @@ -46,12 +117,12 @@ struct OutputIndex { value: V, } -type OutputItem = OutputIndex>; +type OutputItem = OutputIndex>; /// The sender for the result of batch scanning a specific transaction output. -struct OutputReplier(OutputIndex>>); +struct OutputReplier(OutputIndex>>); -impl DynamicUsage for OutputReplier { +impl DynamicUsage for OutputReplier { #[inline(always)] fn dynamic_usage(&self) -> usize { // We count the memory usage of items in the channel on the receiver side. @@ -65,9 +136,9 @@ impl DynamicUsage for OutputReplier { } /// The receiver for the result of batch scanning a specific transaction. -struct BatchReceiver(channel::Receiver>); +struct BatchReceiver(channel::Receiver>); -impl DynamicUsage for BatchReceiver { +impl DynamicUsage for BatchReceiver { fn dynamic_usage(&self) -> usize { // We count the memory usage of items in the channel on the receiver side. let num_items = self.0.len(); @@ -84,7 +155,7 @@ impl DynamicUsage for BatchReceiver { // - Space for an item. // - The state of the slot, stored as an AtomicUsize. const PTR_SIZE: usize = std::mem::size_of::(); - let item_size = std::mem::size_of::>(); + let item_size = std::mem::size_of::>(); const ATOMIC_USIZE_SIZE: usize = std::mem::size_of::(); let block_size = PTR_SIZE + ITEMS_PER_BLOCK * (item_size + ATOMIC_USIZE_SIZE); @@ -208,7 +279,7 @@ impl Task for WithUsageTask { } /// A batch of outputs to trial decrypt. -pub(crate) struct Batch> { +pub(crate) struct Batch> { tags: Vec, ivks: Vec, /// We currently store outputs and repliers as parallel vectors, because @@ -219,15 +290,16 @@ pub(crate) struct Batch, - repliers: Vec>, + repliers: Vec>, } -impl DynamicUsage for Batch +impl DynamicUsage for Batch where A: DynamicUsage, D: BatchDomain + DynamicUsage, D::IncomingViewingKey: DynamicUsage, - Output: ShieldedOutput + DynamicUsage, + Output: DynamicUsage, + Dec: Decryptor, { fn dynamic_usage(&self) -> usize { self.tags.dynamic_usage() @@ -253,11 +325,11 @@ where } } -impl Batch +impl Batch where A: Clone, D: BatchDomain, - Output: ShieldedOutput, + Dec: Decryptor, { /// Constructs a new batch. fn new(tags: Vec, ivks: Vec) -> Self { @@ -276,7 +348,7 @@ where } } -impl Task for Batch +impl Task for Batch where A: Clone + Send + 'static, D: BatchDomain + Send + 'static, @@ -284,7 +356,9 @@ where D::Memo: Send, D::Note: Send, D::Recipient: Send, - Output: ShieldedOutput + Send + 'static, + Output: Send + 'static, + Dec: Decryptor + 'static, + Dec::Memo: Send, { /// Runs the batch of trial decryptions, and reports the results. fn run(self) { @@ -298,20 +372,16 @@ where assert_eq!(outputs.len(), repliers.len()); - let decryption_results = batch::try_compact_note_decryption(&ivks, &outputs); + let decryption_results = Dec::batch_decrypt(&tags, &ivks, &outputs); for (decryption_result, OutputReplier(replier)) in decryption_results.into_iter().zip(repliers.into_iter()) { // If `decryption_result` is `None` then we will just drop `replier`, // indicating to the parent `BatchRunner` that this output was not for us. - if let Some(((note, recipient), ivk_idx)) = decryption_result { + if let Some(value) = decryption_result { let result = OutputIndex { output_index: replier.output_index, - value: DecryptedNote { - ivk_tag: tags[ivk_idx].clone(), - recipient, - note, - }, + value, }; if replier.value.send(result).is_err() { @@ -323,7 +393,12 @@ where } } -impl + Clone> Batch { +impl Batch +where + D: BatchDomain, + Output: Clone, + Dec: Decryptor, +{ /// Adds the given outputs to this batch. /// /// `replier` will be called with the result of every output. @@ -331,7 +406,7 @@ impl + Clone> Ba &mut self, domain: impl Fn() -> D, outputs: &[Output], - replier: channel::Sender>, + replier: channel::Sender>, ) { self.outputs .extend(outputs.iter().cloned().map(|output| (domain(), output))); @@ -361,28 +436,29 @@ impl DynamicUsage for ResultKey { } /// Logic to run batches of trial decryptions on the global threadpool. -pub(crate) struct BatchRunner +pub(crate) struct BatchRunner where D: BatchDomain, - Output: ShieldedOutput, - T: Tasks>, + Dec: Decryptor, + T: Tasks>, { batch_size_threshold: usize, // The batch currently being accumulated. - acc: Batch, + acc: Batch, // The running batches. running_tasks: T, // Receivers for the results of the running batches. - pending_results: HashMap>, + pending_results: HashMap>, } -impl DynamicUsage for BatchRunner +impl DynamicUsage for BatchRunner where A: DynamicUsage, D: BatchDomain + DynamicUsage, D::IncomingViewingKey: DynamicUsage, - Output: ShieldedOutput + DynamicUsage, - T: Tasks> + DynamicUsage, + Output: DynamicUsage, + Dec: Decryptor, + T: Tasks> + DynamicUsage, { fn dynamic_usage(&self) -> usize { self.acc.dynamic_usage() @@ -408,12 +484,12 @@ where } } -impl BatchRunner +impl BatchRunner where A: Clone, D: BatchDomain, - Output: ShieldedOutput, - T: Tasks>, + Dec: Decryptor, + T: Tasks>, { /// Constructs a new batch runner for the given incoming viewing keys. pub(crate) fn new( @@ -430,7 +506,7 @@ where } } -impl BatchRunner +impl BatchRunner where A: Clone + Send + 'static, D: BatchDomain + Send + 'static, @@ -438,8 +514,9 @@ where D::Memo: Send, D::Note: Send, D::Recipient: Send, - Output: ShieldedOutput + Clone + Send + 'static, - T: Tasks>, + Output: Clone + Send + 'static, + Dec: Decryptor, + T: Tasks>, { /// Batches the given outputs for trial decryption. /// @@ -447,9 +524,9 @@ where /// batch, or the all-zeros hash to indicate that no block triggered it (i.e. it was a /// mempool change). /// - /// If after adding the given outputs, the accumulated batch size is at least - /// `BATCH_SIZE_THRESHOLD`, `Self::flush` is called. Subsequent calls to - /// `Self::add_outputs` will be accumulated into a new batch. + /// If after adding the given outputs, the accumulated batch size is at least the size + /// threshold that was set via `Self::new`, `Self::flush` is called. Subsequent calls + /// to `Self::add_outputs` will be accumulated into a new batch. pub(crate) fn add_outputs( &mut self, block_tag: BlockHash, @@ -487,7 +564,7 @@ where &mut self, block_tag: BlockHash, txid: TxId, - ) -> HashMap<(TxId, usize), DecryptedNote> { + ) -> HashMap<(TxId, usize), DecryptedOutput> { self.pending_results .remove(&ResultKey(block_tag, txid)) // We won't have a pending result if the transaction didn't have outputs of diff --git a/zcash_client_backend/src/scanning.rs b/zcash_client_backend/src/scanning.rs index 1e3d0bd59b..6e521165d1 100644 --- a/zcash_client_backend/src/scanning.rs +++ b/zcash_client_backend/src/scanning.rs @@ -17,7 +17,7 @@ use zcash_primitives::{consensus, zip32::Scope}; use crate::data_api::{AccountId, BlockMetadata, ScannedBlock, ScannedBundles}; use crate::{ proto::compact_formats::CompactBlock, - scan::{Batch, BatchRunner, Tasks}, + scan::{Batch, BatchRunner, CompactDecryptor, Tasks}, wallet::{WalletSaplingOutput, WalletSaplingSpend, WalletTx}, ShieldedProtocol, }; @@ -268,9 +268,10 @@ where } type TaggedBatch = - Batch<(AccountId, S), SaplingDomain, CompactOutputDescription>; + + Batch<(AccountId, S), SaplingDomain, CompactOutputDescription, CompactDecryptor>; type TaggedBatchRunner = - BatchRunner<(AccountId, S), SaplingDomain, CompactOutputDescription, T>; + BatchRunner<(AccountId, S), SaplingDomain, CompactOutputDescription, CompactDecryptor, T>; #[tracing::instrument(skip_all, fields(height = block.height))] pub(crate) fn add_block_to_runner( @@ -528,13 +529,13 @@ where 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 = vks.get(&d_note.ivk_tag).expect( + decrypted.remove(&(txid, i)).map(|d_out| { + let a = d_out.ivk_tag.0; + let nk = vks.get(&d_out.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()) + (d_out.note, a, d_out.ivk_tag.1, (*nk).clone()) }) }) .collect() @@ -850,7 +851,7 @@ mod tests { assert_eq!(cb.vtx.len(), 2); let mut batch_runner = if scan_multithreaded { - let mut runner = BatchRunner::<_, _, _, ()>::new( + let mut runner = BatchRunner::<_, _, _, _, ()>::new( 10, dfvk.to_sapling_keys() .iter() @@ -937,7 +938,7 @@ mod tests { assert_eq!(cb.vtx.len(), 3); let mut batch_runner = if scan_multithreaded { - let mut runner = BatchRunner::<_, _, _, ()>::new( + let mut runner = BatchRunner::<_, _, _, _, ()>::new( 10, dfvk.to_sapling_keys() .iter()