From 2b4a88be363cbe6f0a9ccbb61a123d2ba6b98019 Mon Sep 17 00:00:00 2001 From: Jack Grigg Date: Fri, 6 Aug 2021 16:20:42 +0100 Subject: [PATCH] zcash_note_encryption: Add batched trial decryption APIs Extracted from: https://github.com/zcash/librustzcash/commit/8a615c43938a5d8237835cd686247cebfa7ae0ea --- src/batch.rs | 72 ++++++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 44 +++++++++++++++++++++++++++++--- 2 files changed, 112 insertions(+), 4 deletions(-) create mode 100644 src/batch.rs diff --git a/src/batch.rs b/src/batch.rs new file mode 100644 index 0000000..6eabc7f --- /dev/null +++ b/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>( + ivks: &[D::IncomingViewingKey], + outputs: &[(D, Output)], +) -> Vec> { + 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>( + ivks: &[D::IncomingViewingKey], + outputs: &[(D, Output)], +) -> Vec> { + batch_note_decryption(ivks, outputs, try_compact_note_decryption_inner) +} + +fn batch_note_decryption, F, FR>( + ivks: &[D::IncomingViewingKey], + outputs: &[(D, Output)], + decrypt_inner: F, +) -> Vec> +where + F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, D::SymmetricKey) -> Option, +{ + // 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() +} diff --git a/src/lib.rs b/src/lib.rs index a18964b..af784c3 100644 --- a/src/lib.rs +++ b/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, &'a EphemeralKeyBytes)>, + ) -> Vec> { + // 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>( 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>( + 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>( 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>( 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>( + 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>( parse_note_plaintext_without_memo_ivk( domain, ivk, - &ephemeral_key, + ephemeral_key, &output.cmstar_bytes(), &plaintext, )