masp_note_encryption/
batch.rs

1//! APIs for batch trial decryption.
2
3use alloc::vec::Vec; // module is alloc only
4
5use crate::{
6    try_compact_note_decryption_inner, try_note_decryption_inner, BatchDomain, EphemeralKeyBytes,
7    ShieldedOutput, COMPACT_NOTE_SIZE, ENC_CIPHERTEXT_SIZE,
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 recipients.
27///
28/// This is the batched version of [`crate::try_compact_note_decryption`].
29///
30/// Returns a vector containing the decrypted result for each output,
31/// with the same length and in the same order as the outputs were
32/// provided, along with the index in the `ivks` slice associated with
33/// the IVK that successfully decrypted the output.
34#[allow(clippy::type_complexity)]
35pub fn try_compact_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>>(
36    ivks: &[D::IncomingViewingKey],
37    outputs: &[(D, Output)],
38) -> Vec<Option<((D::Note, D::Recipient), usize)>> {
39    batch_note_decryption(ivks, outputs, try_compact_note_decryption_inner)
40}
41
42fn batch_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, CS>, F, FR, const CS: usize>(
43    ivks: &[D::IncomingViewingKey],
44    outputs: &[(D, Output)],
45    decrypt_inner: F,
46) -> Vec<Option<(FR, usize)>>
47where
48    F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, &D::SymmetricKey) -> Option<FR>,
49{
50    if ivks.is_empty() {
51        return (0..outputs.len()).map(|_| None).collect();
52    };
53
54    // Fetch the ephemeral keys for each output, and batch-parse and prepare them.
55    let ephemeral_keys = D::batch_epk(outputs.iter().map(|(_, output)| output.ephemeral_key()));
56
57    // Derive the shared secrets for all combinations of (ivk, output).
58    // The scalar multiplications cannot benefit from batching.
59    let items = ephemeral_keys.iter().flat_map(|(epk, ephemeral_key)| {
60        ivks.iter().map(move |ivk| {
61            (
62                epk.as_ref().map(|epk| D::ka_agree_dec(ivk, epk)),
63                ephemeral_key,
64            )
65        })
66    });
67
68    // Run the batch-KDF to obtain the symmetric keys from the shared secrets.
69    let keys = D::batch_kdf(items);
70
71    // Finish the trial decryption!
72    keys.chunks(ivks.len())
73        .zip(ephemeral_keys.iter().zip(outputs.iter()))
74        .map(|(key_chunk, ((_, ephemeral_key), (domain, output)))| {
75            key_chunk
76                .iter()
77                .zip(ivks.iter().enumerate())
78                .filter_map(|(key, (i, ivk))| {
79                    key.as_ref()
80                        .and_then(|key| decrypt_inner(domain, ivk, ephemeral_key, output, key))
81                        .map(|out| (out, i))
82                })
83                .next()
84        })
85        .collect::<Vec<Option<_>>>()
86}
masp_note_encryption/batch.rs

masp_note_encryption/
batch.rs
