zcash_note_encryption: Add batched trial decryption APIs

Extracted from: 8a615c4393

src/batch.rs

@@ -0,0 +1,72 @@
+//! APIs for batch trial decryption.
+
+use std::iter;
+
+use crate::{
+    try_compact_note_decryption_inner, try_note_decryption_inner, Domain, EphemeralKeyBytes,
+    ShieldedOutput,
+};
+
+/// Trial decryption of a batch of notes with a set of recipients.
+///
+/// This is the batched version of [`zcash_note_encryption::try_note_decryption`].
+pub fn try_note_decryption<D: Domain, Output: ShieldedOutput<D>>(
+    ivks: &[D::IncomingViewingKey],
+    outputs: &[(D, Output)],
+) -> Vec<Option<(D::Note, D::Recipient, D::Memo)>> {
+    batch_note_decryption(ivks, outputs, try_note_decryption_inner)
+}
+
+/// Trial decryption of a batch of notes for light clients with a set of recipients.
+///
+/// This is the batched version of [`zcash_note_encryption::try_compact_note_decryption`].
+pub fn try_compact_note_decryption<D: Domain, Output: ShieldedOutput<D>>(
+    ivks: &[D::IncomingViewingKey],
+    outputs: &[(D, Output)],
+) -> Vec<Option<(D::Note, D::Recipient)>> {
+    batch_note_decryption(ivks, outputs, try_compact_note_decryption_inner)
+}
+
+fn batch_note_decryption<D: Domain, Output: ShieldedOutput<D>, F, FR>(
+    ivks: &[D::IncomingViewingKey],
+    outputs: &[(D, Output)],
+    decrypt_inner: F,
+) -> Vec<Option<FR>>
+where
+    F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, D::SymmetricKey) -> Option<FR>,
+{
+    // Fetch the ephemeral keys for each output.
+    let ephemeral_keys: Vec<_> = outputs
+        .iter()
+        .map(|(_, output)| output.ephemeral_key())
+        .collect();
+
+    // Derive the shared secrets for all combinations of (ivk, output).
+    // None of this work can benefit from batching.
+    let items = ivks.iter().flat_map(|ivk| {
+        ephemeral_keys.iter().map(move |ephemeral_key| {
+            (
+                D::epk(ephemeral_key).map(|epk| D::ka_agree_dec(ivk, &epk)),
+                ephemeral_key,
+            )
+        })
+    });
+
+    // Run the batch-KDF to obtain the symmetric keys from the shared secrets.
+    let keys = D::batch_kdf(items);
+
+    // Finish the trial decryption!
+    ivks.iter()
+        .flat_map(|ivk| {
+            // Reconstruct the matrix of (ivk, output) combinations.
+            iter::repeat(ivk)
+                .zip(ephemeral_keys.iter())
+                .zip(outputs.iter())
+        })
+        .zip(keys)
+        .map(|(((ivk, ephemeral_key), (domain, output)), key)| {
+            // The `and_then` propagates any potential rejection from `D::epk`.
+            key.and_then(|key| decrypt_inner(domain, ivk, ephemeral_key, output, key))
+        })
+        .collect()
+}

src/lib.rs

@@ -7,6 +7,8 @@ use crypto_api_chachapoly::{ChaCha20Ietf, ChachaPolyIetf};
 use rand_core::RngCore;
 use subtle::{Choice, ConstantTimeEq};
 
+pub mod batch;
+
 pub const COMPACT_NOTE_SIZE: usize = 1 + // version
     11 + // diversifier
     8  + // value
