randomness #2: real dkg from randomnet (aptos-labs#12127)

zjma · danielxiangzl · web-flow · commit e8b9b289aa28 · 2024-02-22T16:31:43.000-08:00
* types update from randomnet

* update

* lint

* real dkg and rounding

* rounding

---------

Co-authored-by: danielxiangzl &lt;xiangzhuolun@gmail.com&gt;
diff --git a/dkg/src/dkg_manager/mod.rs b/dkg/src/dkg_manager/mod.rs
@@ -16,7 +16,7 @@ use aptos_validator_transaction_pool::{TxnGuard, VTxnPoolState};
 use futures_channel::oneshot;
 use futures_util::{future::AbortHandle, FutureExt, StreamExt};
 use move_core_types::account_address::AccountAddress;
-use rand::thread_rng;
+use rand::{prelude::StdRng, thread_rng, SeedableRng};
 use std::sync::Arc;
 
 #[allow(dead_code)]
@@ -207,12 +207,12 @@ impl<DKG: DKGTrait> DKGManager<DKG> {
         self.state = match &self.state {
             InnerState::NotStarted => {
                 let public_params = DKG::new_public_params(dkg_session_metadata);
-                let mut rng = thread_rng();
-                let input_secret = if cfg!(feature = "smoke-test") {
-                    DKG::generate_predictable_input_secret_for_testing(self.dealer_sk.as_ref())
+                let mut rng = if cfg!(feature = "smoke-test") {
+                    StdRng::from_seed(self.my_addr.into_bytes())
                 } else {
-                    DKG::InputSecret::generate(&mut rng)
+                    StdRng::from_rng(thread_rng()).unwrap()
                 };
+                let input_secret = DKG::InputSecret::generate(&mut rng);
 
                 let trx = DKG::generate_transcript(
                     &mut rng,
diff --git a/dkg/src/transcript_aggregation/mod.rs b/dkg/src/transcript_aggregation/mod.rs
@@ -76,8 +76,7 @@ impl<S: DKGTrait> BroadcastStatus<DKGMessage> for Arc<TranscriptAggregationState
         // All checks passed. Aggregating.
         trx_aggregator.contributors.insert(metadata.author);
         if let Some(agg_trx) = trx_aggregator.trx.as_mut() {
-            let acc = std::mem::take(agg_trx);
-            *agg_trx = S::aggregate_transcripts(&self.dkg_pub_params, vec![acc, transcript]);
+            S::aggregate_transcripts(&self.dkg_pub_params, agg_trx, transcript);
         } else {
             trx_aggregator.trx = Some(transcript);
         }
diff --git a/state-sync/inter-component/event-notifications/src/tests.rs b/state-sync/inter-component/event-notifications/src/tests.rs
@@ -436,7 +436,7 @@ fn count_event_notifications_and_ensure_ordering(listener: &mut EventNotificatio
             notification_count += 1;
             assert_lt!(
                 last_version_received,
-                event_notification.version.try_into().unwrap()
+                std::convert::TryInto::<i64>::try_into(event_notification.version).unwrap()
             );
             last_version_received = event_notification.version.try_into().unwrap();
         } else {
diff --git a/types/src/dkg/dummy_dkg/mod.rs b/types/src/dkg/dummy_dkg/mod.rs
@@ -14,6 +14,7 @@ pub struct DummyDKG {}
 
 impl DKGTrait for DummyDKG {
     type DealerPrivateKey = bls12381::PrivateKey;
+    type DealtPubKeyShare = ();
     type DealtSecret = DummySecret;
     type DealtSecretShare = DummySecret;
     type InputSecret = DummySecret;
@@ -29,11 +30,14 @@ impl DKGTrait for DummyDKG {
         DummySecret::aggregate(secrets)
     }
 
-    fn dealt_secret_from_input(input: &Self::InputSecret) -> Self::DealtSecret {
+    fn dealt_secret_from_input(
+        _pub_params: &Self::PublicParams,
+        input: &Self::InputSecret,
+    ) -> Self::DealtSecret {
         *input
     }
 
-    fn generate_transcript<R: CryptoRng>(
+    fn generate_transcript<R: CryptoRng + RngCore>(
         _rng: &mut R,
         _params: &Self::PublicParams,
         input_secret: &Self::InputSecret,
@@ -63,31 +67,27 @@ impl DKGTrait for DummyDKG {
 
     fn aggregate_transcripts(
         _params: &Self::PublicParams,
-        transcripts: Vec<DummyDKGTranscript>,
-    ) -> DummyDKGTranscript {
-        let mut all_secrets = vec![];
-        let mut agg_contributions_by_dealer = BTreeMap::new();
-        for transcript in transcripts {
-            let DummyDKGTranscript {
-                secret,
-                contributions_by_dealer,
-            } = transcript;
-            all_secrets.push(secret);
-            agg_contributions_by_dealer.extend(contributions_by_dealer);
-        }
-        DummyDKGTranscript {
-            secret: DummySecret::aggregate(all_secrets),
-            contributions_by_dealer: agg_contributions_by_dealer,
-        }
+        accumulator: &mut Self::Transcript,
+        element: Self::Transcript,
+    ) {
+        let DummyDKGTranscript {
+            secret,
+            contributions_by_dealer,
+        } = element;
+        accumulator
+            .contributions_by_dealer
+            .extend(contributions_by_dealer);
+        accumulator.secret =
+            DummySecret::aggregate(vec![std::mem::take(&mut accumulator.secret), secret]);
     }
 
     fn decrypt_secret_share_from_transcript(
         _pub_params: &Self::PublicParams,
         transcript: &DummyDKGTranscript,
         _player_idx: u64,
         _dk: &Self::NewValidatorDecryptKey,
-    ) -> anyhow::Result<DummySecret> {
-        Ok(transcript.secret)
+    ) -> anyhow::Result<(DummySecret, ())> {
+        Ok((transcript.secret, ()))
     }
 
     fn reconstruct_secret_from_shares(
@@ -108,15 +108,6 @@ impl DKGTrait for DummyDKG {
     fn get_dealers(transcript: &DummyDKGTranscript) -> BTreeSet<u64> {
         transcript.contributions_by_dealer.keys().copied().collect()
     }
-
-    fn generate_predictable_input_secret_for_testing(
-        dealer_sk: &bls12381::PrivateKey,
-    ) -> DummySecret {
-        let bytes_8: [u8; 8] = dealer_sk.to_bytes()[0..8].try_into().unwrap();
-        DummySecret {
-            val: u64::from_be_bytes(bytes_8),
-        }
-    }
 }
 
 #[derive(Copy, Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
diff --git a/types/src/dkg/dummy_dkg/tests.rs b/types/src/dkg/dummy_dkg/tests.rs
@@ -118,7 +118,11 @@ fn test_dummy_dkg_correctness() {
         .iter()
         .map(|state| state.transcript.clone().unwrap())
         .collect();
-    let agg_transcript = DummyDKG::aggregate_transcripts(&pub_params, all_transcripts);
+    let mut agg_transcript = DummyDKGTranscript::default();
+    all_transcripts.into_iter().for_each(|trx| {
+        DummyDKG::aggregate_transcripts(&pub_params, &mut agg_transcript, trx);
+    });
+
     assert!(DummyDKG::verify_transcript(&pub_params, &agg_transcript).is_ok());
 
     // Optional check: bad transcript should be rejected.
@@ -128,13 +132,14 @@ fn test_dummy_dkg_correctness() {
 
     // Every new validator decrypt their own secret share.
     for (idx, nvi) in new_validator_states.iter_mut().enumerate() {
-        nvi.secret_share = DummyDKG::decrypt_secret_share_from_transcript(
+        let (secret, _pub_key) = DummyDKG::decrypt_secret_share_from_transcript(
             &pub_params,
             &agg_transcript,
             idx as u64,
             &nvi.sk,
         )
-        .ok()
+        .unwrap();
+        nvi.secret_share = Some(secret);
     }
 
     // The dealt secret should be reconstructable.
@@ -147,7 +152,9 @@ fn test_dummy_dkg_correctness() {
         DummyDKG::reconstruct_secret_from_shares(&pub_params, player_share_pairs).unwrap();
 
     let all_input_secrets = dealer_states.iter().map(|ds| ds.input_secret).collect();
-    let dealt_secret_from_input =
-        DummyDKG::dealt_secret_from_input(&DummyDKG::aggregate_input_secret(all_input_secrets));
+    let dealt_secret_from_input = DummyDKG::dealt_secret_from_input(
+        &pub_params,
+        &DummyDKG::aggregate_input_secret(all_input_secrets),
+    );
     assert_eq!(dealt_secret_from_reconstruct, dealt_secret_from_input);
 }
diff --git a/types/src/dkg/mod.rs b/types/src/dkg/mod.rs
@@ -1,8 +1,9 @@
 // Copyright © Aptos Foundation
 
 use crate::{
-    dkg::dummy_dkg::DummyDKG, on_chain_config::OnChainConfig,
-    validator_verifier::ValidatorConsensusInfoMoveStruct,
+    dkg::real_dkg::RealDKG,
+    on_chain_config::OnChainConfig,
+    validator_verifier::{ValidatorConsensusInfo, ValidatorConsensusInfoMoveStruct},
 };
 use anyhow::Result;
 use aptos_crypto::Uniform;
@@ -12,9 +13,12 @@ use move_core_types::{
     move_resource::MoveStructType,
 };
 use once_cell::sync::Lazy;
-use rand::CryptoRng;
+use rand::{CryptoRng, RngCore};
 use serde::{Deserialize, Serialize};
-use std::{collections::BTreeSet, fmt::Debug};
+use std::{
+    collections::BTreeSet,
+    fmt::{Debug, Formatter},
+};
 
 #[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Eq, CryptoHasher, BCSCryptoHash)]
 pub struct DKGTranscriptMetadata {
@@ -37,13 +41,22 @@ pub static DKG_START_EVENT_MOVE_TYPE_TAG: Lazy<TypeTag> =
     Lazy::new(|| TypeTag::Struct(Box::new(DKGStartEvent::struct_tag())));
 
 /// DKG transcript and its metadata.
-#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Eq)]
+#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct DKGTranscript {
     pub metadata: DKGTranscriptMetadata,
     #[serde(with = "serde_bytes")]
     pub transcript_bytes: Vec<u8>,
 }
 
+impl Debug for DKGTranscript {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("DKGTranscript")
+            .field("metadata", &self.metadata)
+            .field("transcript_bytes_len", &self.transcript_bytes.len())
+            .finish()
+    }
+}
+
 impl DKGTranscript {
     pub fn new(epoch: u64, author: AccountAddress, transcript_bytes: Vec<u8>) -> Self {
         Self {
@@ -63,52 +76,88 @@ impl DKGTranscript {
     }
 }
 
-// The input of DKG.
+/// Reflection of `0x1::dkg::DKGSessionMetadata` in rust.
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub struct DKGSessionMetadata {
     pub dealer_epoch: u64,
     pub dealer_validator_set: Vec<ValidatorConsensusInfoMoveStruct>,
     pub target_validator_set: Vec<ValidatorConsensusInfoMoveStruct>,
 }
 
-// The input and the run state of DKG.
+impl DKGSessionMetadata {
+    pub fn target_validator_consensus_infos_cloned(&self) -> Vec<ValidatorConsensusInfo> {
+        self.target_validator_set
+            .clone()
+            .into_iter()
+            .map(|obj| obj.try_into().unwrap())
+            .collect()
+    }
+
+    pub fn dealer_consensus_infos_cloned(&self) -> Vec<ValidatorConsensusInfo> {
+        self.dealer_validator_set
+            .clone()
+            .into_iter()
+            .map(|obj| obj.try_into().unwrap())
+            .collect()
+    }
+}
+
 /// Reflection of Move type `0x1::dkg::DKGSessionState`.
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub struct DKGSessionState {
     pub metadata: DKGSessionMetadata,
     pub start_time_us: u64,
-    pub result: Vec<u8>,
-    pub deadline_microseconds: u64,
+    pub transcript: Vec<u8>,
 }
 
+impl DKGSessionState {
+    pub fn target_epoch(&self) -> u64 {
+        self.metadata.dealer_epoch + 1
+    }
+}
 /// Reflection of Move type `0x1::dkg::DKGState`.
 #[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
 pub struct DKGState {
-    pub last_complete: Option<DKGSessionState>,
+    pub last_completed: Option<DKGSessionState>,
     pub in_progress: Option<DKGSessionState>,
 }
 
+impl DKGState {
+    pub fn maybe_last_complete(&self, epoch: u64) -> Option<&DKGSessionState> {
+        match &self.last_completed {
+            Some(session) if session.target_epoch() == epoch => Some(session),
+            _ => None,
+        }
+    }
+
+    pub fn last_complete(&self) -> &DKGSessionState {
+        self.last_completed.as_ref().unwrap()
+    }
+}
+
 impl OnChainConfig for DKGState {
     const MODULE_IDENTIFIER: &'static str = "dkg";
     const TYPE_IDENTIFIER: &'static str = "DKGState";
 }
 
+/// NOTE: this is a subset of the full scheme. Some data items/algorithms are not used in DKG and are omitted.
 pub trait DKGTrait: Debug {
     type DealerPrivateKey;
     type PublicParams: Clone + Debug + Send + Sync;
-    type Transcript: Clone + Default + Send + Sync + Serialize + for<'a> Deserialize<'a>;
+    type Transcript: Clone + Send + Sync + Serialize + for<'a> Deserialize<'a>;
     type InputSecret: Uniform;
     type DealtSecret;
     type DealtSecretShare;
-    type NewValidatorDecryptKey;
+    type DealtPubKeyShare;
+    type NewValidatorDecryptKey: Uniform;
 
     fn new_public_params(dkg_session_metadata: &DKGSessionMetadata) -> Self::PublicParams;
-    fn generate_predictable_input_secret_for_testing(
-        dealer_sk: &Self::DealerPrivateKey,
-    ) -> Self::InputSecret;
     fn aggregate_input_secret(secrets: Vec<Self::InputSecret>) -> Self::InputSecret;
-    fn dealt_secret_from_input(input: &Self::InputSecret) -> Self::DealtSecret;
-    fn generate_transcript<R: CryptoRng>(
+    fn dealt_secret_from_input(
+        pub_params: &Self::PublicParams,
+        input: &Self::InputSecret,
+    ) -> Self::DealtSecret;
+    fn generate_transcript<R: CryptoRng + RngCore>(
         rng: &mut R,
         params: &Self::PublicParams,
         input_secret: &Self::InputSecret,
@@ -120,14 +169,15 @@ pub trait DKGTrait: Debug {
 
     fn aggregate_transcripts(
         params: &Self::PublicParams,
-        transcripts: Vec<Self::Transcript>,
-    ) -> Self::Transcript;
+        accumulator: &mut Self::Transcript,
+        element: Self::Transcript,
+    );
     fn decrypt_secret_share_from_transcript(
         pub_params: &Self::PublicParams,
         trx: &Self::Transcript,
         player_idx: u64,
         dk: &Self::NewValidatorDecryptKey,
-    ) -> Result<Self::DealtSecretShare>;
+    ) -> Result<(Self::DealtSecretShare, Self::DealtPubKeyShare)>;
     fn reconstruct_secret_from_shares(
         pub_params: &Self::PublicParams,
         player_share_pairs: Vec<(u64, Self::DealtSecretShare)>,
@@ -136,6 +186,6 @@ pub trait DKGTrait: Debug {
 }
 
 pub mod dummy_dkg;
+pub mod real_dkg;
 
-// TODO: replace with RealDKG.
-pub type DefaultDKG = DummyDKG;
+pub type DefaultDKG = RealDKG;
diff --git a/types/src/dkg/real_dkg/mod.rs b/types/src/dkg/real_dkg/mod.rs
diff --git a/types/src/dkg/real_dkg/rounding/mod.rs b/types/src/dkg/real_dkg/rounding/mod.rs
diff --git a/types/src/dkg/real_dkg/rounding/tests.rs b/types/src/dkg/real_dkg/rounding/tests.rs
