Skip to content

Commit

Permalink
make confirmations also versioned
Browse files Browse the repository at this point in the history
  • Loading branch information
benr-ml committed Jun 14, 2024
1 parent 4895178 commit 03ce277
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 35 deletions.
42 changes: 32 additions & 10 deletions crates/sui-core/src/authority/authority_per_epoch_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -96,11 +96,11 @@ use sui_types::message_envelope::TrustedEnvelope;
use sui_types::messages_checkpoint::{
CheckpointContents, CheckpointSequenceNumber, CheckpointSignatureMessage, CheckpointSummary,
};
use sui_types::messages_consensus::VersionedDkgMessage;
use sui_types::messages_consensus::{
check_total_jwk_size, AuthorityCapabilities, ConsensusTransaction, ConsensusTransactionKey,
ConsensusTransactionKind,
};
use sui_types::messages_consensus::{VersionedDkgConfimation, VersionedDkgMessage};
use sui_types::storage::GetSharedLocks;
use sui_types::sui_system_state::epoch_start_sui_system_state::{
EpochStartSystemState, EpochStartSystemStateTrait,
Expand Down Expand Up @@ -508,17 +508,23 @@ pub struct AuthorityEpochTables {
pub(crate) dkg_versioned_processed_messages: DBMap<PartyId, VersionedProcessedMessage>,
/// This table is no longer used (can be removed when DBMap supports removing tables)
#[allow(dead_code)]
#[deprecated]
pub(crate) dkg_processed_messages: DBMap<PartyId, dkg_v0::ProcessedMessage<PkG, EncG>>,

/// Records messages used to generate a DKG confirmation. Updated when enough DKG
/// messages are received to progress to the next phase.
pub(crate) dkg_versioned_used_messages: DBMap<u64, VersionedUsedProcessedMessages>,
/// This table is no longer used (can be removed when DBMap supports removing tables)
#[allow(dead_code)]
#[deprecated]
pub(crate) dkg_used_messages: DBMap<u64, dkg_v0::UsedProcessedMessages<PkG, EncG>>,

/// Records confirmations received from other nodes. Updated when receiving a new
/// dkg::Confirmation via consensus.
pub(crate) dkg_versioned_confirmations: DBMap<PartyId, VersionedDkgConfimation>,
/// This table is no longer used (can be removed when DBMap supports removing tables)
#[allow(dead_code)]
#[deprecated]
pub(crate) dkg_confirmations: DBMap<PartyId, dkg::Confirmation<EncG>>,
/// Records the final output of DKG after completion, including the public VSS key and
/// any local private shares.
Expand Down Expand Up @@ -3455,17 +3461,33 @@ impl AuthorityPerEpochStore {
"Received RandomnessDkgConfirmation from {:?}",
authority.concise()
);
match bcs::from_bytes(bytes) {
Ok(confirmation) => randomness_manager.add_confirmation(
batch,
authority,
confirmation,
)?,
Err(e) => {
warn!(
"Failed to deserialize RandomnessDkgMessage from {:?}: {e:?}",

if self.protocol_config.dkg_version() == 0 {
// old message was not an enum
match bcs::from_bytes(bytes) {
Ok(message) => randomness_manager.add_confirmation(
batch,
authority,
VersionedDkgConfimation::V0(message),
)?,

Err(e) => {
warn!(
"Failed to deserialize RandomnessDkgConfirmation from {:?}: {e:?}",
authority.concise(),
);
}
}
} else {
match bcs::from_bytes(bytes) {
Ok(message) => randomness_manager
.add_confirmation(batch, authority, message)?,
Err(e) => {
warn!(
"Failed to deserialize versioned RandomnessDkgMessage from {:?}: {e:?}",
authority.concise(),
);
}
}
}
} else {
Expand Down
81 changes: 59 additions & 22 deletions crates/sui-core/src/epoch/randomness.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use fastcrypto::error::FastCryptoResult;
use fastcrypto::groups::bls12381;
use fastcrypto::serde_helpers::ToFromByteArray;
use fastcrypto::traits::{KeyPair, ToFromBytes};
use fastcrypto_tbls::dkg::{Confirmation, Output};
use fastcrypto_tbls::dkg::Output;
use fastcrypto_tbls::nodes::PartyId;
use fastcrypto_tbls::{dkg, dkg_v0, dkg_v1, nodes};
use futures::stream::FuturesUnordered;
Expand All @@ -27,8 +27,8 @@ use sui_types::base_types::AuthorityName;
use sui_types::committee::{Committee, EpochId, StakeUnit};
use sui_types::crypto::{AuthorityKeyPair, RandomnessRound};
use sui_types::error::{SuiError, SuiResult};
use sui_types::messages_consensus::ConsensusTransaction;
use sui_types::messages_consensus::VersionedDkgMessage;
use sui_types::messages_consensus::{ConsensusTransaction, VersionedDkgConfimation};
use sui_types::sui_system_state::epoch_start_sui_system_state::EpochStartSystemStateTrait;
use tokio::sync::OnceCell;
use tokio::task::JoinHandle;
Expand Down Expand Up @@ -100,7 +100,7 @@ pub struct RandomnessManager {
enqueued_messages: BTreeMap<PartyId, JoinHandle<Option<VersionedProcessedMessage>>>,
processed_messages: BTreeMap<PartyId, VersionedProcessedMessage>,
used_messages: OnceCell<VersionedUsedProcessedMessages>,
confirmations: BTreeMap<PartyId, dkg::Confirmation<EncG>>,
confirmations: BTreeMap<PartyId, VersionedDkgConfimation>,
dkg_output: OnceCell<Option<dkg::Output<PkG, EncG>>>,

// State for randomness generation.
Expand Down Expand Up @@ -286,7 +286,7 @@ impl RandomnessManager {
}
rm.confirmations.extend(
tables
.dkg_confirmations
.dkg_versioned_confirmations
.safe_iter()
.map(|result| result.expect("typed_store should not fail")),
);
Expand Down Expand Up @@ -336,7 +336,7 @@ impl RandomnessManager {
fn merge_dkg_messages(
&self,
dkg_version: u64,
) -> FastCryptoResult<(Confirmation<EncG>, VersionedUsedProcessedMessages)> {
) -> FastCryptoResult<(VersionedDkgConfimation, VersionedUsedProcessedMessages)> {
match dkg_version {
0 => {
let (conf, msgs) = self.party.merge(
Expand All @@ -353,7 +353,10 @@ impl RandomnessManager {
.cloned()
.collect::<Vec<_>>(),
)?;
Ok((conf, VersionedUsedProcessedMessages::V0(msgs)))
Ok((
VersionedDkgConfimation::V0(conf),
VersionedUsedProcessedMessages::V0(msgs),
))
}
1 => {
let (conf, msgs) = self.party.merge_v1(
Expand All @@ -370,14 +373,16 @@ impl RandomnessManager {
.cloned()
.collect::<Vec<_>>(),
)?;
Ok((conf, VersionedUsedProcessedMessages::V1(msgs)))
Ok((
VersionedDkgConfimation::V1(conf),
VersionedUsedProcessedMessages::V1(msgs),
))
}
_ => panic!("BUG: invalid DKG version"),
}
}

fn complete_dkg(&self, dkg_version: u64) -> FastCryptoResult<Output<PkG, EncG>> {
let confs = &self.confirmations.values().cloned().collect::<Vec<_>>();
let rng = &mut StdRng::from_rng(OsRng).expect("RNG construction should not fail");
match dkg_version {
0 => self.party.complete(
Expand All @@ -390,7 +395,18 @@ impl RandomnessManager {
} else {
panic!("BUG: used_messages should be V0")
},
confs,
&self
.confirmations
.values()
.filter_map(|vm| {
if let VersionedDkgConfimation::V0(msg) = vm {
Some(msg)
} else {
panic!("BUG: invalid versioned message")
}
})
.cloned()
.collect::<Vec<_>>(),
rng,
),
1 => self.party.complete_v1(
Expand All @@ -403,7 +419,18 @@ impl RandomnessManager {
} else {
panic!("BUG: used_messages should be V1")
},
confs,
&self
.confirmations
.values()
.filter_map(|vm| {
if let VersionedDkgConfimation::V1(msg) = vm {
Some(msg)
} else {
panic!("BUG: invalid versioned message")
}
})
.cloned()
.collect::<Vec<_>>(),
rng,
),
_ => panic!("BUG: invalid DKG version"),
Expand Down Expand Up @@ -489,7 +516,7 @@ impl RandomnessManager {
Ok((conf, used_msgs)) => {
info!(
"random beacon: sending DKG Confirmation with {} complaints",
conf.complaints.len()
conf.num_of_complaints()
);
if self.used_messages.set(used_msgs.clone()).is_err() {
error!("BUG: used_messages should only ever be set once");
Expand Down Expand Up @@ -647,7 +674,7 @@ impl RandomnessManager {
&mut self,
batch: &mut DBBatch,
authority: &AuthorityName,
conf: dkg::Confirmation<EncG>,
conf: VersionedDkgConfimation,
) -> SuiResult {
if self.dkg_output.initialized() {
// Once we have completed DKG, no more `Confirmation`s are needed.
Expand All @@ -659,14 +686,14 @@ impl RandomnessManager {
);
return Ok(());
};
if *party_id != conf.sender {
if *party_id != conf.sender() {
warn!("ignoring equivocating DKG Confirmation from authority {authority:?} pretending to be PartyId {party_id:?}");
return Ok(());
}
self.confirmations.insert(conf.sender, conf.clone());
self.confirmations.insert(conf.sender(), conf.clone());
batch.insert_batch(
&self.tables()?.dkg_confirmations,
std::iter::once((conf.sender, conf)),
&self.tables()?.dkg_versioned_confirmations,
std::iter::once((conf.sender(), conf)),
)?;
Ok(())
}
Expand Down Expand Up @@ -940,7 +967,7 @@ mod tests {
let mut batch = epoch_stores[i]
.tables()
.unwrap()
.dkg_processed_messages
.dkg_versioned_processed_messages
.batch();
for (j, dkg_message) in dkg_messages.iter().cloned().enumerate() {
randomness_managers[i]
Expand All @@ -961,15 +988,25 @@ mod tests {
assert!(dkg_confirmation.len() == 1);
match dkg_confirmation.remove(0).kind {
ConsensusTransactionKind::RandomnessDkgConfirmation(_, bytes) => {
let msg: fastcrypto_tbls::dkg::Confirmation<EncG> = bcs::from_bytes(&bytes)
.expect("DKG confirmation deserialization should not fail");
dkg_confirmations.push(msg);
if version == 0 {
let msg: fastcrypto_tbls::dkg::Confirmation<EncG> = bcs::from_bytes(&bytes)
.expect("DKG confirmation deserialization should not fail");
dkg_confirmations.push(VersionedDkgConfimation::V0(msg));
} else {
let msg: VersionedDkgConfimation = bcs::from_bytes(&bytes)
.expect("DKG message deserialization should not fail");
dkg_confirmations.push(msg);
}
}
_ => panic!("wrong type of message sent"),
}
}
for i in 0..randomness_managers.len() {
let mut batch = epoch_stores[i].tables().unwrap().dkg_confirmations.batch();
let mut batch = epoch_stores[i]
.tables()
.unwrap()
.dkg_versioned_confirmations
.batch();
for (j, dkg_confirmation) in dkg_confirmations.iter().cloned().enumerate() {
randomness_managers[i]
.add_confirmation(&mut batch, &epoch_stores[j].name, dkg_confirmation)
Expand Down Expand Up @@ -1088,7 +1125,7 @@ mod tests {
let mut batch = epoch_stores[i]
.tables()
.unwrap()
.dkg_processed_messages
.dkg_versioned_processed_messages
.batch();
for (j, dkg_message) in dkg_messages.iter().cloned().enumerate() {
randomness_managers[i]
Expand Down
34 changes: 31 additions & 3 deletions crates/sui-types/src/messages_consensus.rs
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,11 @@ pub enum VersionedDkgMessage {
V1(dkg_v1::Message<bls12381::G2Element, bls12381::G2Element>),
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum VersionedDkgConfimation {
V0(dkg::Confirmation<bls12381::G2Element>),
V1(dkg::Confirmation<bls12381::G2Element>),
}
impl Debug for VersionedDkgMessage {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Expand Down Expand Up @@ -252,6 +257,22 @@ impl VersionedDkgMessage {
}
}

impl VersionedDkgConfimation {
pub fn sender(&self) -> u16 {
match self {
VersionedDkgConfimation::V0(msg) => msg.sender,
VersionedDkgConfimation::V1(msg) => msg.sender,
}
}

pub fn num_of_complaints(&self) -> usize {
match self {
VersionedDkgConfimation::V0(msg) => msg.complaints.len(),
VersionedDkgConfimation::V1(msg) => msg.complaints.len(),
}
}
}

impl ConsensusTransaction {
pub fn new_certificate_message(
authority: &AuthorityName,
Expand Down Expand Up @@ -346,10 +367,17 @@ impl ConsensusTransaction {
}
pub fn new_randomness_dkg_confirmation(
authority: AuthorityName,
confirmation: &dkg::Confirmation<bls12381::G2Element>,
versioned_confirmation: &VersionedDkgConfimation,
) -> Self {
let confirmation =
bcs::to_bytes(confirmation).expect("message serialization should not fail");
let confirmation = match versioned_confirmation {
VersionedDkgConfimation::V0(msg) => {
// Old version does not use the enum, so we need to serialize it separately.
bcs::to_bytes(msg).expect("message serialization should not fail")
}
_ => bcs::to_bytes(versioned_confirmation)
.expect("message serialization should not fail"),
};

let mut hasher = DefaultHasher::new();
confirmation.hash(&mut hasher);
let tracking_id = hasher.finish().to_le_bytes();
Expand Down

0 comments on commit 03ce277

Please sign in to comment.