Restructure batch decryption to avoid cartesian product of results.

While it is necessary in the worst case to perform `m * n` decryptions,
where `m` is the number of outputs being decrypted and `n` is the number
of IVKs, it is possible to stop performing trial decryptions when the
first successful decryption is performed. Also, it's inconvenient and
unnecessary to return the full cartesian product of these results, as
only one IVK will decrypt a given output. This commit modifies batch
trial decryption to stop on the first successful decryption, and instead
of returning the cartesian product of results we return the index of the
input IVK along with the output it decrypted. Note that this means that
trial decryption is not constant-time with respect to the number and/or
order of IVKs.


Extracted from: 5873950648

CHANGELOG.md

@@ -7,6 +7,15 @@ and this library adheres to Rust's notion of
 
 ## [Unreleased]
 
+- Changes to batch decryption APIs:
+  - The return types of `batch::try_note_decryption` and
+    `batch::try_compact_note_decryption` have changed. Now, instead of
+    returning entries corresponding to the cartesian product of the IVKs used for
+    decryption with the outputs being decrypted, this now returns a vector of
+    decryption results of the same length and in the same order as the `outputs`
+    argument to the function. Each successful result includes the index of the
+    entry in `ivks` used to decrypt the value.
+
 ### Changed
 - MSRV is now 1.56.1.
 

src/batch.rs

@@ -1,7 +1,6 @@
 //! APIs for batch trial decryption.
 
 use alloc::vec::Vec; // module is alloc only
-use core::iter;
 
 use crate::{
     try_compact_note_decryption_inner, try_note_decryption_inner, BatchDomain, EphemeralKeyBytes,
@@ -11,21 +10,32 @@ use crate::{
 /// Trial decryption of a batch of notes with a set of recipients.
 ///
 /// This is the batched version of [`crate::try_note_decryption`].
+///
+/// Returns a vector containing the decrypted result for each output,
+/// with the same length and in the same order as the outputs were
+/// provided, along with the index in the `ivks` slice associated with
+/// the IVK that successfully decrypted the output.
 #[allow(clippy::type_complexity)]
 pub fn try_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, ENC_CIPHERTEXT_SIZE>>(
     ivks: &[D::IncomingViewingKey],
     outputs: &[(D, Output)],
-) -> Vec<Option<(D::Note, D::Recipient, D::Memo)>> {
+) -> Vec<Option<((D::Note, D::Recipient, D::Memo), usize)>> {
     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 [`crate::try_compact_note_decryption`].
+///
+/// Returns a vector containing the decrypted result for each output,
+/// with the same length and in the same order as the outputs were
+/// provided, along with the index in the `ivks` slice associated with
+/// the IVK that successfully decrypted the output.
+#[allow(clippy::type_complexity)]
 pub fn try_compact_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>>(
     ivks: &[D::IncomingViewingKey],
     outputs: &[(D, Output)],
-) -> Vec<Option<(D::Note, D::Recipient)>> {
+) -> Vec<Option<((D::Note, D::Recipient), usize)>> {
     batch_note_decryption(ivks, outputs, try_compact_note_decryption_inner)
 }
 
@@ -33,17 +43,17 @@ fn batch_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, CS>, F, FR, c
     ivks: &[D::IncomingViewingKey],
     outputs: &[(D, Output)],
     decrypt_inner: F,
-) -> Vec<Option<FR>>
+) -> Vec<Option<(FR, usize)>>
 where
-    F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, D::SymmetricKey) -> Option<FR>,
+    F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, &D::SymmetricKey) -> Option<FR>,
 {
     // Fetch the ephemeral keys for each output and batch-parse them.
     let ephemeral_keys = D::batch_epk(outputs.iter().map(|(_, output)| output.ephemeral_key()));
 
     // Derive the shared secrets for all combinations of (ivk, output).
     // The scalar multiplications cannot benefit from batching.
-    let items = ivks.iter().flat_map(|ivk| {
-        ephemeral_keys.iter().map(move |(epk, ephemeral_key)| {
+    let items = ephemeral_keys.iter().flat_map(|(epk, ephemeral_key)| {
+        ivks.iter().map(move |ivk| {
             (
                 epk.as_ref().map(|epk| D::ka_agree_dec(ivk, epk)),
                 ephemeral_key,
@@ -55,17 +65,18 @@ where
     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())
+    keys.chunks(ivks.len())
+        .zip(ephemeral_keys.iter().zip(outputs.iter()))
+        .map(|(key_chunk, ((_, ephemeral_key), (domain, output)))| {
+            key_chunk
+                .iter()
+                .zip(ivks.iter().enumerate())
+                .filter_map(|(key, (i, ivk))| {
+                    key.as_ref()
+                        .and_then(|key| decrypt_inner(domain, ivk, ephemeral_key, output, key))
+                        .map(|out| (out, i))
+                })
+                .next()
         })
-        .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()
+        .collect::<Vec<Option<_>>>()
 }

src/lib.rs

@@ -520,7 +520,7 @@ pub fn try_note_decryption<D: Domain, Output: ShieldedOutput<D, ENC_CIPHERTEXT_S
     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)
+    try_note_decryption_inner(domain, ivk, &ephemeral_key, output, &key)
 }
 
 fn try_note_decryption_inner<D: Domain, Output: ShieldedOutput<D, ENC_CIPHERTEXT_SIZE>>(
@@ -528,7 +528,7 @@ fn try_note_decryption_inner<D: Domain, Output: ShieldedOutput<D, ENC_CIPHERTEXT
     ivk: &D::IncomingViewingKey,
     ephemeral_key: &EphemeralKeyBytes,
     output: &Output,
-    key: D::SymmetricKey,
+    key: &D::SymmetricKey,
 ) -> Option<(D::Note, D::Recipient, D::Memo)> {
     let enc_ciphertext = output.enc_ciphertext();
 
@@ -617,7 +617,7 @@ pub fn try_compact_note_decryption<D: Domain, Output: ShieldedOutput<D, COMPACT_
     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)
+    try_compact_note_decryption_inner(domain, ivk, &ephemeral_key, output, &key)
 }
 
 fn try_compact_note_decryption_inner<D: Domain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>>(
@@ -625,7 +625,7 @@ fn try_compact_note_decryption_inner<D: Domain, Output: ShieldedOutput<D, COMPAC
     ivk: &D::IncomingViewingKey,
     ephemeral_key: &EphemeralKeyBytes,
     output: &Output,
-    key: D::SymmetricKey,
+    key: &D::SymmetricKey,
 ) -> Option<(D::Note, D::Recipient)> {
     // Start from block 1 to skip over Poly1305 keying output
     let mut plaintext = [0; COMPACT_NOTE_SIZE];