zcash_client_backend: Introduce `Decryptor` trait
This generalizes over batch decryption of either compact or full outputs. Closes zcash/librustzcash#1171.
This commit is contained in:
parent
11db94bf42
commit
d56536c1d1
|
@ -303,7 +303,7 @@ where
|
||||||
.get_sapling_nullifiers(NullifierQuery::Unspent)
|
.get_sapling_nullifiers(NullifierQuery::Unspent)
|
||||||
.map_err(Error::Wallet)?;
|
.map_err(Error::Wallet)?;
|
||||||
|
|
||||||
let mut batch_runner = BatchRunner::<_, _, _, ()>::new(
|
let mut batch_runner = BatchRunner::<_, _, _, _, ()>::new(
|
||||||
100,
|
100,
|
||||||
dfvks
|
dfvks
|
||||||
.iter()
|
.iter()
|
||||||
|
|
|
@ -8,7 +8,9 @@ use std::sync::{
|
||||||
};
|
};
|
||||||
|
|
||||||
use memuse::DynamicUsage;
|
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};
|
use zcash_primitives::{block::BlockHash, transaction::TxId};
|
||||||
|
|
||||||
/// A decrypted transaction output.
|
/// A decrypted transaction output.
|
||||||
|
@ -41,6 +43,72 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A decryptor of transaction outputs.
|
||||||
|
pub(crate) trait Decryptor<D: BatchDomain, Output> {
|
||||||
|
type Memo;
|
||||||
|
|
||||||
|
// Once we reach MSRV 1.75.0, this can return `impl Iterator`.
|
||||||
|
fn batch_decrypt<A: Clone>(
|
||||||
|
tags: &[A],
|
||||||
|
ivks: &[D::IncomingViewingKey],
|
||||||
|
outputs: &[(D, Output)],
|
||||||
|
) -> Vec<Option<DecryptedOutput<A, D, Self::Memo>>>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A decryptor of outputs as encoded in transactions.
|
||||||
|
pub(crate) struct FullDecryptor;
|
||||||
|
|
||||||
|
impl<D: BatchDomain, Output: ShieldedOutput<D, ENC_CIPHERTEXT_SIZE>> Decryptor<D, Output>
|
||||||
|
for FullDecryptor
|
||||||
|
{
|
||||||
|
type Memo = D::Memo;
|
||||||
|
|
||||||
|
fn batch_decrypt<A: Clone>(
|
||||||
|
tags: &[A],
|
||||||
|
ivks: &[D::IncomingViewingKey],
|
||||||
|
outputs: &[(D, Output)],
|
||||||
|
) -> Vec<Option<DecryptedOutput<A, D, Self::Memo>>> {
|
||||||
|
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<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> Decryptor<D, Output>
|
||||||
|
for CompactDecryptor
|
||||||
|
{
|
||||||
|
type Memo = ();
|
||||||
|
|
||||||
|
fn batch_decrypt<A: Clone>(
|
||||||
|
tags: &[A],
|
||||||
|
ivks: &[D::IncomingViewingKey],
|
||||||
|
outputs: &[(D, Output)],
|
||||||
|
) -> Vec<Option<DecryptedOutput<A, D, Self::Memo>>> {
|
||||||
|
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.
|
/// A value correlated with an output index.
|
||||||
struct OutputIndex<V> {
|
struct OutputIndex<V> {
|
||||||
/// The index of the output within the corresponding shielded bundle.
|
/// The index of the output within the corresponding shielded bundle.
|
||||||
|
@ -211,7 +279,7 @@ impl<Item: Task> Task for WithUsageTask<Item> {
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A batch of outputs to trial decrypt.
|
/// A batch of outputs to trial decrypt.
|
||||||
pub(crate) struct Batch<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> {
|
pub(crate) struct Batch<A, D: BatchDomain, Output, Dec: Decryptor<D, Output>> {
|
||||||
tags: Vec<A>,
|
tags: Vec<A>,
|
||||||
ivks: Vec<D::IncomingViewingKey>,
|
ivks: Vec<D::IncomingViewingKey>,
|
||||||
/// We currently store outputs and repliers as parallel vectors, because
|
/// We currently store outputs and repliers as parallel vectors, because
|
||||||
|
@ -222,15 +290,16 @@ pub(crate) struct Batch<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOT
|
||||||
/// all be part of the same struct, which would also track the output index
|
/// all be part of the same struct, which would also track the output index
|
||||||
/// (that is captured in the outer `OutputIndex` of each `OutputReplier`).
|
/// (that is captured in the outer `OutputIndex` of each `OutputReplier`).
|
||||||
outputs: Vec<(D, Output)>,
|
outputs: Vec<(D, Output)>,
|
||||||
repliers: Vec<OutputReplier<A, D, ()>>,
|
repliers: Vec<OutputReplier<A, D, Dec::Memo>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<A, D, Output> DynamicUsage for Batch<A, D, Output>
|
impl<A, D, Output, Dec> DynamicUsage for Batch<A, D, Output, Dec>
|
||||||
where
|
where
|
||||||
A: DynamicUsage,
|
A: DynamicUsage,
|
||||||
D: BatchDomain + DynamicUsage,
|
D: BatchDomain + DynamicUsage,
|
||||||
D::IncomingViewingKey: DynamicUsage,
|
D::IncomingViewingKey: DynamicUsage,
|
||||||
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + DynamicUsage,
|
Output: DynamicUsage,
|
||||||
|
Dec: Decryptor<D, Output>,
|
||||||
{
|
{
|
||||||
fn dynamic_usage(&self) -> usize {
|
fn dynamic_usage(&self) -> usize {
|
||||||
self.tags.dynamic_usage()
|
self.tags.dynamic_usage()
|
||||||
|
@ -256,11 +325,11 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<A, D, Output> Batch<A, D, Output>
|
impl<A, D, Output, Dec> Batch<A, D, Output, Dec>
|
||||||
where
|
where
|
||||||
A: Clone,
|
A: Clone,
|
||||||
D: BatchDomain,
|
D: BatchDomain,
|
||||||
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
|
Dec: Decryptor<D, Output>,
|
||||||
{
|
{
|
||||||
/// Constructs a new batch.
|
/// Constructs a new batch.
|
||||||
fn new(tags: Vec<A>, ivks: Vec<D::IncomingViewingKey>) -> Self {
|
fn new(tags: Vec<A>, ivks: Vec<D::IncomingViewingKey>) -> Self {
|
||||||
|
@ -279,7 +348,7 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<A, D, Output> Task for Batch<A, D, Output>
|
impl<A, D, Output, Dec> Task for Batch<A, D, Output, Dec>
|
||||||
where
|
where
|
||||||
A: Clone + Send + 'static,
|
A: Clone + Send + 'static,
|
||||||
D: BatchDomain + Send + 'static,
|
D: BatchDomain + Send + 'static,
|
||||||
|
@ -287,7 +356,9 @@ where
|
||||||
D::Memo: Send,
|
D::Memo: Send,
|
||||||
D::Note: Send,
|
D::Note: Send,
|
||||||
D::Recipient: Send,
|
D::Recipient: Send,
|
||||||
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Send + 'static,
|
Output: Send + 'static,
|
||||||
|
Dec: Decryptor<D, Output> + 'static,
|
||||||
|
Dec::Memo: Send,
|
||||||
{
|
{
|
||||||
/// Runs the batch of trial decryptions, and reports the results.
|
/// Runs the batch of trial decryptions, and reports the results.
|
||||||
fn run(self) {
|
fn run(self) {
|
||||||
|
@ -301,21 +372,16 @@ where
|
||||||
|
|
||||||
assert_eq!(outputs.len(), repliers.len());
|
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
|
for (decryption_result, OutputReplier(replier)) in
|
||||||
decryption_results.into_iter().zip(repliers.into_iter())
|
decryption_results.into_iter().zip(repliers.into_iter())
|
||||||
{
|
{
|
||||||
// If `decryption_result` is `None` then we will just drop `replier`,
|
// If `decryption_result` is `None` then we will just drop `replier`,
|
||||||
// indicating to the parent `BatchRunner` that this output was not for us.
|
// 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 {
|
let result = OutputIndex {
|
||||||
output_index: replier.output_index,
|
output_index: replier.output_index,
|
||||||
value: DecryptedOutput {
|
value,
|
||||||
ivk_tag: tags[ivk_idx].clone(),
|
|
||||||
recipient,
|
|
||||||
note,
|
|
||||||
memo: (),
|
|
||||||
},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
if replier.value.send(result).is_err() {
|
if replier.value.send(result).is_err() {
|
||||||
|
@ -327,7 +393,12 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone> Batch<A, D, Output> {
|
impl<A, D, Output, Dec> Batch<A, D, Output, Dec>
|
||||||
|
where
|
||||||
|
D: BatchDomain,
|
||||||
|
Output: Clone,
|
||||||
|
Dec: Decryptor<D, Output>,
|
||||||
|
{
|
||||||
/// Adds the given outputs to this batch.
|
/// Adds the given outputs to this batch.
|
||||||
///
|
///
|
||||||
/// `replier` will be called with the result of every output.
|
/// `replier` will be called with the result of every output.
|
||||||
|
@ -335,7 +406,7 @@ impl<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone> Ba
|
||||||
&mut self,
|
&mut self,
|
||||||
domain: impl Fn() -> D,
|
domain: impl Fn() -> D,
|
||||||
outputs: &[Output],
|
outputs: &[Output],
|
||||||
replier: channel::Sender<OutputItem<A, D, ()>>,
|
replier: channel::Sender<OutputItem<A, D, Dec::Memo>>,
|
||||||
) {
|
) {
|
||||||
self.outputs
|
self.outputs
|
||||||
.extend(outputs.iter().cloned().map(|output| (domain(), output)));
|
.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.
|
/// Logic to run batches of trial decryptions on the global threadpool.
|
||||||
pub(crate) struct BatchRunner<A, D, Output, T>
|
pub(crate) struct BatchRunner<A, D, Output, Dec, T>
|
||||||
where
|
where
|
||||||
D: BatchDomain,
|
D: BatchDomain,
|
||||||
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
|
Dec: Decryptor<D, Output>,
|
||||||
T: Tasks<Batch<A, D, Output>>,
|
T: Tasks<Batch<A, D, Output, Dec>>,
|
||||||
{
|
{
|
||||||
batch_size_threshold: usize,
|
batch_size_threshold: usize,
|
||||||
// The batch currently being accumulated.
|
// The batch currently being accumulated.
|
||||||
acc: Batch<A, D, Output>,
|
acc: Batch<A, D, Output, Dec>,
|
||||||
// The running batches.
|
// The running batches.
|
||||||
running_tasks: T,
|
running_tasks: T,
|
||||||
// Receivers for the results of the running batches.
|
// Receivers for the results of the running batches.
|
||||||
pending_results: HashMap<ResultKey, BatchReceiver<A, D, ()>>,
|
pending_results: HashMap<ResultKey, BatchReceiver<A, D, Dec::Memo>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<A, D, Output, T> DynamicUsage for BatchRunner<A, D, Output, T>
|
impl<A, D, Output, Dec, T> DynamicUsage for BatchRunner<A, D, Output, Dec, T>
|
||||||
where
|
where
|
||||||
A: DynamicUsage,
|
A: DynamicUsage,
|
||||||
D: BatchDomain + DynamicUsage,
|
D: BatchDomain + DynamicUsage,
|
||||||
D::IncomingViewingKey: DynamicUsage,
|
D::IncomingViewingKey: DynamicUsage,
|
||||||
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + DynamicUsage,
|
Output: DynamicUsage,
|
||||||
T: Tasks<Batch<A, D, Output>> + DynamicUsage,
|
Dec: Decryptor<D, Output>,
|
||||||
|
T: Tasks<Batch<A, D, Output, Dec>> + DynamicUsage,
|
||||||
{
|
{
|
||||||
fn dynamic_usage(&self) -> usize {
|
fn dynamic_usage(&self) -> usize {
|
||||||
self.acc.dynamic_usage()
|
self.acc.dynamic_usage()
|
||||||
|
@ -412,12 +484,12 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<A, D, Output, T> BatchRunner<A, D, Output, T>
|
impl<A, D, Output, Dec, T> BatchRunner<A, D, Output, Dec, T>
|
||||||
where
|
where
|
||||||
A: Clone,
|
A: Clone,
|
||||||
D: BatchDomain,
|
D: BatchDomain,
|
||||||
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
|
Dec: Decryptor<D, Output>,
|
||||||
T: Tasks<Batch<A, D, Output>>,
|
T: Tasks<Batch<A, D, Output, Dec>>,
|
||||||
{
|
{
|
||||||
/// Constructs a new batch runner for the given incoming viewing keys.
|
/// Constructs a new batch runner for the given incoming viewing keys.
|
||||||
pub(crate) fn new(
|
pub(crate) fn new(
|
||||||
|
@ -434,7 +506,7 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<A, D, Output, T> BatchRunner<A, D, Output, T>
|
impl<A, D, Output, Dec, T> BatchRunner<A, D, Output, Dec, T>
|
||||||
where
|
where
|
||||||
A: Clone + Send + 'static,
|
A: Clone + Send + 'static,
|
||||||
D: BatchDomain + Send + 'static,
|
D: BatchDomain + Send + 'static,
|
||||||
|
@ -442,8 +514,9 @@ where
|
||||||
D::Memo: Send,
|
D::Memo: Send,
|
||||||
D::Note: Send,
|
D::Note: Send,
|
||||||
D::Recipient: Send,
|
D::Recipient: Send,
|
||||||
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone + Send + 'static,
|
Output: Clone + Send + 'static,
|
||||||
T: Tasks<Batch<A, D, Output>>,
|
Dec: Decryptor<D, Output>,
|
||||||
|
T: Tasks<Batch<A, D, Output, Dec>>,
|
||||||
{
|
{
|
||||||
/// Batches the given outputs for trial decryption.
|
/// Batches the given outputs for trial decryption.
|
||||||
///
|
///
|
||||||
|
@ -491,7 +564,7 @@ where
|
||||||
&mut self,
|
&mut self,
|
||||||
block_tag: BlockHash,
|
block_tag: BlockHash,
|
||||||
txid: TxId,
|
txid: TxId,
|
||||||
) -> HashMap<(TxId, usize), DecryptedOutput<A, D, ()>> {
|
) -> HashMap<(TxId, usize), DecryptedOutput<A, D, Dec::Memo>> {
|
||||||
self.pending_results
|
self.pending_results
|
||||||
.remove(&ResultKey(block_tag, txid))
|
.remove(&ResultKey(block_tag, txid))
|
||||||
// We won't have a pending result if the transaction didn't have outputs of
|
// We won't have a pending result if the transaction didn't have outputs of
|
||||||
|
|
|
@ -21,7 +21,7 @@ use zcash_primitives::{
|
||||||
use crate::data_api::{BlockMetadata, ScannedBlock, ScannedBundles};
|
use crate::data_api::{BlockMetadata, ScannedBlock, ScannedBundles};
|
||||||
use crate::{
|
use crate::{
|
||||||
proto::compact_formats::CompactBlock,
|
proto::compact_formats::CompactBlock,
|
||||||
scan::{Batch, BatchRunner, Tasks},
|
scan::{Batch, BatchRunner, CompactDecryptor, Tasks},
|
||||||
wallet::{WalletSaplingOutput, WalletSaplingSpend, WalletTx},
|
wallet::{WalletSaplingOutput, WalletSaplingSpend, WalletTx},
|
||||||
ShieldedProtocol,
|
ShieldedProtocol,
|
||||||
};
|
};
|
||||||
|
@ -268,9 +268,10 @@ pub fn scan_block<P: consensus::Parameters + Send + 'static, K: ScanningKey>(
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
type TaggedBatch<S> = Batch<(AccountId, S), SaplingDomain, CompactOutputDescription>;
|
type TaggedBatch<S> =
|
||||||
|
Batch<(AccountId, S), SaplingDomain, CompactOutputDescription, CompactDecryptor>;
|
||||||
type TaggedBatchRunner<S, T> =
|
type TaggedBatchRunner<S, T> =
|
||||||
BatchRunner<(AccountId, S), SaplingDomain, CompactOutputDescription, T>;
|
BatchRunner<(AccountId, S), SaplingDomain, CompactOutputDescription, CompactDecryptor, T>;
|
||||||
|
|
||||||
#[tracing::instrument(skip_all, fields(height = block.height))]
|
#[tracing::instrument(skip_all, fields(height = block.height))]
|
||||||
pub(crate) fn add_block_to_runner<P, S, T>(
|
pub(crate) fn add_block_to_runner<P, S, T>(
|
||||||
|
@ -834,7 +835,7 @@ mod tests {
|
||||||
assert_eq!(cb.vtx.len(), 2);
|
assert_eq!(cb.vtx.len(), 2);
|
||||||
|
|
||||||
let mut batch_runner = if scan_multithreaded {
|
let mut batch_runner = if scan_multithreaded {
|
||||||
let mut runner = BatchRunner::<_, _, _, ()>::new(
|
let mut runner = BatchRunner::<_, _, _, _, ()>::new(
|
||||||
10,
|
10,
|
||||||
dfvk.to_sapling_keys()
|
dfvk.to_sapling_keys()
|
||||||
.iter()
|
.iter()
|
||||||
|
@ -921,7 +922,7 @@ mod tests {
|
||||||
assert_eq!(cb.vtx.len(), 3);
|
assert_eq!(cb.vtx.len(), 3);
|
||||||
|
|
||||||
let mut batch_runner = if scan_multithreaded {
|
let mut batch_runner = if scan_multithreaded {
|
||||||
let mut runner = BatchRunner::<_, _, _, ()>::new(
|
let mut runner = BatchRunner::<_, _, _, _, ()>::new(
|
||||||
10,
|
10,
|
||||||
dfvk.to_sapling_keys()
|
dfvk.to_sapling_keys()
|
||||||
.iter()
|
.iter()
|
||||||
|
|
Loading…
Reference in New Issue