Merge branch 'ax/reshare-core' into ax/reshare-vess

alxiong · alxiong · commit 2987525e930c · 2025-08-05T15:43:17.000+08:00
diff --git a/timeboost-crypto/src/feldman.rs b/timeboost-crypto/src/feldman.rs
@@ -10,7 +10,7 @@ use multisig::Committee;
 use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
-use std::{iter::successors, num::NonZeroUsize};
+use std::{iter::successors, num::NonZeroUsize, ops::Add};
 
 use crate::{
     interpolation::{interpolate, interpolate_in_exponent},
@@ -158,27 +158,26 @@ impl<C: CurveGroup> VerifiableSecretSharing for FeldmanVss<C> {
 
     fn reconstruct(
         pp: &Self::PublicParam,
-        shares: impl Iterator<Item = (usize, Self::SecretShare)>,
+        shares: impl ExactSizeIterator<Item = (usize, Self::SecretShare)> + Clone,
     ) -> Result<Self::Secret, VssError> {
-        let shares = shares.collect::<Vec<_>>();
         let n = pp.n.get();
         let t = pp.t.get();
         // input validation
         if shares.len() != t {
             return Err(VssError::MismatchedSharesCount(t, shares.len()));
         }
-        for (idx, _) in shares.iter() {
-            if *idx >= n {
-                return Err(VssError::IndexOutOfBound(n - 1, *idx));
+        for (idx, _) in shares.clone() {
+            if idx >= n {
+                return Err(VssError::IndexOutOfBound(n - 1, idx));
             }
         }
 
         // Lagrange interpolate to get back the secret
         let eval_points: Vec<_> = shares
-            .iter()
-            .map(|&(idx, _)| C::ScalarField::from(idx as u64 + 1))
+            .clone()
+            .map(|(idx, _)| C::ScalarField::from(idx as u64 + 1))
             .collect();
-        let evals: Vec<_> = shares.iter().map(|&(_, share)| share).collect();
+        let evals: Vec<_> = shares.map(|(_, share)| share).collect();
         interpolate::<C>(&eval_points, &evals)
             .map_err(|e| VssError::FailedReconstruction(e.to_string()))
     }
@@ -217,6 +216,36 @@ impl<C: CurveGroup> FeldmanCommitment<C> {
     }
 }
 
+impl<C: CurveGroup> Add<FeldmanCommitment<C>> for FeldmanCommitment<C> {
+    type Output = FeldmanCommitment<C>;
+
+    fn add(self, other: FeldmanCommitment<C>) -> Self::Output {
+        &self + &other
+    }
+}
+
+impl<C: CurveGroup> Add<&FeldmanCommitment<C>> for FeldmanCommitment<C> {
+    type Output = FeldmanCommitment<C>;
+
+    fn add(self, other: &FeldmanCommitment<C>) -> Self::Output {
+        &self + other
+    }
+}
+
+impl<C: CurveGroup> Add<&FeldmanCommitment<C>> for &FeldmanCommitment<C> {
+    type Output = FeldmanCommitment<C>;
+
+    fn add(self, other: &FeldmanCommitment<C>) -> Self::Output {
+        let combined: Vec<C> = self
+            .comm
+            .iter()
+            .zip(other.comm.iter())
+            .map(|(x, y)| *x + y)
+            .collect();
+        C::normalize_batch(&combined).into()
+    }
+}
+
 impl<C: CurveGroup> KeyResharing<Self> for FeldmanVss<C> {
     fn reshare<R: Rng>(
         new_pp: &FeldmanVssPublicParam,
@@ -252,19 +281,17 @@ impl<C: CurveGroup> KeyResharing<Self> for FeldmanVss<C> {
     fn combine(
         old_pp: &FeldmanVssPublicParam,
         new_pp: &FeldmanVssPublicParam,
-        send_node_indices: &[usize],
-        row_commitments: &[FeldmanCommitment<C>],
         recv_node_idx: usize,
-        recv_reshares: &[C::ScalarField],
+        reshares: impl ExactSizeIterator<Item = (usize, C::ScalarField, FeldmanCommitment<C>)> + Clone,
     ) -> Result<(C::ScalarField, FeldmanCommitment<C>), VssError> {
         // input validation
         let n = old_pp.n.get();
-        if send_node_indices.is_empty() || row_commitments.is_empty() || recv_reshares.is_empty() {
+        if reshares.len() == 0 {
             return Err(VssError::EmptyReshare);
         }
-        for idx in send_node_indices.iter() {
-            if *idx >= n {
-                return Err(VssError::IndexOutOfBound(n - 1, *idx));
+        for (idx, _, _) in reshares.clone() {
+            if idx >= n {
+                return Err(VssError::IndexOutOfBound(n - 1, idx));
             }
         }
 
@@ -273,24 +300,23 @@ impl<C: CurveGroup> KeyResharing<Self> for FeldmanVss<C> {
         if recv_node_idx >= new_n {
             return Err(VssError::IndexOutOfBound(new_n - 1, recv_node_idx));
         }
-        if row_commitments.iter().any(|cm| cm.len() != new_t) {
-            return Err(VssError::InvalidCommitment);
-        }
-
-        let subset_size = recv_reshares.len();
-        if send_node_indices.len() != subset_size || row_commitments.len() != subset_size {
-            return Err(VssError::MismatchedInputLength);
+        for (_, _, row_commitment) in reshares.clone() {
+            if row_commitment.len() != new_t {
+                return Err(VssError::InvalidCommitment);
+            }
         }
 
         // interpolate reshares to get new secret share
-        let eval_points: Vec<_> = send_node_indices
-            .iter()
-            .map(|&idx| C::ScalarField::from(idx as u64 + 1))
+        let eval_points: Vec<_> = reshares
+            .clone()
+            .map(|(idx, _, _)| C::ScalarField::from(idx as u64 + 1))
             .collect();
-        let new_secret = interpolate::<C>(&eval_points, recv_reshares)
+        let recv_reshares: Vec<_> = reshares.clone().map(|(_, share, _)| share).collect();
+        let new_secret = interpolate::<C>(&eval_points, &recv_reshares)
             .map_err(|e| VssError::FailedCombine(e.to_string()))?;
 
         // interpolate in the exponent to get new Feldman commitment
+        let row_commitments: Vec<_> = reshares.map(|(_, _, commitment)| commitment).collect();
         let new_commitment = (0..new_t)
             .into_par_iter()
             .map(|j| {
@@ -461,15 +487,10 @@ mod tests {
             let selected_row_commitments: Vec<FeldmanCommitment<_>> =
                 (0..old_t).map(|i| row_commitments[i].clone()).collect();
 
-            let (new_secret_share, new_commitment) = FeldmanVss::<G1Projective>::combine(
-                &old_pp,
-                &new_pp,
-                &(0..old_t).collect::<Vec<_>>(),
-                &selected_row_commitments,
-                j,
-                &recv_reshares,
-            )
-            .unwrap();
+            let reshares_iter =
+                (0..old_t).map(|i| (i, recv_reshares[i], selected_row_commitments[i].clone()));
+            let (new_secret_share, new_commitment) =
+                FeldmanVss::<G1Projective>::combine(&old_pp, &new_pp, j, reshares_iter).unwrap();
 
             new_shares.push(new_secret_share);
             new_commitments.push(new_commitment);
diff --git a/timeboost-crypto/src/traits/dkg.rs b/timeboost-crypto/src/traits/dkg.rs
@@ -1,6 +1,7 @@
 //! Traits related to Distributed Key Generation (DKG) and Key Resharing
 
 use ark_std::rand::Rng;
+use std::ops::Add;
 use thiserror::Error;
 
 /// A trait for (t, n)-Verifiable Secret Sharing (VSS) schemes.
@@ -51,8 +52,24 @@ pub trait VerifiableSecretSharing {
     /// Returns `Ok(secret)` if reconstruction succeeds, or an appropriate `VssError` otherwise.
     fn reconstruct(
         pp: &Self::PublicParam,
-        shares: impl Iterator<Item = (usize, Self::SecretShare)>,
+        shares: impl ExactSizeIterator<Item = (usize, Self::SecretShare)> + Clone,
     ) -> Result<Self::Secret, VssError>;
+
+    /// Aggregates multiple commitments and secret shares into a single commitment and secret share.
+    ///
+    /// This is commonly used in DKG protocols to combine multiple dealings/contributions.
+    ///
+    /// Returns `Ok((secret_share, commitment))` if aggregation succeeds
+    fn aggregate<I>(dealings: I) -> Result<(Self::SecretShare, Self::Commitment), VssError>
+    where
+        I: Iterator<Item = (Self::SecretShare, Self::Commitment)>,
+        Self::Commitment: Add<Self::Commitment, Output = Self::Commitment>,
+        Self::SecretShare: Add<Self::SecretShare, Output = Self::SecretShare>,
+    {
+        dealings
+            .reduce(|(acc_share, acc_comm), (share, comm)| (acc_share + share, acc_comm + comm))
+            .ok_or(VssError::EmptyAggInput)
+    }
 }
 
 /// Publicly verifiable key resharing scheme for a VSS where existing share holders of a Shamir
@@ -94,10 +111,8 @@ pub trait KeyResharing<VSS: VerifiableSecretSharing> {
     fn combine(
         old_pp: &VSS::PublicParam,
         new_pp: &VSS::PublicParam,
-        send_node_indices: &[usize],
-        row_commitments: &[VSS::Commitment],
         recv_node_idx: usize,
-        recv_reshares: &[VSS::SecretShare],
+        reshares: impl ExactSizeIterator<Item = (usize, VSS::SecretShare, VSS::Commitment)> + Clone,
     ) -> Result<(VSS::Secret, VSS::Commitment), VssError>;
 }
 
@@ -118,6 +133,8 @@ pub enum VssError {
     FailedReconstruction(String),
     #[error("internal err: {0}")]
     InternalError(String),
+    #[error("aggregation input is empty")]
+    EmptyAggInput,
 
     #[error("reshare data is empty")]
     EmptyReshare,
diff --git a/timeboost-sequencer/Cargo.toml b/timeboost-sequencer/Cargo.toml
@@ -24,11 +24,11 @@ serde = { workspace = true }
 thiserror = { workspace = true }
 timeboost-crypto = { path = "../timeboost-crypto" }
 timeboost-types = { path = "../timeboost-types" }
+timeboost-utils = { path = "../timeboost-utils" }
 tokio = { workspace = true }
 tracing = { workspace = true }
 
 [dev-dependencies]
 ark-std = { workspace = true }
 bs58 = { workspace = true }
 portpicker = { workspace = true }
-timeboost-utils = { path = "../timeboost-utils" }
diff --git a/timeboost-sequencer/src/decrypt.rs b/timeboost-sequencer/src/decrypt.rs
@@ -20,6 +20,7 @@ use timeboost_crypto::{DecryptionScheme, Plaintext};
 use timeboost_types::{
     DecryptionKey, DkgAccumulator, DkgBundle, DkgKeyStore, InclusionList, Subset,
 };
+use timeboost_utils::ResultIter;
 use tokio::spawn;
 use tokio::sync::mpsc::{Receiver, Sender, channel};
 use tokio::task::JoinHandle;
@@ -662,24 +663,23 @@ impl Worker {
                 // TODO: centralize these constant, redeclared in DkgAccumulator.try_add()
                 let aad: &[u8; 3] = b"dkg";
                 let vess = Vess::new_fast();
-                let (shares, commitments) = subset
-                    .bundles()
-                    .iter()
-                    .map(|b| {
-                        vess.decrypt_share(committee, &self.dkg_sk, b.vess_ct(), aad)
-                            .map(|s| (s, b.comm().clone()))
-                            .map_err(|e| DecrypterError::Dkg(e.to_string()))
-                    })
-                    .collect::<Result<(Vec<_>, Vec<_>)>>()?;
+                let mut dealings_iter = ResultIter::new(subset.bundles().iter().map(|b| {
+                    vess.decrypt_share(committee, &self.dkg_sk, b.vess_ct(), aad)
+                        .map(|s| (s, b.comm().clone()))
+                }));
 
                 let dec_sk = DecryptionKey::from_dkg(
                     committee.size().into(),
                     self.dkg_sk.node_idx(),
-                    &commitments,
-                    &shares,
+                    &mut dealings_iter,
                 )
                 .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
 
+                // in case of early-return of ResultIter
+                dealings_iter
+                    .result()
+                    .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
+
                 self.enc_key.set(dec_sk.pubkey().clone());
                 self.dkg_state = DkgState::Completed(dec_sk);
                 info!(committee_id = %committee.id(), node = %self.label, "DKG finished (node successfully recovered)");
@@ -770,24 +770,23 @@ impl Worker {
                 // TODO:(alex) centralize these constant, redeclared in DkgAccumulator.try_add()
                 let aad: &[u8; 3] = b"dkg";
                 let vess = Vess::new_fast();
-                let (shares, commitments) = subset
-                    .bundles()
-                    .iter()
-                    .map(|b| {
-                        vess.decrypt_share(committee, &self.dkg_sk, b.vess_ct(), aad)
-                            .map(|s| (s, b.comm().clone()))
-                            .map_err(|e| DecrypterError::Dkg(e.to_string()))
-                    })
-                    .collect::<Result<(Vec<_>, Vec<_>)>>()?;
+                let mut dealings_iter = ResultIter::new(subset.bundles().iter().map(|b| {
+                    vess.decrypt_share(committee, &self.dkg_sk, b.vess_ct(), aad)
+                        .map(|s| (s, b.comm().clone()))
+                }));
 
                 let dec_sk = DecryptionKey::from_dkg(
                     committee.size().into(),
                     self.dkg_sk.node_idx(),
-                    &commitments,
-                    &shares,
+                    &mut dealings_iter,
                 )
                 .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
 
+                // in case of early-return of ResultIter
+                dealings_iter
+                    .result()
+                    .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
+
                 self.enc_key.set(dec_sk.pubkey().clone());
                 self.dkg_state = DkgState::Completed(dec_sk);
                 info!(committee_id = %committee.id(), node = %self.label, "DKG finished");
@@ -1418,8 +1417,12 @@ mod tests {
             .iter()
             .enumerate()
             .map(|(node_idx, shares)| {
-                super::DecryptionKey::from_dkg(COMMITTEE_SIZE, node_idx, &commitments, shares)
-                    .expect("threshold key derivation should succeed")
+                super::DecryptionKey::from_dkg(
+                    COMMITTEE_SIZE,
+                    node_idx,
+                    shares.iter().cloned().zip(commitments.iter().cloned()),
+                )
+                .expect("threshold key derivation should succeed")
             })
             .collect();
         tracing::info!(
diff --git a/timeboost-types/src/decryption.rs b/timeboost-types/src/decryption.rs
@@ -1,10 +1,9 @@
-use std::collections::{BTreeMap, btree_map};
-
 use anyhow::anyhow;
-use ark_ec::{AffineRepr, CurveGroup};
+use ark_ec::AffineRepr;
 use multisig::{Committee, CommitteeId, KeyId};
 use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
+use std::collections::{BTreeMap, btree_map};
 use timeboost_crypto::{
     DecryptionScheme,
     prelude::{DkgEncKey, Vess, Vss},
@@ -43,47 +42,34 @@ impl DecryptionKey {
     /// # Parameters
     /// - `committee_size`: size of the threshold committee
     /// - `node_idx`: in 0..committee_size, currently same as KeyId
-    /// - `commitments`: the Feldman Commitments: multiple output of `ShoupVess::encrypted_shares()`
-    /// - `key_shares`: multiple decrypted secret shares from `ShoupVess::decrypt_share()`
-    pub fn from_dkg(
+    /// - `dealings`: ResultIter containing decrypted shares and commitments
+    pub fn from_dkg<I>(
         committee_size: usize,
         node_idx: usize,
-        commitments: &[<Vss as VerifiableSecretSharing>::Commitment],
-        key_shares: &[<Vss as VerifiableSecretSharing>::SecretShare],
-    ) -> anyhow::Result<Self> {
-        anyhow::ensure!(
-            commitments.len() == key_shares.len(),
-            "mismatched input length"
-        );
-
-        // aggregate selected dealings/contributions
-        let agg_comm = commitments
-            .par_iter()
-            .cloned()
-            .reduce_with(|a, b| {
-                let combined: Vec<_> = a
-                    .into_iter()
-                    .zip(b.into_iter())
-                    // NOTE: ideally we can use C::normalize_batch(), but C is not exposed,
-                    // minor optimization, so ignore for now.
-                    .map(|(x, y)| (x + y).into_affine())
-                    .collect();
-                combined.into()
-            })
-            .ok_or_else(|| anyhow!("no commitments provided"))?;
-        let agg_key_share = key_shares.iter().sum();
+        mut dealings: I,
+    ) -> anyhow::Result<Self>
+    where
+        I: Iterator<
+            Item = (
+                <Vss as VerifiableSecretSharing>::SecretShare,
+                <Vss as VerifiableSecretSharing>::Commitment,
+            ),
+        >,
+    {
+        // aggregate selected dealings
+        let (agg_key_share, agg_comm) = Vss::aggregate(&mut dealings)?;
 
         // derive key material
-        Self::from_single_dkg(committee_size, node_idx, &agg_comm, agg_key_share)
+        Self::from_single_dkg(committee_size, node_idx, agg_key_share, &agg_comm)
     }
 
     /// inner routine to construct from a single (aggregated or interpolated) DKG output,
     /// shared in both DKG and resharing logic.
     fn from_single_dkg(
         committee_size: usize,
         node_idx: usize,
-        commitment: &<Vss as VerifiableSecretSharing>::Commitment,
         key_share: <Vss as VerifiableSecretSharing>::SecretShare,
+        commitment: &<Vss as VerifiableSecretSharing>::Commitment,
     ) -> anyhow::Result<Self> {
         // note: all .into() are made available via derive_more::From on those structs
         let pk: PublicKey = commitment
diff --git a/timeboost-utils/src/lib.rs b/timeboost-utils/src/lib.rs
