DKG Round 2: return the public packages into a new CombinedPublicPackage struct

This struct is intended to group public packages coming from the same
sender and has a more compact serialization than the concatenation of
the serializations of multiple `PublicPackage` structs.

Author: andiflabs

src/dkg/round2.rs

@@ -3,7 +3,6 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
 use crate::checksum::Checksum;
-
 use crate::checksum::ChecksumError;
 use crate::checksum::ChecksumHasher;
 use crate::checksum::CHECKSUM_LEN;
@@ -30,7 +29,6 @@ use crate::serde::write_variable_length_bytes;
 use rand_core::CryptoRng;
 use rand_core::RngCore;
 use std::borrow::Borrow;
-
 use std::collections::BTreeMap;
 use std::hash::Hasher;
 use std::io;
@@ -229,15 +227,25 @@ impl PublicPackage {
 
     pub fn serialize_into<W: io::Write>(&self, mut writer: W) -> io::Result<()> {
         self.sender_identity.serialize_into(&mut writer)?;
+        self.serialize_without_sender_into(writer)
+    }
+
+    fn serialize_without_sender_into<W: io::Write>(&self, mut writer: W) -> io::Result<()> {
         self.recipient_identity.serialize_into(&mut writer)?;
         let frost_package = self.frost_package.serialize().map_err(io::Error::other)?;
         write_variable_length_bytes(&mut writer, &frost_package)?;
-        writer.write_all(&self.checksum.to_le_bytes())?;
-        Ok(())
+        writer.write_all(&self.checksum.to_le_bytes())
     }
 
     pub fn deserialize_from<R: io::Read>(mut reader: R) -> io::Result<Self> {
         let sender_identity = Identity::deserialize_from(&mut reader)?;
+        Self::deserialize_without_sender_from(reader, sender_identity)
+    }
+
+    fn deserialize_without_sender_from<R: io::Read>(
+        mut reader: R,
+        sender_identity: Identity,
+    ) -> io::Result<Self> {
         let recipient_identity = Identity::deserialize_from(&mut reader)?;
 
         let frost_package = read_variable_length_bytes(&mut reader)?;
@@ -256,12 +264,85 @@ impl PublicPackage {
     }
 }
 
+/// A collection of [`PublicPackage`] structs, all from the same sender.
+#[derive(Default, Clone, PartialEq, Eq, Debug)]
+pub struct CombinedPublicPackage {
+    packages: Vec<PublicPackage>,
+}
+
+impl CombinedPublicPackage {
+    // This struct should not be constructed directly, hence `new` does not need to be `pub`.
+    // Keeping `new` private has the advantage that the implementation does not need to strictly
+    // enforce the same `sender_identity`, but it can omit this check (here we still check in debug
+    // builds just to catch bugs).
+    fn new(packages: Vec<PublicPackage>) -> Self {
+        // The serialization expects at least 1 package to be present
+        debug_assert!(!packages.is_empty());
+
+        let first_identity = &packages[0].sender_identity;
+        for pkg in &packages {
+            debug_assert_eq!(&pkg.sender_identity, first_identity);
+        }
+
+        Self { packages }
+    }
+
+    #[inline]
+    #[must_use]
+    pub fn packages(&self) -> &[PublicPackage] {
+        &self.packages
+    }
+
+    #[inline]
+    pub fn packages_for<'a>(
+        &'a self,
+        recipient_identity: &'a Identity,
+    ) -> impl Iterator<Item = &'a PublicPackage> + 'a {
+        self.packages
+            .iter()
+            .filter(move |pkg| &pkg.recipient_identity == recipient_identity)
+    }
+
+    pub fn serialize(&self) -> Vec<u8> {
+        let mut buf = Vec::new();
+        self.serialize_into(&mut buf).expect("serialization failed");
+        buf
+    }
+
+    pub fn serialize_into<W: io::Write>(&self, mut writer: W) -> io::Result<()> {
+        let sender_identity = &self.packages[0].sender_identity;
+        sender_identity.serialize_into(&mut writer)?;
+        write_variable_length(writer, &self.packages, |writer, pkg| {
+            pkg.serialize_without_sender_into(writer)
+        })
+    }
+
+    pub fn deserialize_from<R: io::Read>(mut reader: R) -> io::Result<Self> {
+        let sender_identity = Identity::deserialize_from(&mut reader)?;
+
+        let packages = read_variable_length(reader, move |reader| {
+            PublicPackage::deserialize_without_sender_from(reader, sender_identity.clone())
+        })?;
+
+        Ok(Self { packages })
+    }
+}
+
+impl IntoIterator for CombinedPublicPackage {
+    type Item = PublicPackage;
+    type IntoIter = <Vec<PublicPackage> as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.packages.into_iter()
+    }
+}
+
 pub fn round2<'a, P, R>(
     secret: &participant::Secret,
     round1_secret_package: &[u8],
     round1_public_packages: P,
     mut csrng: R,
-) -> Result<(Vec<u8>, Vec<PublicPackage>), Error>
+) -> Result<(Vec<u8>, CombinedPublicPackage), Error>
 where
     P: IntoIterator<Item = &'a round1::PublicPackage>,
     R: RngCore + CryptoRng,
@@ -354,7 +435,10 @@ where
         round2_public_packages.push(public_package);
     }
 