@@ -99,6 +101,19 @@ pub trait Domain {
 
     fn kdf(secret: Self::SharedSecret, ephemeral_key: &EphemeralKeyBytes) -> Self::SymmetricKey;
 
+    /// Computes `Self::kdf` on a batch of items.
+    ///
+    /// For each item in the batch, if the shared secret is `None`, this returns `None` at
+    /// that position.
+    fn batch_kdf<'a>(
+        items: impl Iterator<Item = (Option<Self::SharedSecret>, &'a EphemeralKeyBytes)>,
+    ) -> Vec<Option<Self::SymmetricKey>> {
+        // Default implementation: do the non-batched thing.
+        items
+            .map(|(secret, ephemeral_key)| secret.map(|secret| Self::kdf(secret, ephemeral_key)))
+            .collect()
+    }
+
     // for right now, we just need `recipient` to get `d`; in the future when we
     // can get that from a Sapling note, the recipient parameter will be able
     // to be removed.
@@ -334,13 +349,23 @@ pub fn try_note_decryption<D: Domain, Output: ShieldedOutput<D>>(
     ivk: &D::IncomingViewingKey,
     output: &Output,
 ) -> Option<(D::Note, D::Recipient, D::Memo)> {
-    assert_eq!(output.enc_ciphertext().len(), ENC_CIPHERTEXT_SIZE);
     let ephemeral_key = output.ephemeral_key();
 
     let epk = D::epk(&ephemeral_key)?;
     let shared_secret = D::ka_agree_dec(ivk, &epk);
     let key = D::kdf(shared_secret, &ephemeral_key);
 
+    try_note_decryption_inner(domain, ivk, &ephemeral_key, output, key)
+}
+
+fn try_note_decryption_inner<D: Domain, Output: ShieldedOutput<D>>(
+    domain: &D,
+    ivk: &D::IncomingViewingKey,
+    ephemeral_key: &EphemeralKeyBytes,
+    output: &Output,
+    key: D::SymmetricKey,
+) -> Option<(D::Note, D::Recipient, D::Memo)> {
+    assert_eq!(output.enc_ciphertext().len(), ENC_CIPHERTEXT_SIZE);
     let mut plaintext = [0; ENC_CIPHERTEXT_SIZE];
     assert_eq!(
         ChachaPolyIetf::aead_cipher()
@@ -358,7 +383,7 @@ pub fn try_note_decryption<D: Domain, Output: ShieldedOutput<D>>(
     let (note, to) = parse_note_plaintext_without_memo_ivk(
         domain,
         ivk,
-        &ephemeral_key,
+        ephemeral_key,
         &output.cmstar_bytes(),
         &plaintext,
     )?;
@@ -419,13 +444,24 @@ pub fn try_compact_note_decryption<D: Domain, Output: ShieldedOutput<D>>(
     ivk: &D::IncomingViewingKey,
     output: &Output,
 ) -> Option<(D::Note, D::Recipient)> {
-    assert_eq!(output.enc_ciphertext().len(), COMPACT_NOTE_SIZE);
     let ephemeral_key = output.ephemeral_key();
 
     let epk = D::epk(&ephemeral_key)?;
     let shared_secret = D::ka_agree_dec(&ivk, &epk);
     let key = D::kdf(shared_secret, &ephemeral_key);
 
+    try_compact_note_decryption_inner(domain, ivk, &ephemeral_key, output, key)
+}
+
+fn try_compact_note_decryption_inner<D: Domain, Output: ShieldedOutput<D>>(
+    domain: &D,
+    ivk: &D::IncomingViewingKey,
+    ephemeral_key: &EphemeralKeyBytes,
+    output: &Output,
+    key: D::SymmetricKey,
+) -> Option<(D::Note, D::Recipient)> {
+    assert_eq!(output.enc_ciphertext().len(), COMPACT_NOTE_SIZE);
+
     // Start from block 1 to skip over Poly1305 keying output
     let mut plaintext = [0; COMPACT_NOTE_SIZE];
     plaintext.copy_from_slice(output.enc_ciphertext());
@@ -434,7 +470,7 @@ pub fn try_compact_note_decryption<D: Domain, Output: ShieldedOutput<D>>(
     parse_note_plaintext_without_memo_ivk(
         domain,
         ivk,
-        &ephemeral_key,
+        ephemeral_key,
         &output.cmstar_bytes(),
         &plaintext,
     )