Use batch decryption in wallet scanning.

This modifies wallet scanning to perform per-block batched
decryption. It also alters the structure of the `ScanningKey`
trait to correctly include internal (change) keys in the scan
process.

components/zcash_note_encryption/src/batch.rs

@@ -47,6 +47,10 @@ fn batch_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, CS>, F, FR, c
 where
     F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, &D::SymmetricKey) -> Option<FR>,
 {
+    if ivks.is_empty() {
+        return (0..outputs.len()).map(|_| None).collect();
+    };
+
     // Fetch the ephemeral keys for each output and batch-parse them.
     let ephemeral_keys = D::batch_epk(outputs.iter().map(|(_, output)| output.ephemeral_key()));
 

zcash_client_backend/CHANGELOG.md

@@ -6,6 +6,7 @@ and this library adheres to Rust's notion of
 [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+
 ### Added
 - Functionality that enables the receiving and spending of transparent funds,
   behind the new `transparent-inputs` feature flag.
@@ -107,6 +108,12 @@ and this library adheres to Rust's notion of
   - `Zip321Error::TransparentMemo(usize)`
   - `Zip321Error::RecipientMissing(usize)`
   - `Zip321Error::ParseError(String)`
+- The api of `welding_rig::ScanningKey` has changed to accommodate batch
+  decryption and to correctly handle scanning with the internal (change) keys
+  derived from ZIP 316 UFVKs and UIVKs.
+- `welding_rig::scan_block` now uses batching for trial-decryption of
+  transaction outputs.
+
 
 ### Removed
 - `zcash_client_backend::data_api`:

zcash_client_backend/src/welding_rig.rs

@@ -4,15 +4,15 @@ use ff::PrimeField;
 use std::collections::HashSet;
 use std::convert::TryFrom;
 use subtle::{ConditionallySelectable, ConstantTimeEq, CtOption};
-use zcash_note_encryption::{ShieldedOutput, COMPACT_NOTE_SIZE};
+use zcash_note_encryption::batch;
 use zcash_primitives::{
-    consensus::{self, BlockHeight},
+    consensus,
     merkle_tree::{CommitmentTree, IncrementalWitness},
     sapling::{
         self,
-        keys::DiversifiableFullViewingKey,
+        keys::{DiversifiableFullViewingKey, Scope},
-        note_encryption::{try_sapling_compact_note_decryption, SaplingDomain},
+        note_encryption::SaplingDomain,
-        Node, Note, Nullifier, PaymentAddress, SaplingIvk,
+        Node, Note, Nullifier, NullifierDerivingKey, SaplingIvk,
     },
     transaction::components::sapling::CompactOutputDescription,
     zip32::{AccountId, ExtendedFullViewingKey},
@@ -21,58 +21,6 @@ use zcash_primitives::{
 use crate::proto::compact_formats::CompactBlock;
 use crate::wallet::{WalletShieldedOutput, WalletShieldedSpend, WalletTx};
 
-/// Scans a [`CompactSaplingOutput`] with a set of [`ScanningKey`]s.
-///
-/// Returns a [`WalletShieldedOutput`] and corresponding [`IncrementalWitness`] if this
-/// output belongs to any of the given [`ScanningKey`]s.
-///
-/// The given [`CommitmentTree`] and existing [`IncrementalWitness`]es are incremented
-/// with this output's commitment.
-///
-/// [`ScanningKey`]: crate::welding_rig::ScanningKey
-#[allow(clippy::too_many_arguments)]
-fn scan_output<P: consensus::Parameters, K: ScanningKey>(
-    params: &P,
-    height: BlockHeight,
-    index: usize,
-    output: CompactOutputDescription,
-    vks: &[(&AccountId, &K)],
-    spent_from_accounts: &HashSet<AccountId>,
-    tree: &mut CommitmentTree<Node>,
-) -> Option<WalletShieldedOutput<K::Nf>> {
-    for (account, vk) in vks.iter() {
-        let (note, to) = match vk.try_decryption(params, height, &output) {
-            Some(ret) => ret,
-            None => continue,
-        };
-
-        // 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 witness = IncrementalWitness::from_tree(tree);
-        let nf = vk.nf(&note, &witness);
-
-        return Some(WalletShieldedOutput {
-            index,
-            cmu: output.cmu,
-            ephemeral_key: output.ephemeral_key,
-            account: **account,
-            note,
-            to,
-            is_change,
-            witness,
-            nf,
-        });
-    }
-
-    None
-}
-
 /// A key that can be used to perform trial decryption and nullifier
 /// computation for a Sapling [`CompactSaplingOutput`]
 ///
@@ -87,47 +35,45 @@ fn scan_output<P: consensus::Parameters, K: ScanningKey>(
 /// [`CompactSaplingOutput`]: crate::proto::compact_formats::CompactSaplingOutput
 /// [`scan_block`]: crate::welding_rig::scan_block
 pub trait ScanningKey {
+    /// The type of key that is used to decrypt Sapling outputs;
+    type SaplingNk;
+
     /// The type of nullifier extracted when a note is successfully
     /// obtained by trial decryption.
     type Nf;
 
-    /// Attempts to decrypt a Sapling note and payment address
+    /// Obtain the underlying Sapling incoming viewing key(s) for this scanning key.
-    /// from the specified ciphertext using this scanning key.
+    fn to_sapling_keys(&self) -> Vec<(SaplingIvk, Self::SaplingNk)>;
-    fn try_decryption<
-        P: consensus::Parameters,
-        Output: ShieldedOutput<SaplingDomain<P>, COMPACT_NOTE_SIZE>,
-    >(
-        &self,
-        params: &P,
-        height: BlockHeight,
-        output: &Output,
-    ) -> Option<(Note, PaymentAddress)>;
 
     /// Produces the nullifier for the specified note and witness, if possible.
     ///
     /// IVK-based implementations of this trait cannot successfully derive
     /// nullifiers, in which case `Self::Nf` should be set to the unit type
     /// and this function is a no-op.
-    fn nf(&self, note: &Note, witness: &IncrementalWitness<Node>) -> Self::Nf;
+    fn sapling_nf(
+        key: &Self::SaplingNk,
+        note: &Note,
+        witness: &IncrementalWitness<Node>,
+    ) -> Self::Nf;
 }
 
 impl ScanningKey for DiversifiableFullViewingKey {
+    type SaplingNk = NullifierDerivingKey;
     type Nf = sapling::Nullifier;
 
-    fn try_decryption<
+    fn to_sapling_keys(&self) -> Vec<(SaplingIvk, Self::SaplingNk)> {
-        P: consensus::Parameters,
+        vec![
-        Output: ShieldedOutput<SaplingDomain<P>, COMPACT_NOTE_SIZE>,
+            (self.to_ivk(Scope::External), self.to_nk(Scope::External)),
-    >(
+            (self.to_ivk(Scope::Internal), self.to_nk(Scope::Internal)),
-        &self,
+        ]
-        params: &P,
-        height: BlockHeight,
-        output: &Output,
-    ) -> Option<(Note, PaymentAddress)> {
-        try_sapling_compact_note_decryption(params, height, &self.fvk().vk.ivk(), output)
     }
 
-    fn nf(&self, note: &Note, witness: &IncrementalWitness<Node>) -> Self::Nf {
+    fn sapling_nf(
-        note.nf(&self.fvk().vk.nk, witness.position() as u64)
+        key: &Self::SaplingNk,
+        note: &Note,
+        witness: &IncrementalWitness<Node>,
+    ) -> Self::Nf {
+        note.nf(key, witness.position() as u64)
     }
 }
 
@@ -136,22 +82,19 @@ impl ScanningKey for DiversifiableFullViewingKey {
 ///
 /// [`ExtendedFullViewingKey`]: zcash_primitives::zip32::ExtendedFullViewingKey
 impl ScanningKey for ExtendedFullViewingKey {
-    type Nf = Nullifier;
+    type SaplingNk = NullifierDerivingKey;
+    type Nf = sapling::Nullifier;
 
-    fn try_decryption<
+    fn to_sapling_keys(&self) -> Vec<(SaplingIvk, Self::SaplingNk)> {
-        P: consensus::Parameters,
+        vec![(self.fvk.vk.ivk(), self.fvk.vk.nk)]
-        Output: ShieldedOutput<SaplingDomain<P>, COMPACT_NOTE_SIZE>,
-    >(
-        &self,
-        params: &P,
-        height: BlockHeight,
-        output: &Output,
-    ) -> Option<(Note, PaymentAddress)> {
-        try_sapling_compact_note_decryption(params, height, &self.fvk.vk.ivk(), output)
     }
 
-    fn nf(&self, note: &Note, witness: &IncrementalWitness<Node>) -> Self::Nf {
+    fn sapling_nf(
-        note.nf(&self.fvk.vk.nk, witness.position() as u64)
+        key: &Self::SaplingNk,
+        note: &Note,
+        witness: &IncrementalWitness<Node>,
+    ) -> Self::Nf {
+        note.nf(key, witness.position() as u64)
     }
 }
 
@@ -160,21 +103,14 @@ impl ScanningKey for ExtendedFullViewingKey {
 ///
 /// [`SaplingIvk`]: zcash_primitives::sapling::SaplingIvk
 impl ScanningKey for SaplingIvk {
+    type SaplingNk = ();
     type Nf = ();
 
-    fn try_decryption<
+    fn to_sapling_keys(&self) -> Vec<(SaplingIvk, Self::SaplingNk)> {
-        P: consensus::Parameters,
+        vec![(self.clone(), ())]
-        Output: ShieldedOutput<SaplingDomain<P>, COMPACT_NOTE_SIZE>,
-    >(
-        &self,
-        params: &P,
-        height: BlockHeight,
-        output: &Output,
-    ) -> Option<(Note, PaymentAddress)> {
-        try_sapling_compact_note_decryption(params, height, self, output)
     }
 
-    fn nf(&self, _note: &Note, _witness: &IncrementalWitness<Node>) {}
+    fn sapling_nf(_key: &Self::SaplingNk, _note: &Note, _witness: &IncrementalWitness<Node>) {}
 }
 
 /// Scans a [`CompactBlock`] with a set of [`ScanningKey`]s.
@@ -265,17 +201,43 @@ pub fn scan_block<P: consensus::Parameters, K: ScanningKey>(
                 })
                 .collect();
 
-            for (index, output) in tx.outputs.into_iter().enumerate() {
+            let decoded = &tx
+                .outputs
+                .into_iter()
+                .map(|output| {
+                    (
+                        SaplingDomain::for_height(params.clone(), block_height),
+                        CompactOutputDescription::try_from(output)
+                            .expect("Invalid output found in compact block decoding."),
+                    )
+                })
+                .collect::<Vec<_>>();
+
+            let vks = vks
+                .iter()
+                .flat_map(|(a, k)| {
+                    k.to_sapling_keys()
+                        .into_iter()
+                        .map(move |(ivk, nk)| (**a, ivk, nk))
+                })
+                .collect::<Vec<_>>();
+
+            let ivks = vks
+                .iter()
+                .map(|(_, ivk, _)| (*ivk).clone())
+                .collect::<Vec<_>>();
+
+            let decrypted = batch::try_compact_note_decryption(&ivks, decoded);
+
+            for (index, ((_, output), dec_output)) in decoded.iter().zip(decrypted).enumerate() {
                 // Grab mutable references to new witnesses from previous outputs
                 // in this transaction so that we can update them. Scoped so we
                 // don't hold mutable references to shielded_outputs for too long.
                 let new_witnesses: Vec<_> = shielded_outputs
                     .iter_mut()
-                    .map(|output| &mut output.witness)
+                    .map(|out| &mut out.witness)
                     .collect();
 
-                let output = CompactOutputDescription::try_from(output).ok().unwrap();
-
                 // Increment tree and witnesses
                 let node = Node::new(output.cmu.to_repr());
                 for witness in &mut *existing_witnesses {
@@ -289,16 +251,29 @@ pub fn scan_block<P: consensus::Parameters, K: ScanningKey>(
                 }
                 tree.append(node).unwrap();
 
-                if let Some(output) = scan_output(
+                if let Some(((note, to), ivk_idx)) = dec_output {
-                    params,
+                    // A note is marked as "change" if the account that received it
-                    block_height,
+                    // also spent notes in the same transaction. This will catch,
-                    index,
+                    // for instance:
-                    output,
+                    // - Change created by spending fractions of notes.
-                    vks,
+                    // - Notes created by consolidation transactions.
-                    &spent_from_accounts,
+                    // - Notes sent from one account to itself.
-                    tree,
+                    let (account, _, nk) = &vks[ivk_idx];
-                ) {
+                    let is_change = spent_from_accounts.contains(account);
-                    shielded_outputs.push(output);
+                    let witness = IncrementalWitness::from_tree(tree);
+                    let nf = K::sapling_nf(nk, &note, &witness);
+
+                    shielded_outputs.push(WalletShieldedOutput {
+                        index,
+                        cmu: output.cmu,
+                        ephemeral_key: output.ephemeral_key.clone(),
+                        account: *account,
+                        note,
+                        to,
+                        is_change,
+                        witness,
+                        nf,
+                    })
                 }
             }
         }