Skip to content

Commit

Permalink
zcash_client_backend: Factor out common note decryption from `scan_bl…
Browse files Browse the repository at this point in the history 
…ock_with_runner`
  • Loading branch information
@nuttycom
nuttycom committed Feb 23, 2024
1 parent 7f754ce commit ebcb3dd
Show file tree
Hide file tree
Showing 2 changed files with 197 additions and 104 deletions.
23 changes: 22 additions & 1 deletion zcash_client_backend/src/scan.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -506,6 +506,14 @@ where
}
}

pub(crate) trait BatchResults<A, D: Domain, Memo> {
fn collect_results(
&mut self,
block_tag: BlockHash,
txid: TxId,
) -> HashMap<(TxId, usize), DecryptedOutput<A, D, Memo>>;
}

impl<A, D, Output, Dec, T> BatchRunner<A, D, Output, Dec, T>
where
A: Clone + Send + 'static,
Expand Down Expand Up @@ -554,13 +562,26 @@ where
self.running_tasks.run_task(batch);
}
}
}

impl<A, D, Output, Dec, T> BatchResults<A, D, Dec::Memo> for BatchRunner<A, D, Output, Dec, T>
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<D, Output>,
T: Tasks<Batch<A, D, Output, Dec>>,
{
/// 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,
Expand Down
278 changes: 175 additions & 103 deletions zcash_client_backend/src/scanning.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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<WalletSaplingOutput<SK::Nf, SK::Scope>> = 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| {
(
Expand All @@ -526,108 +534,32 @@ pub(crate) fn scan_block_with_runner<
.expect("Invalid output found in compact block decoding."),
)
})
.collect::<Vec<_>>();

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::<HashMap<_, _>>();

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::<Vec<_>>();

let ivks = sapling_keys
.iter()
.map(|(_, _, ivk, _)| PreparedIncomingViewingKey::new(ivk))
.collect::<Vec<_>>();

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, &note, 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::<Vec<_>>(),
&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,
});
}

Expand Down Expand Up @@ -714,6 +646,146 @@ fn check_nullifiers<Spend, Nf: ConstantTimeEq + Copy, WS>(
(found_spent, unlinked_nullifiers)
}

fn find_notes<
M,
D: BatchDomain,
SK: ScanningKey<Note = D::Note>,
Output: ShieldedOutput<D, 52>,
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<AccountId>,
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<WalletOutput>,
Vec<(NoteCommitment, Retention<BlockHeight>)>,
) {

Check failure on line 683 in zcash_client_backend/src/scanning.rs

View workflow job for this annotation

GitHub Actions / Clippy (MSRV)

this function has too many arguments (13/7) 

error: this function has too many arguments (13/7)
   --> zcash_client_backend/src/scanning.rs:649:1
    |
649 | / fn find_notes<
650 | |     M,
651 | |     D: BatchDomain,
652 | |     SK: ScanningKey<Note = D::Note>,
...   |
682 | |     Vec<(NoteCommitment, Retention<BlockHeight>)>,
683 | | ) {
    | |_^
    |
    = note: `-D clippy::too-many-arguments` implied by `-D warnings`
    = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#too_many_arguments

Check failure on line 683 in zcash_client_backend/src/scanning.rs

View workflow job for this annotation

GitHub Actions / Clippy (MSRV)

this function has too many arguments (13/7) 

error: this function has too many arguments (13/7)
   --> zcash_client_backend/src/scanning.rs:649:1
    |
649 | / fn find_notes<
650 | |     M,
651 | |     D: BatchDomain,
652 | |     SK: ScanningKey<Note = D::Note>,
...   |
682 | |     Vec<(NoteCommitment, Retention<BlockHeight>)>,
683 | | ) {
    | |_^
    |
    = note: `-D clippy::too-many-arguments` implied by `-D warnings`
    = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#too_many_arguments
// 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::<HashMap<_, _>>();

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::<Vec<_>>(),
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::<Vec<_>>();

let ivks = tagged_keys
.iter()
.map(|(_, _, ivk, _)| prepare_key(ivk))
.collect::<Vec<_>>();

let mut decrypted_len = 0;
(
batch::try_compact_note_decryption(&ivks, &decoded[..])

Check failure on line 729 in zcash_client_backend/src/scanning.rs

View workflow job for this annotation

GitHub Actions / Clippy (MSRV)

redundant slicing of the whole range 

error: redundant slicing of the whole range
   --> zcash_client_backend/src/scanning.rs:729:55
    |
729 |             batch::try_compact_note_decryption(&ivks, &decoded[..])
    |                                                       ^^^^^^^^^^^^ help: use the original value instead: `decoded`
    |
    = note: `-D clippy::redundant-slicing` implied by `-D warnings`
    = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#redundant_slicing

Check failure on line 729 in zcash_client_backend/src/scanning.rs

View workflow job for this annotation

GitHub Actions / Clippy (MSRV)

redundant slicing of the whole range 

error: redundant slicing of the whole range
   --> zcash_client_backend/src/scanning.rs:729:55
    |
729 |             batch::try_compact_note_decryption(&ivks, &decoded[..])
    |                                                       ^^^^^^^^^^^^ help: use the original value instead: `decoded`
    |
    = note: `-D clippy::redundant-slicing` implied by `-D warnings`
    = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#redundant_slicing
.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::<Vec<_>>(),
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, &note, 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::{
Expand Down

0 comments on commit ebcb3dd

Please sign in to comment.