zcash_note_encryption/
batch.rs

1//! APIs for batch trial decryption.
2
3use alloc::vec::Vec; // module is alloc only
4
5use crate::{
6    BatchDomain, COMPACT_NOTE_SIZE, ENC_CIPHERTEXT_SIZE, EphemeralKeyBytes, ShieldedOutput,
7    try_compact_note_decryption_inner, try_note_decryption_inner,
8};
9
10/// Trial decryption of a batch of notes with a set of recipients.
11///
12/// This is the batched version of [`crate::try_note_decryption`].
13///
14/// Returns a vector containing the decrypted result for each output,
15/// with the same length and in the same order as the outputs were
16/// provided, along with the index in the `ivks` slice associated with
17/// the IVK that successfully decrypted the output.
18#[allow(clippy::type_complexity)]
19pub fn try_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, ENC_CIPHERTEXT_SIZE>>(
20    ivks: &[D::IncomingViewingKey],
21    outputs: &[(D, Output)],
22) -> Vec<Option<((D::Note, D::Recipient, D::Memo), usize)>> {
23    batch_note_decryption(ivks, outputs, try_note_decryption_inner)
24}
25
26/// Trial decryption of a batch of notes for light clients with a set of
27/// recipients.
28///
29/// This is the batched version of [`crate::try_compact_note_decryption`].
30///
31/// Returns a vector containing the decrypted result for each output,
32/// with the same length and in the same order as the outputs were
33/// provided, along with the index in the `ivks` slice associated with
34/// the IVK that successfully decrypted the output.
35#[allow(clippy::type_complexity)]
36pub fn try_compact_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>>(
37    ivks: &[D::IncomingViewingKey],
38    outputs: &[(D, Output)],
39) -> Vec<Option<((D::Note, D::Recipient), usize)>> {
40    batch_note_decryption(ivks, outputs, try_compact_note_decryption_inner)
41}
42
43fn batch_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, CS>, F, FR, const CS: usize>(
44    ivks: &[D::IncomingViewingKey],
45    outputs: &[(D, Output)],
46    decrypt_inner: F,
47) -> Vec<Option<(FR, usize)>>
48where
49    F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, &D::SymmetricKey) -> Option<FR>,
50{
51    if ivks.is_empty() {
52        return (0..outputs.len()).map(|_| None).collect();
53    };
54
55    // Fetch the ephemeral keys for each output, and batch-parse and prepare them.
56    let ephemeral_keys = D::batch_epk(outputs.iter().map(|(_, output)| output.ephemeral_key()));
57
58    // Derive the shared secrets for all combinations of (ivk, output), one batched
59    // same-key agreement per ivk: domains for which same-scalar multiplications can
60    // share work accelerate here, and the default `batch_ka_agree_dec` implementation
61    // is exactly the previous per-item computation.
62    // Reassembly below is in the (output-major, ivk-minor) order the batch-KDF
63    // expects, moving values out of the per-ivk columns (`SharedSecret` need not be
64    // `Clone`).
65    let mut columns: Vec<_> = ivks
66        .iter()
67        .map(|ivk| {
68            D::batch_ka_agree_dec(ivk, ephemeral_keys.iter().map(|(epk, _)| epk.as_ref()))
69                .into_iter()
70        })
71        .collect();
72    let mut secrets: Vec<Option<D::SharedSecret>> =
73        Vec::with_capacity(ephemeral_keys.len() * ivks.len());
74    for _ in 0..ephemeral_keys.len() {
75        for column in columns.iter_mut() {
76            secrets.push(
77                column
78                    .next()
79                    .expect("all columns have one entry per output"),
80            );
81        }
82    }
83    // Preserve the upstream iterator code during the ownership migration.
84    #[allow(clippy::manual_repeat_n)]
85    let items = secrets.into_iter().zip(
86        ephemeral_keys
87            .iter()
88            .flat_map(|(_, ephemeral_key)| core::iter::repeat(ephemeral_key).take(ivks.len())),
89    );
90
91    // Run the batch-KDF to obtain the symmetric keys from the shared secrets.
92    let keys = D::batch_kdf(items);
93
94    // Finish the trial decryption!
95    keys.chunks(ivks.len())
96        .zip(ephemeral_keys.iter().zip(outputs.iter()))
97        .map(|(key_chunk, ((_, ephemeral_key), (domain, output)))| {
98            key_chunk
99                .iter()
100                .zip(ivks.iter().enumerate())
101                .find_map(|(key, (i, ivk))| {
102                    key.as_ref()
103                        .and_then(|key| decrypt_inner(domain, ivk, ephemeral_key, output, key))
104                        .map(|out| (out, i))
105                })
106        })
107        .collect::<Vec<Option<_>>>()
108}
zcash_note_encryption/batch.rs

zcash_note_encryption/
batch.rs
