From d56536c1d18e223dc9cec9f93dd92aa556ef1e24 Mon Sep 17 00:00:00 2001 From: Jack Grigg Date: Wed, 7 Feb 2024 16:50:33 +0000 Subject: [PATCH] zcash_client_backend: Introduce `Decryptor` trait This generalizes over batch decryption of either compact or full outputs. Closes zcash/librustzcash#1171. --- zcash_client_backend/src/data_api/chain.rs | 2 +- zcash_client_backend/src/scan.rs | 141 ++++++++++++++++----- zcash_client_backend/src/scanning.rs | 11 +- 3 files changed, 114 insertions(+), 40 deletions(-) diff --git a/zcash_client_backend/src/data_api/chain.rs b/zcash_client_backend/src/data_api/chain.rs index 82e027af3..1555785f8 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 6c5fad1c2..41e85e94c 100644 --- a/zcash_client_backend/src/scan.rs +++ b/zcash_client_backend/src/scan.rs @@ -8,7 +8,9 @@ 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 transaction output. @@ -41,6 +43,72 @@ where } } +/// 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. @@ -211,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 @@ -222,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() @@ -256,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 { @@ -279,7 +348,7 @@ where } } -impl Task for Batch +impl Task for Batch where A: Clone + Send + 'static, D: BatchDomain + Send + 'static, @@ -287,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) { @@ -301,21 +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: DecryptedOutput { - ivk_tag: tags[ivk_idx].clone(), - recipient, - note, - memo: (), - }, + value, }; if replier.value.send(result).is_err() { @@ -327,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. @@ -335,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))); @@ -365,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() @@ -412,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( @@ -434,7 +506,7 @@ where } } -impl BatchRunner +impl BatchRunner where A: Clone + Send + 'static, D: BatchDomain + Send + 'static, @@ -442,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. /// @@ -491,7 +564,7 @@ where &mut self, block_tag: BlockHash, txid: TxId, - ) -> HashMap<(TxId, usize), DecryptedOutput> { + ) -> 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 b0bad39cc..38d03881f 100644 --- a/zcash_client_backend/src/scanning.rs +++ b/zcash_client_backend/src/scanning.rs @@ -21,7 +21,7 @@ use zcash_primitives::{ use crate::data_api::{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 @@ pub fn scan_block( ) } -type TaggedBatch = Batch<(AccountId, S), SaplingDomain, CompactOutputDescription>; +type TaggedBatch = + 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( @@ -834,7 +835,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() @@ -921,7 +922,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()