-    Ok((encrypted_secret_package, round2_public_packages))
+    Ok((
+        encrypted_secret_package,
+        CombinedPublicPackage::new(round2_public_packages),
+    ))
 }
 
 #[cfg(test)]
@@ -558,12 +642,12 @@ mod tests {
             .expect("round 2 secret package import failed");
 
         round2_public_packages
-            .iter()
-            .find(|&p| p.recipient_identity() == &identity2)
+            .packages_for(&identity2)
+            .next()
             .expect("round 2 public packages missing package for identity2");
         round2_public_packages
-            .iter()
-            .find(|&p| p.recipient_identity() == &identity3)
+            .packages_for(&identity3)
+            .next()
             .expect("round 2 public packages missing package for identity3");
     }
 

src/dkg/round3.rs

@@ -24,13 +24,16 @@ pub fn round3<'a, P, Q>(
 ) -> Result<(KeyPackage, PublicKeyPackage, GroupSecretKey), Error>
 where
     P: IntoIterator<Item = &'a round1::PublicPackage>,
-    Q: IntoIterator<Item = &'a round2::PublicPackage>,
+    Q: IntoIterator<Item = &'a round2::CombinedPublicPackage>,
 {
     let identity = secret.to_identity();
     let round2_secret_package =
         import_secret_package(round2_secret_package, secret).map_err(Error::DecryptionError)?;
     let round1_public_packages = round1_public_packages.into_iter().collect::<Vec<_>>();
-    let round2_public_packages = round2_public_packages.into_iter().collect::<Vec<_>>();
+    let round2_public_packages = round2_public_packages
+        .into_iter()
+        .flat_map(|combo| combo.packages_for(&identity))
+        .collect::<Vec<_>>();
 
     let (min_signers, max_signers) = round2::get_secret_package_signers(&round2_secret_package);
 
@@ -76,7 +79,7 @@ where
             .is_some()
         {
             return Err(Error::InvalidInput(format!(
-                "multiple public packages provided for identity {}",
+                "multiple round 1 public packages provided for identity {}",
                 public_package.identity()
             )));
         }
@@ -95,7 +98,9 @@ where
     // inputs
     round1_frost_packages
         .remove(&identity.to_frost_identifier())
-        .expect("missing public package for identity");
+        .ok_or_else(|| {
+            Error::InvalidInput("missing round 1 public package for own identity".to_string())
+        })?;
 
     let expected_round2_checksum =
         round2::input_checksum(round1_public_packages.iter().map(Borrow::borrow));
@@ -108,7 +113,7 @@ where
 
         if !identity.eq(public_package.recipient_identity()) {
             return Err(Error::InvalidInput(format!(
-                "public package does not have the correct recipient identity {:?}",
+                "round 2 public package does not have the correct recipient identity {:?}",
                 public_package.recipient_identity().serialize()
             )));
         }
@@ -121,7 +126,7 @@ where
             .is_some()
         {
             return Err(Error::InvalidInput(format!(
-                "multiple public packages provided for identity {}",
+                "multiple round 2 public packages provided for identity {}",
                 public_package.sender_identity()
             )));
         }
@@ -176,24 +181,19 @@ mod tests {
         )
         .expect("round 2 failed");
 
-        let (_, round2_public_packages_2) = round2::round2(
+        let (_, round2_public_packages) = round2::round2(
             &secret2,
             &round1_secret_package_2,
             [&package1, &package2],
             thread_rng(),
         )
         .expect("round 2 failed");
 
-        let round2_public_packages = [round2_public_packages_2
-            .iter()
-            .find(|p| p.recipient_identity().eq(&identity1))
-            .expect("should have package for identity1")];
-
         let result = round3(
             &secret1,
             &encrypted_secret_package,
             [&package2],
-            round2_public_packages,
+            [&round2_public_packages],
         );
 
         match result {
@@ -225,24 +225,19 @@ mod tests {
         )
         .expect("round 2 failed");
 
-        let (_, round2_public_packages_2) = round2::round2(
+        let (_, round2_public_packages) = round2::round2(
             &secret2,
             &round1_secret_package_2,
             [&package1, &package2],
             thread_rng(),
         )
         .expect("round 2 failed");
 
-        let round2_public_packages = [round2_public_packages_2
-            .iter()
-            .find(|p| p.recipient_identity().eq(&identity1))
-            .expect("should have package for identity1")];
-
         let result = round3(
             &secret1,
             &encrypted_secret_package,
             [&package1, &package1],
-            round2_public_packages,
+            [&round2_public_packages],
         );
 
         match result {
@@ -308,22 +303,11 @@ mod tests {
         )
         .expect("round 2 failed");
 
-        let round2_public_packages = [
-            round2_public_packages_2
-                .iter()
-                .find(|p| p.recipient_identity().eq(&identity1))
-                .expect("should have package for identity1"),
-            round2_public_packages_3
-                .iter()
-                .find(|p| p.recipient_identity().eq(&identity1))
-                .expect("should have package for identity1"),
-        ];
-
         round3(
             &secret1,
             &encrypted_secret_package,
             [&package1, &package2, &package3],
-            round2_public_packages,
+            [&round2_public_packages_2, &round2_public_packages_3],
         )
         .expect("round 3 failed");
     }