From 4d3f8e8574fdd6eefe9586406af59bf356d4c898 Mon Sep 17 00:00:00 2001 From: Hanh Date: Tue, 15 Mar 2022 13:57:10 +0800 Subject: [PATCH] Batch trial decryption --- src/chain.rs | 138 ++++++++++++++++++++++++++++++++++++++++----------- src/scan.rs | 2 + 2 files changed, 112 insertions(+), 28 deletions(-) diff --git a/src/chain.rs b/src/chain.rs index 0dade68..efdb3a8 100644 --- a/src/chain.rs +++ b/src/chain.rs @@ -7,17 +7,19 @@ use ff::PrimeField; use log::info; use rayon::prelude::*; use std::collections::HashMap; +use std::marker::PhantomData; use std::time::Instant; use thiserror::Error; use tonic::transport::{Certificate, Channel, ClientTlsConfig}; use tonic::Request; +use zcash_note_encryption::batch::try_compact_note_decryption; use zcash_primitives::consensus::{BlockHeight, Network, NetworkUpgrade, Parameters}; use zcash_primitives::merkle_tree::{CommitmentTree, IncrementalWitness}; -use zcash_primitives::sapling::note_encryption::try_sapling_compact_note_decryption; +use zcash_primitives::sapling::note_encryption::{SaplingDomain, try_sapling_compact_note_decryption}; use zcash_primitives::sapling::{Node, Note, PaymentAddress}; use zcash_primitives::transaction::components::sapling::CompactOutputDescription; use zcash_primitives::zip32::ExtendedFullViewingKey; -use zcash_note_encryption::EphemeralKeyBytes; +use zcash_note_encryption::{COMPACT_NOTE_SIZE, Domain, EphemeralKeyBytes, ShieldedOutput}; const MAX_CHUNK: u32 = 50000; @@ -96,6 +98,7 @@ pub struct DecryptedBlock<'a> { pub count_outputs: u32, pub spends: Vec, pub compact_block: &'a CompactBlock, + pub elapsed: usize, } #[derive(Clone)] @@ -130,15 +133,64 @@ pub fn to_output_description(co: &CompactOutput) -> CompactOutputDescription { od } +struct AccountOutput<'a, N: Parameters> { + epk: EphemeralKeyBytes, + cmu: as Domain>::ExtractedCommitmentBytes, + ciphertext: [u8; COMPACT_NOTE_SIZE], + output_index: usize, + block_output_index: usize, + vtx: &'a CompactTx, + _phantom: PhantomData, +} + +impl <'a, N: Parameters> AccountOutput<'a, N> { + fn new(output_index: usize, block_output_index: usize, vtx: &'a CompactTx, co: &CompactOutput) -> Self { + let mut epk_bytes = [0u8; 32]; + epk_bytes.copy_from_slice(&co.epk); + let epk = EphemeralKeyBytes::from(epk_bytes); + let mut cmu_bytes = [0u8; 32]; + cmu_bytes.copy_from_slice(&co.cmu); + let cmu = as Domain>::ExtractedCommitmentBytes::from(cmu_bytes); + let mut ciphertext_bytes = [0u8; COMPACT_NOTE_SIZE]; + ciphertext_bytes.copy_from_slice(&co.ciphertext); + + AccountOutput { + output_index, + block_output_index, + vtx, + epk, + cmu, + ciphertext: ciphertext_bytes, + _phantom: PhantomData::default(), + } + } +} + +impl <'a, N: Parameters> ShieldedOutput, COMPACT_NOTE_SIZE> for AccountOutput<'a, N> { + fn ephemeral_key(&self) -> EphemeralKeyBytes { + self.epk.clone() + } + + fn cmstar_bytes(&self) -> as Domain>::ExtractedCommitmentBytes { + self.cmu + } + + fn enc_ciphertext(&self) -> &[u8; COMPACT_NOTE_SIZE] { + &self.ciphertext + } +} + fn decrypt_notes<'a, N: Parameters>( network: &N, block: &'a CompactBlock, - vks: &HashMap, + vks: &[(&u32, &AccountViewKey)], ) -> DecryptedBlock<'a> { let height = BlockHeight::from_u32(block.height as u32); let mut count_outputs = 0u32; let mut spends: Vec = vec![]; let mut notes: Vec = vec![]; + let vvks: Vec<_> = vks.iter().map(|vk| vk.1.ivk.clone()).collect(); + let mut outputs: Vec<(SaplingDomain, AccountOutput)> = vec![]; for (tx_index, vtx) in block.vtx.iter().enumerate() { for cs in vtx.spends.iter() { let mut nf = [0u8; 32]; @@ -146,40 +198,67 @@ fn decrypt_notes<'a, N: Parameters>( spends.push(Nf(nf)); } - // let _epks: Vec<_> = vtx.outputs.iter().map(|o| { - // &o.epk - // }).collect(); - - for (output_index, co) in vtx.outputs.iter().enumerate() { - let od = to_output_description(co); - for (&account, vk) in vks.iter() { - if let Some((note, pa)) = - try_sapling_compact_note_decryption(network, height, &vk.ivk, &od) - { - notes.push(DecryptedNote { - account, - ivk: vk.fvk.clone(), - note, - pa, - viewonly: vk.viewonly, - position_in_block: count_outputs as usize, - height: block.height as u32, - tx_index, - txid: vtx.hash.clone(), - output_index, - }); - } - } + let domain = SaplingDomain::::for_height(network.clone(), height); + let output = AccountOutput::::new(output_index, count_outputs as usize, vtx, co); + outputs.push((domain, output)); + + // let od = to_output_description(co); + // + // for (&account, vk) in vks.iter() { + // if let Some((note, pa)) = + // try_sapling_compact_note_decryption(network, height, &vk.ivk, &od) + // { + // notes.push(DecryptedNote { + // account, + // ivk: vk.fvk.clone(), + // note, + // pa, + // viewonly: vk.viewonly, + // position_in_block: count_outputs as usize, + // height: block.height as u32, + // tx_index, + // txid: vtx.hash.clone(), + // output_index, + // }); + // } + // } + count_outputs += 1; } } + + let start = Instant::now(); + let notes_decrypted = try_compact_note_decryption::, AccountOutput> + (&vvks, &outputs); + let elapsed = start.elapsed().as_millis() as usize; + + for (pos, opt_note) in notes_decrypted.iter().enumerate() { + if let Some((note, pa)) = opt_note { + let vk = &vks[pos / outputs.len()]; + let output = &outputs[pos % outputs.len()]; + notes.push(DecryptedNote { + account: *vk.0, + ivk: vk.1.fvk.clone(), + note: note.clone(), + pa: pa.clone(), + viewonly: vk.1.viewonly, + position_in_block: output.1.block_output_index, + height: block.height as u32, + tx_index: output.1.vtx.index as usize, + txid: output.1.vtx.hash.clone(), + output_index: output.1.output_index, + }); + } + } + DecryptedBlock { height: block.height as u32, spends, notes, count_outputs, compact_block: block, + elapsed, } } @@ -189,9 +268,10 @@ impl DecryptNode { } pub fn decrypt_blocks<'a>(&self, network: &Network, blocks: &'a [CompactBlock]) -> Vec> { + let vks: Vec<_> = self.vks.iter().collect(); let mut decrypted_blocks: Vec = blocks .par_iter() - .map(|b| decrypt_notes(network, b, &self.vks)) + .map(|b| decrypt_notes(network, b, &vks)) .collect(); decrypted_blocks.sort_by(|a, b| a.height.cmp(&b.height)); decrypted_blocks @@ -331,6 +411,8 @@ pub async fn sync(network: &Network, vks: HashMap, ld_url: let start = Instant::now(); let blocks = decrypter.decrypt_blocks(network, &cbs); eprintln!("Decrypt Notes: {} ms", start.elapsed().as_millis()); + let batch_decrypt_elapsed: usize = blocks.iter().map(|b| b.elapsed).sum(); + eprintln!(" Batch Decrypt: {} ms", batch_decrypt_elapsed); let start = Instant::now(); let witnesses = calculate_tree_state_v2(&cbs, &blocks); diff --git a/src/scan.rs b/src/scan.rs index 566f362..6528bdf 100644 --- a/src/scan.rs +++ b/src/scan.rs @@ -165,6 +165,8 @@ pub async fn sync_async( // db tx scope let db_tx = db.begin_transaction()?; let dec_blocks = decrypter.decrypt_blocks(&network, &blocks.0); + let batch_decrypt_elapsed: usize = dec_blocks.iter().map(|b| b.elapsed).sum(); + log::info!(" Batch Decrypt: {} ms", batch_decrypt_elapsed); for b in dec_blocks.iter() { let mut my_nfs: Vec = vec![]; for nf in b.spends.iter() {