diff --git a/timeboost-crypto/src/feldman.rs b/timeboost-crypto/src/feldman.rs index 7445adc5..bd55dd96 100644 --- a/timeboost-crypto/src/feldman.rs +++ b/timeboost-crypto/src/feldman.rs @@ -302,17 +302,25 @@ impl KeyResharing for FeldmanVss { old_pp: &FeldmanVssPublicParam, new_pp: &FeldmanVssPublicParam, recv_node_idx: usize, - reshares: impl ExactSizeIterator)> + Clone, + reshares: impl Iterator)>, ) -> Result<(C::ScalarField, FeldmanCommitment), VssError> { - // input validation let n = old_pp.n.get(); - if reshares.len() == 0 { - return Err(VssError::EmptyReshare); - } - for (idx, _, _) in reshares.clone() { + + let mut eval_points = Vec::new(); + let mut recv_shares = Vec::new(); + let mut row_commitments = Vec::new(); + + for (idx, share, commitment) in reshares { if idx >= n { return Err(VssError::IndexOutOfBound(n - 1, idx)); } + eval_points.push(C::ScalarField::from(idx as u64 + 1)); + recv_shares.push(share); + row_commitments.push(commitment); + } + + if eval_points.is_empty() { + return Err(VssError::EmptyReshare); } let new_n = new_pp.n.get(); @@ -320,23 +328,16 @@ impl KeyResharing for FeldmanVss { if recv_node_idx >= new_n { return Err(VssError::IndexOutOfBound(new_n - 1, recv_node_idx)); } - for (_, _, row_commitment) in reshares.clone() { - if row_commitment.len() != new_t { - return Err(VssError::InvalidCommitment); - } + + if row_commitments.iter().any(|c| c.len() != new_t) { + return Err(VssError::InvalidCommitment); } // interpolate reshares to get new secret share - let eval_points: Vec<_> = reshares - .clone() - .map(|(idx, _, _)| C::ScalarField::from(idx as u64 + 1)) - .collect(); - let recv_reshares: Vec<_> = reshares.clone().map(|(_, share, _)| share).collect(); - let new_secret = interpolate::(&eval_points, &recv_reshares) + let new_secret = interpolate::(&eval_points, &recv_shares) .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| { @@ -344,9 +345,9 @@ impl KeyResharing for FeldmanVss { row_commitments.iter().map(|row| row[j]).collect(); interpolate_in_exponent::(&eval_points, &j_th_coeffs) .map_err(|e| VssError::FailedCombine(e.to_string())) + .map(|p| p.into_affine()) }) .collect::, VssError>>()?; - let new_commitment = C::normalize_batch(&new_commitment); Ok((new_secret, new_commitment.into())) } diff --git a/timeboost-crypto/src/traits/dkg.rs b/timeboost-crypto/src/traits/dkg.rs index d0a2f0fa..0bab5397 100644 --- a/timeboost-crypto/src/traits/dkg.rs +++ b/timeboost-crypto/src/traits/dkg.rs @@ -112,7 +112,7 @@ pub trait KeyResharing { old_pp: &VSS::PublicParam, new_pp: &VSS::PublicParam, recv_node_idx: usize, - reshares: impl ExactSizeIterator + Clone, + reshares: impl Iterator, ) -> Result<(VSS::Secret, VSS::Commitment), VssError>; } diff --git a/timeboost-sequencer/Cargo.toml b/timeboost-sequencer/Cargo.toml index 78c0643c..89724b17 100644 --- a/timeboost-sequencer/Cargo.toml +++ b/timeboost-sequencer/Cargo.toml @@ -33,5 +33,5 @@ tracing = { workspace = true } ark-std = { workspace = true } bs58 = { workspace = true } futures = { workspace = true } -test-utils = { path = "../test-utils", features = ["ports"]} +test-utils = { path = "../test-utils", features = ["ports"] } timeboost-utils = { path = "../timeboost-utils" } diff --git a/timeboost-sequencer/src/decrypt.rs b/timeboost-sequencer/src/decrypt.rs index 6b534bc5..923ec155 100644 --- a/timeboost-sequencer/src/decrypt.rs +++ b/timeboost-sequencer/src/decrypt.rs @@ -1,3 +1,7 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; +use std::result::Result as StdResult; +use std::sync::Arc; + use ark_std::{UniformRand, rand::thread_rng}; use bon::Builder; use bytes::{BufMut, Bytes, BytesMut}; @@ -5,24 +9,20 @@ use cliquenet::overlay::{Data, DataError, NetworkDown, Overlay}; use cliquenet::{ AddressableCommittee, MAX_MESSAGE_SIZE, Network, NetworkError, NetworkMetrics, Role, }; -use rayon::iter::ParallelIterator; -use rayon::prelude::*; - use multisig::{Committee, CommitteeId, PublicKey}; use parking_lot::RwLock; +use rayon::iter::ParallelIterator; +use rayon::prelude::*; use sailfish::types::{Evidence, Round, RoundNumber}; use serde::{Deserialize, Serialize}; -use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; -use std::result::Result as StdResult; -use std::sync::Arc; use timeboost_config::DECRYPTER_PORT_OFFSET; use timeboost_crypto::prelude::{ DkgDecKey, LabeledDkgDecKey, Plaintext, ThresholdCiphertext, ThresholdDecShare, ThresholdEncError, ThresholdEncScheme, ThresholdScheme, Vess, VssSecret, }; use timeboost_types::{ - AccumulatorMode, DkgAccumulator, DkgBundle, DkgSubset, InclusionList, KeyStore, KeyStoreVec, - ThresholdKey, ThresholdKeyCell, + AccumulatorMode, DkgAccumulator, DkgBundle, DkgSubset, DkgSubsetRef, InclusionList, KeyStore, + KeyStoreVec, ThresholdKey, ThresholdKeyCell, }; use tokio::spawn; use tokio::sync::mpsc::{Receiver, Sender, channel}; @@ -728,21 +728,31 @@ impl Worker { .ok_or(DecrypterError::NoCommittee(self.current))?; let prev = (guard.len() == 2).then(|| guard.last().clone()); - - subsets.insert(src, res.subset.to_owned()); let committee = current.committee(); let threshold: usize = committee.one_honest_threshold().into(); - let mut counts = HashMap::new(); - for subset in subsets.values() { - *counts.entry(subset).or_insert(0) += 1; + subsets.insert(src, res.subset); + + let mut subset_groups: HashMap> = HashMap::new(); + let mut threshold_subset = None; + + for (pk, subset) in subsets.iter() { + let pks = subset_groups.entry(subset.as_ref()).or_default(); + pks.push(pk); + + if pks.len() >= threshold { + threshold_subset = Some((subset.as_ref(), pks[0])); + break; + } } - if let Some((&subset, _)) = counts.iter().find(|(_, count)| **count >= threshold) { - let acc = DkgAccumulator::from_subset(current.clone(), subset.to_owned()); + if let Some((subset_ref, pk)) = threshold_subset { + let subset = subsets.to_owned().remove(pk).expect("subset exists in map"); + let acc = DkgAccumulator::from_subset(current.clone(), subset); let mode = acc.mode().clone(); self.tracker.insert(committee.id(), acc); - let dec_key = subset + + let dec_key = subset_ref .extract_key(current, &self.dkg_sk, prev.as_ref()) .map_err(|e| DecrypterError::Dkg(e.to_string()))?; @@ -830,7 +840,7 @@ impl Worker { if matches!(mode, AccumulatorMode::Dkg) { if let Some(subset) = acc.try_finalize() { let dec_key = subset - .extract_key(&key_store.clone(), &self.dkg_sk, None) + .extract_key(&key_store, &self.dkg_sk, None) .map_err(|e| DecrypterError::Dkg(e.to_string()))?; self.dec_key.set(dec_key); self.state = WorkerState::Running; diff --git a/timeboost-types/src/decryption.rs b/timeboost-types/src/decryption.rs index e8d0b352..d45e8ed0 100644 --- a/timeboost-types/src/decryption.rs +++ b/timeboost-types/src/decryption.rs @@ -81,7 +81,7 @@ impl ThresholdKey { dealings: I, ) -> anyhow::Result where - I: ExactSizeIterator + Clone, + I: Iterator, { let old_pp = FeldmanVssPublicParam::from(old_committee); let new_pp = FeldmanVssPublicParam::from(new_committee); @@ -131,10 +131,15 @@ impl ThresholdKey { /// `DecryptionKeyCell` is a thread-safe container for an optional `DecryptionKey` /// that allows asynchronous notification when the key is set. -#[derive(Debug, Clone, Default)] +#[derive(Debug, Default)] pub struct ThresholdKeyCell { - key: Arc>>, - notify: Arc, + inner: Arc, +} + +#[derive(Debug, Default)] +struct ThresholdKeyCellInner { + key: RwLock>, + notify: Notify, } impl ThresholdKeyCell { @@ -143,12 +148,12 @@ impl ThresholdKeyCell { } pub fn set(&self, key: ThresholdKey) { - *self.key.write() = Some(key); - self.notify.notify_waiters(); + *self.inner.key.write() = Some(key); + self.inner.notify.notify_waiters(); } pub fn get(&self) -> Option { - (*self.key.read()).clone() + (*self.inner.key.read()).clone() } pub fn enc_key(&self) -> Option { @@ -156,12 +161,12 @@ impl ThresholdKeyCell { } pub fn get_ref(&self) -> impl Deref> { - self.key.read() + self.inner.key.read() } pub async fn read(&self) -> ThresholdKey { loop { - let fut = self.notify.notified(); + let fut = self.inner.notify.notified(); if let Some(k) = self.get() { return k; } @@ -170,6 +175,14 @@ impl ThresholdKeyCell { } } +impl Clone for ThresholdKeyCell { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + /// A small, non-empty collection of KeyStores. #[derive(Debug, Default, Clone)] #[allow(clippy::len_without_is_empty)] @@ -280,10 +293,10 @@ impl From for KeyStoreVec { /// A `KeyStore` with committee information and public keys used in the DKG or key resharing #[derive(Debug, Clone)] -pub struct KeyStore(Arc); +pub struct KeyStore(Arc); #[derive(Debug)] -struct Inner { +struct KeyStoreInner { committee: Committee, keys: BTreeMap, } @@ -294,7 +307,7 @@ impl KeyStore { I: IntoIterator, T: Into, { - let this = Self(Arc::new(Inner { + let this = Self(Arc::new(KeyStoreInner { committee: c, keys: keys.into_iter().map(|(i, k)| (i.into(), k)).collect(), })); @@ -446,20 +459,19 @@ impl DkgAccumulator { } /// Try to finalize the accumulator into a subset if enough bundles are collected. - pub fn try_finalize(&mut self) -> Option { + /// Returns a reference to the internal data to avoid cloning the bundles. + pub fn try_finalize(&self) -> Option> { if self.complete { let combkey = match &self.mode { AccumulatorMode::Dkg => None, AccumulatorMode::Resharing(combkey) => Some(combkey.clone()), }; - let subset = DkgSubset { + Some(DkgSubsetRef { committee_id: self.committee().id(), - bundles: self.bundles.clone(), + bundles: &self.bundles, combkey, - }; - - Some(subset) + }) } else { None } @@ -488,6 +500,14 @@ pub struct DkgSubset { combkey: Option, } +/// A reference-based version of DkgSubset to avoid cloning bundles. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct DkgSubsetRef<'a> { + committee_id: CommitteeId, + bundles: &'a [DkgBundle], + combkey: Option, +} + impl DkgSubset { /// Create a new subset with DKG bundles. pub fn new_dkg(committee_id: CommitteeId, bundles: Vec) -> Self { @@ -536,6 +556,27 @@ impl DkgSubset { self.combkey.is_some() } + /// Convert this DkgSubset to a DkgSubsetRef. + pub fn as_ref(&self) -> DkgSubsetRef<'_> { + DkgSubsetRef { + committee_id: self.committee_id, + bundles: &self.bundles, + combkey: self.combkey.as_ref().cloned(), + } + } + + /// Extract the new threshold decryption key from the subset. + pub fn extract_key( + &self, + curr: &KeyStore, + dkg_sk: &LabeledDkgDecKey, + prev: Option<&KeyStore>, + ) -> anyhow::Result { + self.as_ref().extract_key(curr, dkg_sk, prev) + } +} + +impl<'a> DkgSubsetRef<'a> { /// Extract the new threshold decryption key from the subset. pub fn extract_key( &self, @@ -547,7 +588,7 @@ impl DkgSubset { match &self.combkey { None => { - let mut dealings_iter = ResultIter::new(self.bundles().iter().map(|b| { + let mut dealings_iter = ResultIter::new(self.bundles.iter().map(|b| { vess.decrypt_share(curr.committee(), dkg_sk, b.vess_ct(), DKG_AAD) .map(|s| (s, b.comm().clone())) })); @@ -559,38 +600,28 @@ impl DkgSubset { )?; dealings_iter.result()?; - Ok(dec_key) } Some(combkey) => { - let Some(prev) = prev else { - return Err(anyhow!("previous key store missing")); - }; + let prev = prev.ok_or_else(|| anyhow!("previous key store missing"))?; + let mut dealings_iter = ResultIter::new(self.bundles.iter().map(|b| { + let node_idx = b.origin().0.into(); + let pub_share = combkey + .get_pub_share(node_idx) + .ok_or(VessError::FailedVerification)?; + vess.decrypt_reshare(curr.committee(), dkg_sk, b.vess_ct(), DKG_AAD, *pub_share) + .map(|s| (node_idx, s, b.comm().clone())) + })); - let dealings: Vec<_> = self - .bundles() - .iter() - .map(|b| { - let node_idx = b.origin().0.into(); - let pub_share = combkey - .get_pub_share(node_idx) - .ok_or(VessError::FailedVerification)?; - let s = vess.decrypt_reshare( - curr.committee(), - dkg_sk, - b.vess_ct(), - DKG_AAD, - *pub_share, - )?; - Ok((node_idx, s, b.comm().clone())) - }) - .collect::, VessError>>()?; - ThresholdKey::from_resharing( + let dec_key = ThresholdKey::from_resharing( prev.committee(), curr.committee(), dkg_sk.node_idx(), - dealings.into_iter(), - ) + &mut dealings_iter, + )?; + + dealings_iter.result()?; + Ok(dec_key) } } } diff --git a/timeboost-types/src/lib.rs b/timeboost-types/src/lib.rs index f6eaf0df..a04239d9 100644 --- a/timeboost-types/src/lib.rs +++ b/timeboost-types/src/lib.rs @@ -22,8 +22,8 @@ pub use bundle::{ pub use bytes::Bytes; pub use candidate_list::{CandidateList, CandidateListBytes}; pub use decryption::{ - AccumulatorMode, DkgAccumulator, DkgSubset, KeyStore, KeyStoreVec, ResultIter, ThresholdKey, - ThresholdKeyCell, + AccumulatorMode, DkgAccumulator, DkgSubset, DkgSubsetRef, KeyStore, KeyStoreVec, ResultIter, + ThresholdKey, ThresholdKeyCell, }; pub use delayed_inbox::DelayedInboxIndex; pub use inclusion_list::InclusionList;