zec
/
librustzcash
mirror of https://github.com/zcash/librustzcash.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #636 from zcash/batch-scanner-tag-ivks 

zcash_client_backend: Add tags to IVKs in the batch scanner
This commit is contained in:
Kris Nuttycombe 2022-09-13 16:49:10 -06:00 committed by GitHub
parent b5908dc964 f7b7760051
commit d4cbc04c16
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 84 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
zcash_client_backend/src/data_api/chain.rs
View File

@ -232,10 +232,12 @@ where
let mut batch_runner = BatchRunner::new(
100,
dfvks
.iter()
.flat_map(|(_, dfvk)| [dfvk.to_ivk(Scope::External), dfvk.to_ivk(Scope::Internal)])
.collect(),
dfvks.iter().flat_map(|(account, dfvk)| {
[
((**account, Scope::External), dfvk.to_ivk(Scope::External)),
((**account, Scope::Internal), dfvk.to_ivk(Scope::Internal)),
]
}),
);
cache.with_blocks(last_height, limit, |block: CompactBlock| {

58
zcash_client_backend/src/scan.rs
View File

@ -7,17 +7,18 @@ use zcash_note_encryption::{batch, BatchDomain, Domain, ShieldedOutput, COMPACT_
use zcash_primitives::{block::BlockHash, transaction::TxId};
/// A decrypted note.
pub(crate) struct DecryptedNote<D: Domain> {
/// The incoming viewing key used to decrypt the note.
pub(crate) ivk: D::IncomingViewingKey,
pub(crate) struct DecryptedNote<A, D: Domain> {
/// 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,
}
impl<D: Domain> fmt::Debug for DecryptedNote<D>
impl<A, D: Domain> fmt::Debug for DecryptedNote<A, D>
where
A: fmt::Debug,
D::IncomingViewingKey: fmt::Debug,
D::Recipient: fmt::Debug,
D::Note: fmt::Debug,
@ -25,7 +26,7 @@ where
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DecryptedNote")
.field("ivk", &self.ivk)
.field("ivk_tag", &self.ivk_tag)
.field("recipient", &self.recipient)
.field("note", &self.note)
.finish()
@ -40,10 +41,12 @@ struct OutputIndex<V> {
value: V,
}
type OutputReplier<D> = OutputIndex<channel::Sender<OutputIndex<Option<DecryptedNote<D>>>>>;
type OutputReplier<A, D> = OutputIndex<channel::Sender<OutputIndex<Option<DecryptedNote<A, D>>>>>;
type OutputResult<A, D> = channel::Receiver<OutputIndex<Option<DecryptedNote<A, D>>>>;
/// A batch of outputs to trial decrypt.
struct Batch<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> {
struct Batch<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> {
tags: Vec<A>,
ivks: Vec<D::IncomingViewingKey>,
/// We currently store outputs and repliers as parallel vectors, because
/// [`batch::try_note_decryption`] accepts a slice of domain/output pairs
@ -53,18 +56,20 @@ struct Batch<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> {
/// all be part of the same struct, which would also track the output index
/// (that is captured in the outer `OutputIndex` of each `OutputReplier`).
outputs: Vec<(D, Output)>,
repliers: Vec<OutputReplier<D>>,
repliers: Vec<OutputReplier<A, D>>,
}
impl<D, Output> Batch<D, Output>
impl<A, D, Output> Batch<A, D, Output>
where
A: Clone,
D: BatchDomain,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
D::IncomingViewingKey: Clone,
{
/// Constructs a new batch.
fn new(ivks: Vec<D::IncomingViewingKey>) -> Self {
fn new(tags: Vec<A>, ivks: Vec<D::IncomingViewingKey>) -> Self {
assert_eq!(tags.len(), ivks.len());
Self {
tags,
ivks,
outputs: vec![],
repliers: vec![],
@ -86,7 +91,7 @@ where
let result = OutputIndex {
output_index: replier.output_index,
value: decryption_result.map(|((note, recipient), ivk_idx)| DecryptedNote {
ivk: self.ivks[ivk_idx].clone(),
ivk_tag: self.tags[ivk_idx].clone(),
recipient,
note,
}),
@ -100,7 +105,7 @@ where
}
}
impl<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone> Batch<D, Output> {
impl<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone> Batch<A, D, Output> {
/// Adds the given outputs to this batch.
///
/// `replier` will be called with the result of every output.
@ -108,7 +113,7 @@ impl<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone> Batch
&mut self,
domain: impl Fn() -> D,
outputs: &[Output],
replier: channel::Sender<OutputIndex<Option<DecryptedNote<D>>>>,
replier: channel::Sender<OutputIndex<Option<DecryptedNote<A, D>>>>,
) {
self.outputs
.extend(outputs.iter().cloned().map(|output| (domain(), output)));
@ -123,30 +128,35 @@ impl<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone> Batch
type ResultKey = (BlockHash, TxId);
/// Logic to run batches of trial decryptions on the global threadpool.
pub(crate) struct BatchRunner<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> {
pub(crate) struct BatchRunner<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> {
batch_size_threshold: usize,
acc: Batch<D, Output>,
pending_results: HashMap<ResultKey, channel::Receiver<OutputIndex<Option<DecryptedNote<D>>>>>,
acc: Batch<A, D, Output>,
pending_results: HashMap<ResultKey, OutputResult<A, D>>,
}
impl<D, Output> BatchRunner<D, Output>
impl<A, D, Output> BatchRunner<A, D, Output>
where
A: Clone,
D: BatchDomain,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
D::IncomingViewingKey: Clone,
{
/// Constructs a new batch runner for the given incoming viewing keys.
pub(crate) fn new(batch_size_threshold: usize, ivks: Vec<D::IncomingViewingKey>) -> Self {
pub(crate) fn new(
batch_size_threshold: usize,
ivks: impl Iterator<Item = (A, D::IncomingViewingKey)>,
) -> Self {
let (tags, ivks) = ivks.unzip();
Self {
batch_size_threshold,
acc: Batch::new(ivks),
acc: Batch::new(tags, ivks),
pending_results: HashMap::default(),
}
}
}
impl<D, Output> BatchRunner<D, Output>
impl<A, D, Output> BatchRunner<A, D, Output>
where
A: Clone + Send + 'static,
D: BatchDomain + Send + 'static,
D::IncomingViewingKey: Clone + Send,
D::Memo: Send,
@ -184,7 +194,7 @@ where
/// Subsequent calls to `Self::add_outputs` will be accumulated into a new batch.
pub(crate) fn flush(&mut self) {
if !self.acc.is_empty() {
let mut batch = Batch::new(self.acc.ivks.clone());
let mut batch = Batch::new(self.acc.tags.clone(), self.acc.ivks.clone());
mem::swap(&mut batch, &mut self.acc);
rayon::spawn_fifo(|| batch.run());
}
@ -199,7 +209,7 @@ where
&mut self,
block_tag: BlockHash,
txid: TxId,
) -> HashMap<(TxId, usize), DecryptedNote<D>> {
) -> HashMap<(TxId, usize), DecryptedNote<A, D>> {
self.pending_results
.remove(&(block_tag, txid))
// We won't have a pending result if the transaction didn't have outputs of

65
zcash_client_backend/src/welding_rig.rs
View File

@ -38,10 +38,13 @@ use crate::{
/// [`CompactSaplingOutput`]: crate::proto::compact_formats::CompactSaplingOutput
/// [`scan_block`]: crate::welding_rig::scan_block
pub trait ScanningKey {
/// The type representing the scope of the scanning key.
type Scope: Clone + Eq + std::hash::Hash + Send + 'static;
/// The type of key that is used to decrypt Sapling outputs;
type SaplingNk: Clone;
type SaplingKeys: IntoIterator<Item = (SaplingIvk, Self::SaplingNk)>;
type SaplingKeys: IntoIterator<Item = (Self::Scope, SaplingIvk, Self::SaplingNk)>;
/// The type of nullifier extracted when a note is successfully
/// obtained by trial decryption.
@ -63,14 +66,23 @@ pub trait ScanningKey {
}
impl ScanningKey for DiversifiableFullViewingKey {
type Scope = Scope;
type SaplingNk = NullifierDerivingKey;
type SaplingKeys = [(SaplingIvk, Self::SaplingNk); 2];
type SaplingKeys = [(Self::Scope, SaplingIvk, Self::SaplingNk); 2];
type Nf = sapling::Nullifier;
fn to_sapling_keys(&self) -> Self::SaplingKeys {
[
(self.to_ivk(Scope::External), self.to_nk(Scope::External)),
(self.to_ivk(Scope::Internal), self.to_nk(Scope::Internal)),
(
Scope::External,
self.to_ivk(Scope::External),
self.to_nk(Scope::External),
),
(
Scope::Internal,
self.to_ivk(Scope::Internal),
self.to_nk(Scope::Internal),
),
]
}
@ -88,12 +100,13 @@ impl ScanningKey for DiversifiableFullViewingKey {
///
/// [`ExtendedFullViewingKey`]: zcash_primitives::zip32::ExtendedFullViewingKey
impl ScanningKey for ExtendedFullViewingKey {
type Scope = Scope;
type SaplingNk = NullifierDerivingKey;
type SaplingKeys = [(SaplingIvk, Self::SaplingNk); 1];
type SaplingKeys = [(Self::Scope, SaplingIvk, Self::SaplingNk); 1];
type Nf = sapling::Nullifier;
fn to_sapling_keys(&self) -> Self::SaplingKeys {
[(self.fvk.vk.ivk(), self.fvk.vk.nk)]
[(Scope::External, self.fvk.vk.ivk(), self.fvk.vk.nk)]
}
fn sapling_nf(
@ -110,12 +123,13 @@ impl ScanningKey for ExtendedFullViewingKey {
///
/// [`SaplingIvk`]: zcash_primitives::sapling::SaplingIvk
impl ScanningKey for SaplingIvk {
type Scope = ();
type SaplingNk = ();
type SaplingKeys = [(SaplingIvk, Self::SaplingNk); 1];
type SaplingKeys = [(Self::Scope, SaplingIvk, Self::SaplingNk); 1];
type Nf = ();
fn to_sapling_keys(&self) -> Self::SaplingKeys {
[(self.clone(), ())]
[((), self.clone(), ())]
}
fn sapling_nf(_key: &Self::SaplingNk, _note: &Note, _witness: &IncrementalWitness<Node>) {}
@ -163,11 +177,17 @@ pub fn scan_block<P: consensus::Parameters + Send + 'static, K: ScanningKey>(
)
}
pub(crate) fn add_block_to_runner<P: consensus::Parameters + Send + 'static>(
type TaggedBatchRunner<P, S> =
BatchRunner<(AccountId, S), SaplingDomain<P>, CompactOutputDescription>;
pub(crate) fn add_block_to_runner<P, S>(
params: &P,
block: CompactBlock,
batch_runner: &mut BatchRunner<SaplingDomain<P>, CompactOutputDescription>,
) {
batch_runner: &mut TaggedBatchRunner<P, S>,
) where
P: consensus::Parameters + Send + 'static,
S: Clone + Send + 'static,
{
let block_hash = block.hash();
let block_height = block.height();
@ -198,7 +218,7 @@ pub(crate) fn scan_block_with_runner<P: consensus::Parameters + Send + 'static,
nullifiers: &[(AccountId, Nullifier)],
tree: &mut CommitmentTree<Node>,
existing_witnesses: &mut [&mut IncrementalWitness<Node>],
mut batch_runner: Option<&mut BatchRunner<SaplingDomain<P>, CompactOutputDescription>>,
mut batch_runner: Option<&mut TaggedBatchRunner<P, K::Scope>>,
) -> Vec<WalletTx<K::Nf>> {
let mut wtxs: Vec<WalletTx<K::Nf>> = vec![];
let block_height = block.height();
@ -276,7 +296,7 @@ pub(crate) fn scan_block_with_runner<P: consensus::Parameters + Send + 'static,
.flat_map(|(a, k)| {
k.to_sapling_keys()
.into_iter()
.map(move |(ivk, nk)| (ivk.to_repr(), (**a, nk)))
.map(move |(scope, _, nk)| ((**a, scope), nk))
})
.collect::<HashMap<_, _>>();
@ -284,11 +304,12 @@ pub(crate) fn scan_block_with_runner<P: consensus::Parameters + Send + 'static,
(0..decoded.len())
.map(|i| {
decrypted.remove(&(txid, i)).map(|d_note| {
let (a, nk) = vks.get(&d_note.ivk.to_repr()).expect(
let a = d_note.ivk_tag.0;
let nk = vks.get(&d_note.ivk_tag).expect(
"The batch runner and scan_block must use the same set of IVKs.",
);
((d_note.note, d_note.recipient), *a, (*nk).clone())
((d_note.note, d_note.recipient), a, (*nk).clone())
})
})
.collect()
@ -298,7 +319,7 @@ pub(crate) fn scan_block_with_runner<P: consensus::Parameters + Send + 'static,
.flat_map(|(a, k)| {
k.to_sapling_keys()
.into_iter()
.map(move |(ivk, nk)| (**a, ivk, nk))
.map(move |(_, ivk, nk)| (**a, ivk, nk))
})
.collect::<Vec<_>>();
@ -516,6 +537,7 @@ mod tests {
#[test]
fn scan_block_with_my_tx() {
fn go(scan_multithreaded: bool) {
let account = AccountId::from(0);
let extsk = ExtendedSpendingKey::master(&[]);
let extfvk = ExtendedFullViewingKey::from(&extsk);
@ -535,8 +557,7 @@ mod tests {
extfvk
.to_sapling_keys()
.iter()
.map(|(k, _)| k.clone())
.collect(),
.map(|(scope, ivk, _)| ((account, *scope), ivk.clone())),
);
add_block_to_runner(&Network::TestNetwork, cb.clone(), &mut runner);
@ -550,7 +571,7 @@ mod tests {
let txs = scan_block_with_runner(
&Network::TestNetwork,
cb,
&[(&AccountId::from(0), &extfvk)],
&[(&account, &extfvk)],
&[],
&mut tree,
&mut [],
@ -565,7 +586,7 @@ mod tests {
assert_eq!(tx.shielded_spends.len(), 0);
assert_eq!(tx.shielded_outputs.len(), 1);
assert_eq!(tx.shielded_outputs[0].index, 0);
assert_eq!(tx.shielded_outputs[0].account, AccountId::from(0));
assert_eq!(tx.shielded_outputs[0].account, account);
assert_eq!(tx.shielded_outputs[0].note.value, 5);
// Check that the witness root matches
@ -579,6 +600,7 @@ mod tests {
#[test]
fn scan_block_with_txs_after_my_tx() {
fn go(scan_multithreaded: bool) {
let account = AccountId::from(0);
let extsk = ExtendedSpendingKey::master(&[]);
let extfvk = ExtendedFullViewingKey::from(&extsk);
@ -598,8 +620,7 @@ mod tests {
extfvk
.to_sapling_keys()
.iter()
.map(|(k, _)| k.clone())
.collect(),
.map(|(scope, ivk, _)| ((account, *scope), ivk.clone())),
);
add_block_to_runner(&Network::TestNetwork, cb.clone(), &mut runner);

2
zcash_primitives/src/sapling/keys.rs
View File

@ -191,7 +191,7 @@ impl FullViewingKey {
/// to other people. For example, a user can give an external [`SaplingIvk`] to a merchant
/// terminal, enabling it to only detect "real" transactions from customers and not
/// internal transactions from the wallet.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Scope {
/// A scope used for wallet-external operations, namely deriving addresses to give to
/// other users in order to receive funds.