improve vess api to accept IntoIterator instead

Author: alxiong

timeboost-crypto/src/mre.rs

@@ -174,8 +174,8 @@ pub struct Ciphertext<C: CurveGroup, H: Digest = sha2::Sha256> {
 /// - `aad` is the associated data
 /// - `C` is the DL group, `H` is the choice of H_enc whose output space = message space
 ///   - preprocess messages to pad them to proper length before passing in
-pub fn encrypt<C, H, R>(
-    recipients: &[EncryptionKey<C>],
+pub fn encrypt<'a, C, H, R, I>(
+    recipients: I,
     messages: &[Vec<u8>],
     aad: &[u8],
     rng: &mut R,
@@ -184,14 +184,17 @@ where
     C: CurveGroup,
     H: Digest,
     R: Rng + CryptoRng,
+    I: IntoIterator<Item = &'a EncryptionKey<C>>,
+    I::IntoIter: ExactSizeIterator,
 {
     // input validation
-    if recipients.is_empty() || messages.is_empty() {
+    let recipients_iter = recipients.into_iter();
+    if messages.is_empty() {
         return Err(MultiRecvEncError::EmptyInput);
     }
-    if recipients.len() != messages.len() {
+    if recipients_iter.len() != messages.len() {
         return Err(MultiRecvEncError::MismatchedInputLength(
-            recipients.len(),
+            recipients_iter.len(),
             messages.len(),
         ));
     }
@@ -210,8 +213,7 @@ where
     let epk = C::generator().mul(&esk);
 
     // generate recipient-specific ciphertext parts
-    let cts = recipients
-        .iter()
+    let cts = recipients_iter
         .zip(messages.iter())
         .enumerate()
         .map(|(idx, (pk, msg))| {
@@ -228,7 +230,7 @@ where
 
             // TODO(alex): use SIMD vectorized XOR when `std::simd` move out of nightly,
             // or rayon as an intermediate improvement
-            let ct = Output::<H>::from_iter(k.iter().zip(msg).map(|(ki, m)| ki ^ m));
+            let ct = Output::<H>::from_iter(k.iter().zip(msg.iter()).map(|(ki, m)| ki ^ m));
             Ok(ct)
         })
         .collect::<Result<Vec<_>, MultiRecvEncError>>()?;
@@ -276,7 +278,7 @@ impl From<ark_serialize::SerializationError> for MultiRecvEncError {
 
 #[cfg(test)]
 mod tests {
-    use std::iter::repeat_with;
+    use std::{collections::BTreeMap, iter::repeat_with};
 
     use ark_bls12_381::G1Projective;
     use ark_std::rand;
@@ -290,8 +292,12 @@ mod tests {
         let n = 10; // num of recipients
         let recv_sks: Vec<DecryptionKey<G1Projective>> =
             repeat_with(|| DecryptionKey::rand(rng)).take(n).collect();
-        let recv_pks: Vec<EncryptionKey<G1Projective>> =
-            recv_sks.iter().map(EncryptionKey::from).collect();
+        // collecting into a BTreeSet to demonstrate flexible encrypt() input type
+        let recv_pks: BTreeMap<usize, EncryptionKey<G1Projective>> = recv_sks
+            .iter()
+            .enumerate()
+            .map(|(i, sk)| (i, EncryptionKey::from(sk)))
+            .collect();
         let labeled_sks: Vec<LabeledDecryptionKey<G1Projective>> = recv_sks
             .into_iter()
             .enumerate()
@@ -302,7 +308,7 @@ mod tests {
             .collect::<Vec<_>>();
         let aad = b"Alice";
 
-        let mre_ct = encrypt::<G1Projective, H, _>(&recv_pks, &msgs, aad, rng).unwrap();
+        let mre_ct = encrypt::<G1Projective, H, _, _>(recv_pks.values(), &msgs, aad, rng).unwrap();
         for i in 0..n {
             let ct = mre_ct.get_recipient_ct(i).unwrap();
             assert_eq!(

timeboost-crypto/src/vess.rs

@@ -209,11 +209,11 @@ impl<C: CurveGroup> ShoupVess<C> {
 
     // deterministically generate the `i`-th dealing from a random seed
     // each dealing contains (Shamir poly + Feldman commitment + MRE ciphertext)
-    fn new_dealing(
+    fn new_dealing<'a, I>(
         &self,
         ith: usize,
         seed: &[u8; 32],
-        recipients: &[mre::EncryptionKey<C>],
+        recipients: I,
         aad: &[u8],
     ) -> Result<
         (
@@ -222,7 +222,11 @@ impl<C: CurveGroup> ShoupVess<C> {
             MultiRecvCiphertext<C, sha2::Sha256>,
         ),
         VessError,
-    > {
+    >
+    where
+        I: IntoIterator<Item = &'a mre::EncryptionKey<C>>,
+        I::IntoIter: ExactSizeIterator,
+    {
         let mut rng = ChaCha20Rng::from_seed(*seed);
         let vss_secret = C::ScalarField::rand(&mut rng);
 
@@ -231,7 +235,7 @@ impl<C: CurveGroup> ShoupVess<C> {
         let serialized_shares: Vec<Vec<u8>> =
             FeldmanVss::<C>::compute_serialized_shares(&self.vss_pp, &poly).collect();
 
-        let mre_ct = mre::encrypt::<C, sha2::Sha256, _>(
+        let mre_ct = mre::encrypt::<C, sha2::Sha256, _, _>(
             recipients,
             &serialized_shares,
             &self.indexed_aad(aad, ith),
@@ -258,9 +262,9 @@ impl<C: CurveGroup> ShoupVess<C> {
     ///   step 1.a is split as two internal steps in the two APIs above. r_k is 32 bytes and
     ///   SpongeFish's built-in prover private coin toss.
     /// - random subset seed s: see [`Self::map_subset_seed()`]
-    pub fn encrypted_shares(
+    pub fn encrypted_shares<'a, I>(
         &self,
-        recipients: &[mre::EncryptionKey<C>],
+        recipients: I,
         secret: C::ScalarField,
         aad: &[u8],
     ) -> Result<
@@ -269,11 +273,16 @@ impl<C: CurveGroup> ShoupVess<C> {
             <FeldmanVss<C> as VerifiableSecretSharing>::Commitment,
         ),
         VessError,
-    > {
-        // input validation
+    >
+    where
+        I: IntoIterator<Item = &'a mre::EncryptionKey<C>>,
+        I::IntoIter: ExactSizeIterator + Clone + Sync,
+    {
+        // input validation - check length without consuming the iterator
+        let recipients_iter = recipients.into_iter();
         let n = self.vss_pp.n.get();
-        if recipients.len() != n {
-            return Err(VessError::WrongRecipientsLength(n, recipients.len()));
+        if recipients_iter.len() != n {
+            return Err(VessError::WrongRecipientsLength(n, recipients_iter.len()));
         }
 
         let mut prover_state = self.io_pattern(aad).to_prover_state();
@@ -293,7 +302,7 @@ impl<C: CurveGroup> ShoupVess<C> {
         )> = seeds
             .par_iter()
             .enumerate()
-            .map(|(i, r)| self.new_dealing(i, r, recipients, aad))
+            .map(|(i, r)| self.new_dealing(i, r, recipients_iter.clone(), aad))
             .collect::<Result<_, VessError>>()?;
 
         // compute h:= H_compress(aad, dealings)
@@ -392,13 +401,17 @@ impl<C: CurveGroup> ShoupVess<C> {
 
     /// Verify if the ciphertext (for all recipients) correctly encrypting valid secret shares,
     /// verifiable by anyone.
-    pub fn verify(
+    pub fn verify<'a, I>(
         &self,
-        recipients: &[mre::EncryptionKey<C>],
+        recipients: I,
         ct: &VessCiphertext,
         comm: &<FeldmanVss<C> as VerifiableSecretSharing>::Commitment,
         aad: &[u8],
-    ) -> Result<(), VessError> {
+    ) -> Result<(), VessError>
+    where
+        I: IntoIterator<Item = &'a mre::EncryptionKey<C>> + Clone,
+        I::IntoIter: ExactSizeIterator,
+    {
         let mut verifier_state = self.io_pattern(aad).to_verifier_state(&ct.transcript);
 
         // verifier logic until Step 4b
@@ -454,7 +467,8 @@ impl<C: CurveGroup> ShoupVess<C> {
                     let seed = seeds
                         .pop_front()
                         .expect("subset_size < num_repetitions, so seeds.len() > 0");
-                    let (_poly, cm, mre_ct) = self.new_dealing(i, &seed, recipients, aad)?;
+                    let (_poly, cm, mre_ct) =
+                        self.new_dealing(i, &seed, recipients.clone(), aad)?;
 
                     hasher.update(serialize_to_vec![cm]?);
                     hasher.update(mre_ct.to_bytes());
@@ -703,7 +717,11 @@ mod tests {
         UniformRand,
         rand::{SeedableRng, rngs::StdRng},
     };
-    use std::{collections::BTreeSet, iter::repeat_with, num::NonZeroUsize};
+    use std::{
+        collections::{BTreeMap, BTreeSet},
+        iter::repeat_with,
+        num::NonZeroUsize,
+    };
 
     type H = sha2::Sha256;
     type Vss = FeldmanVss<G1Projective>;
@@ -717,18 +735,23 @@ mod tests {
             repeat_with(|| mre::DecryptionKey::rand(rng))
                 .take(n)
                 .collect();
-        let recv_pks: Vec<mre::EncryptionKey<G1Projective>> =
-            recv_sks.iter().map(mre::EncryptionKey::from).collect();
+        let recv_pks: BTreeMap<usize, mre::EncryptionKey<G1Projective>> = recv_sks
+            .iter()
+            .enumerate()
+            .map(|(i, sk)| (i, mre::EncryptionKey::from(sk)))
+            .collect();
         let labeled_sks: Vec<LabeledDecryptionKey<G1Projective>> = recv_sks
             .into_iter()
             .enumerate()
             .map(|(i, sk)| sk.label(i))
             .collect();
 
         let aad = b"Associated data";
-        let (ct, comm) = vess.encrypted_shares(&recv_pks, secret, aad).unwrap();
+        let (ct, comm) = vess
+            .encrypted_shares(recv_pks.values(), secret, aad)
+            .unwrap();
 
-        assert!(vess.verify(&recv_pks, &ct, &comm, aad).is_ok());
+        assert!(vess.verify(recv_pks.values(), &ct, &comm, aad).is_ok());
         for labeled_recv_sk in labeled_sks {
             let share = vess.decrypt_share(&labeled_recv_sk, &ct, aad).unwrap();
             assert!(Vss::verify(&vess.vss_pp, labeled_recv_sk.node_idx, &share, &comm).is_ok());