aptos-labs · zjma · Mar 3, 2024 · Mar 2, 2024 · Mar 2, 2024 · Mar 2, 2024
diff --git a/aptos-move/framework/aptos-framework/doc/dkg.md b/aptos-move/framework/aptos-framework/doc/dkg.md
@@ -6,11 +6,17 @@
 DKG on-chain states and helper functions.
 
 
+-  [Resource `FailureInjectionBlockDKG`](#0x1_dkg_FailureInjectionBlockDKG)
+-  [Resource `FailureInjectionBlockRandomness`](#0x1_dkg_FailureInjectionBlockRandomness)
 -  [Struct `DKGSessionMetadata`](#0x1_dkg_DKGSessionMetadata)
 -  [Struct `DKGStartEvent`](#0x1_dkg_DKGStartEvent)
 -  [Struct `DKGSessionState`](#0x1_dkg_DKGSessionState)
 -  [Resource `DKGState`](#0x1_dkg_DKGState)
 -  [Constants](#@Constants_0)
+-  [Function `block_dkg`](#0x1_dkg_block_dkg)
+-  [Function `unblock_dkg`](#0x1_dkg_unblock_dkg)
+-  [Function `block_randomness`](#0x1_dkg_block_randomness)
+-  [Function `unblock_randomness`](#0x1_dkg_unblock_randomness)
 -  [Function `initialize`](#0x1_dkg_initialize)
 -  [Function `start`](#0x1_dkg_start)
 -  [Function `finish`](#0x1_dkg_finish)
@@ -32,6 +38,64 @@ DKG on-chain states and helper functions.
 
 
 
+<a id="0x1_dkg_FailureInjectionBlockDKG"></a>
+
+## Resource `FailureInjectionBlockDKG`
+
+If this resource is present under 0x1, validators should not do DKG (so the epoch change get stuck).
+This is test-only.
+
+
+<pre><code><b>struct</b> <a href="dkg.md#0x1_dkg_FailureInjectionBlockDKG">FailureInjectionBlockDKG</a> <b>has</b> drop, key
+</code></pre>
+
+
+
+<details>
+<summary>Fields</summary>
+
+
+<dl>
+<dt>
+<code>dummy_field: bool</code>
+</dt>
+<dd>
+
+</dd>
+</dl>
+
+
+</details>
+
+<a id="0x1_dkg_FailureInjectionBlockRandomness"></a>
+
+## Resource `FailureInjectionBlockRandomness`
+
+If this resource is present under 0x1, validators should not provider randomness to block (so the execution get stuck).
+This is test-only.
+
+
+<pre><code><b>struct</b> <a href="dkg.md#0x1_dkg_FailureInjectionBlockRandomness">FailureInjectionBlockRandomness</a> <b>has</b> drop, key
+</code></pre>
+
+
+
+<details>
+<summary>Fields</summary>
+
+
+<dl>
+<dt>
+<code>dummy_field: bool</code>
+</dt>
+<dd>
+
+</dd>
+</dl>
+
+
+</details>
+
 <a id="0x1_dkg_DKGSessionMetadata"></a>
 
 ## Struct `DKGSessionMetadata`
@@ -66,6 +130,18 @@ This can be considered as the public input of DKG.
 </dt>
 <dd>
 
+</dd>
+<dt>
+<code>block_dkg: bool</code>
+</dt>
+<dd>
+
+</dd>
+<dt>
+<code>block_randomness: bool</code>
+</dt>
+<dd>
+
 </dd>
 </dl>
 
@@ -204,6 +280,114 @@ The completed and in-progress DKG sessions.
 
 
 
+<a id="0x1_dkg_block_dkg"></a>
+
+## Function `block_dkg`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="dkg.md#0x1_dkg_block_dkg">block_dkg</a>(framework: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>)
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="dkg.md#0x1_dkg_block_dkg">block_dkg</a>(framework: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>) {
+    <a href="system_addresses.md#0x1_system_addresses_assert_aptos_framework">system_addresses::assert_aptos_framework</a>(framework);
+    <b>if</b> (!<b>exists</b>&lt;<a href="dkg.md#0x1_dkg_FailureInjectionBlockDKG">FailureInjectionBlockDKG</a>&gt;(@aptos_framework)) {
+        <b>move_to</b>(framework, <a href="dkg.md#0x1_dkg_FailureInjectionBlockDKG">FailureInjectionBlockDKG</a> {})
+    }
+}
+</code></pre>
+
+
+
+</details>
+
+<a id="0x1_dkg_unblock_dkg"></a>
+
+## Function `unblock_dkg`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="dkg.md#0x1_dkg_unblock_dkg">unblock_dkg</a>(framework: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>)
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="dkg.md#0x1_dkg_unblock_dkg">unblock_dkg</a>(framework: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>) <b>acquires</b> <a href="dkg.md#0x1_dkg_FailureInjectionBlockDKG">FailureInjectionBlockDKG</a> {
+    <a href="system_addresses.md#0x1_system_addresses_assert_aptos_framework">system_addresses::assert_aptos_framework</a>(framework);
+    <b>if</b> (<b>exists</b>&lt;<a href="dkg.md#0x1_dkg_FailureInjectionBlockDKG">FailureInjectionBlockDKG</a>&gt;(@aptos_framework)) {
+        <b>move_from</b>&lt;<a href="dkg.md#0x1_dkg_FailureInjectionBlockDKG">FailureInjectionBlockDKG</a>&gt;(@aptos_framework);
+    }
+}
+</code></pre>
+
+
+
+</details>
+
+<a id="0x1_dkg_block_randomness"></a>
+
+## Function `block_randomness`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="dkg.md#0x1_dkg_block_randomness">block_randomness</a>(framework: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>)
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="dkg.md#0x1_dkg_block_randomness">block_randomness</a>(framework: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>) {
+    <a href="system_addresses.md#0x1_system_addresses_assert_aptos_framework">system_addresses::assert_aptos_framework</a>(framework);
+    <b>if</b> (!<b>exists</b>&lt;<a href="dkg.md#0x1_dkg_FailureInjectionBlockRandomness">FailureInjectionBlockRandomness</a>&gt;(@aptos_framework)) {
+        <b>move_to</b>(framework, <a href="dkg.md#0x1_dkg_FailureInjectionBlockRandomness">FailureInjectionBlockRandomness</a> {})
+    }
+}
+</code></pre>
+
+
+
+</details>
+
+<a id="0x1_dkg_unblock_randomness"></a>
+
+## Function `unblock_randomness`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="dkg.md#0x1_dkg_unblock_randomness">unblock_randomness</a>(framework: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>)
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="dkg.md#0x1_dkg_unblock_randomness">unblock_randomness</a>(framework: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>) <b>acquires</b> <a href="dkg.md#0x1_dkg_FailureInjectionBlockRandomness">FailureInjectionBlockRandomness</a> {
+    <a href="system_addresses.md#0x1_system_addresses_assert_aptos_framework">system_addresses::assert_aptos_framework</a>(framework);
+    <b>if</b> (!<b>exists</b>&lt;<a href="dkg.md#0x1_dkg_FailureInjectionBlockRandomness">FailureInjectionBlockRandomness</a>&gt;(@aptos_framework)) {
+        <b>move_from</b>&lt;<a href="dkg.md#0x1_dkg_FailureInjectionBlockRandomness">FailureInjectionBlockRandomness</a>&gt;(@aptos_framework);
+    }
+}
+</code></pre>
+
+
+
+</details>
+
 <a id="0x1_dkg_initialize"></a>
 
 ## Function `initialize`
@@ -264,6 +448,8 @@ Abort if a DKG is already in progress.
         dealer_epoch,
         dealer_validator_set,
         target_validator_set,
+        block_dkg: <b>exists</b>&lt;<a href="dkg.md#0x1_dkg_FailureInjectionBlockDKG">FailureInjectionBlockDKG</a>&gt;(@aptos_framework),
+        block_randomness: <b>exists</b>&lt;<a href="dkg.md#0x1_dkg_FailureInjectionBlockRandomness">FailureInjectionBlockRandomness</a>&gt;(@aptos_framework),
     };
     <b>let</b> start_time_us = <a href="timestamp.md#0x1_timestamp_now_microseconds">timestamp::now_microseconds</a>();
     dkg_state.in_progress = std::option::some(<a href="dkg.md#0x1_dkg_DKGSessionState">DKGSessionState</a> {

diff --git a/aptos-move/framework/aptos-framework/sources/dkg.move b/aptos-move/framework/aptos-framework/sources/dkg.move
@@ -13,11 +13,21 @@ module aptos_framework::dkg {
     const EDKG_IN_PROGRESS: u64 = 1;
     const EDKG_NOT_IN_PROGRESS: u64 = 2;
 
+    /// If this resource is present under 0x1, validators should not do DKG (so the epoch change get stuck).
+    /// This is test-only.
+    struct FailureInjectionBlockDKG has drop, key {}
+
+    /// If this resource is present under 0x1, validators should not provider randomness to block (so the execution get stuck).
+    /// This is test-only.
+    struct FailureInjectionBlockRandomness has drop, key {}
+
     /// This can be considered as the public input of DKG.
     struct DKGSessionMetadata has copy, drop, store {
         dealer_epoch: u64,
         dealer_validator_set: vector<ValidatorConsensusInfo>,
         target_validator_set: vector<ValidatorConsensusInfo>,
+        block_dkg: bool,
+        block_randomness: bool,
     }
 
     #[event]
@@ -40,6 +50,34 @@ module aptos_framework::dkg {
         in_progress: Option<DKGSessionState>,
     }
 
+    public fun block_dkg(framework: &signer) {
+        system_addresses::assert_aptos_framework(framework);
+        if (!exists<FailureInjectionBlockDKG>(@aptos_framework)) {
+            move_to(framework, FailureInjectionBlockDKG {})
+        }
+    }
+
+    public fun unblock_dkg(framework: &signer) acquires FailureInjectionBlockDKG {
+        system_addresses::assert_aptos_framework(framework);
+        if (exists<FailureInjectionBlockDKG>(@aptos_framework)) {
+            move_from<FailureInjectionBlockDKG>(@aptos_framework);
+        }
+    }
+
+    public fun block_randomness(framework: &signer) {
+        system_addresses::assert_aptos_framework(framework);
+        if (!exists<FailureInjectionBlockRandomness>(@aptos_framework)) {
+            move_to(framework, FailureInjectionBlockRandomness {})
+        }
+    }
+
+    public fun unblock_randomness(framework: &signer) acquires FailureInjectionBlockRandomness {
+        system_addresses::assert_aptos_framework(framework);
+        if (!exists<FailureInjectionBlockRandomness>(@aptos_framework)) {
+            move_from<FailureInjectionBlockRandomness>(@aptos_framework);
+        }
+    }
+
     /// Called in genesis to initialize on-chain states.
     public fun initialize(aptos_framework: &signer) {
         system_addresses::assert_aptos_framework(aptos_framework);
@@ -65,6 +103,8 @@ module aptos_framework::dkg {
             dealer_epoch,
             dealer_validator_set,
             target_validator_set,
+            block_dkg: exists<FailureInjectionBlockDKG>(@aptos_framework),
+            block_randomness: exists<FailureInjectionBlockRandomness>(@aptos_framework),
         };
         let start_time_us = timestamp::now_microseconds();
         dkg_state.in_progress = std::option::some(DKGSessionState {

diff --git a/aptos-move/vm-genesis/src/lib.rs b/aptos-move/vm-genesis/src/lib.rs
@@ -55,6 +55,7 @@ const GOVERNANCE_MODULE_NAME: &str = "aptos_governance";
 const CODE_MODULE_NAME: &str = "code";
 const VERSION_MODULE_NAME: &str = "version";
 const OIDB_MODULE_NAME: &str = "openid_account";
+#[allow(dead_code)]
 const JWKS_MODULE_NAME: &str = "jwks";
 const CONFIG_BUFFER_MODULE_NAME: &str = "config_buffer";
 const DKG_MODULE_NAME: &str = "dkg";
@@ -556,6 +557,7 @@ fn initialize_reconfiguration_state(session: &mut SessionExt) {
     );
 }
 
+#[allow(dead_code)]
 fn initialize_jwks(session: &mut SessionExt) {
     exec_function(
         session,

diff --git a/consensus/src/dag/dag_fetcher.rs b/consensus/src/dag/dag_fetcher.rs
@@ -370,7 +370,6 @@ impl TDagFetcher for DagFetcher {
                         let bounded_executor = bounded_executor.clone();
                         async move {
                             let nodes = response.certified_nodes();
-                            let _epoch = epoch_state.epoch;
                             ensure!(
                                 stream::iter(nodes.clone())
                                     .concurrent_map(bounded_executor.clone(), move |node| {

diff --git a/consensus/src/epoch_manager.rs b/consensus/src/epoch_manager.rs
@@ -972,6 +972,7 @@ impl<P: OnChainConfigProvider> EpochManager<P> {
             vuf_pp,
             keys,
             dkg_pub_params.pvss_config.wconfig.clone(),
+            dkg_session.metadata.block_randomness,
         );
 
         Ok(rand_config)

diff --git a/consensus/src/rand/rand_gen/rand_manager.rs b/consensus/src/rand/rand_gen/rand_manager.rs
@@ -403,9 +403,11 @@ impl<S: TShare, D: TAugmentedData> RandManager<S, D> {
                 },
 
             }
-            let maybe_ready_blocks = self.block_queue.dequeue_rand_ready_prefix();
-            if !maybe_ready_blocks.is_empty() {
-                self.process_ready_blocks(maybe_ready_blocks);
+            if !self.config.block_randomness {
+                let maybe_ready_blocks = self.block_queue.dequeue_rand_ready_prefix();
+                if !maybe_ready_blocks.is_empty() {
+                    self.process_ready_blocks(maybe_ready_blocks);
+                }
             }
         }
         info!("RandManager stopped");

diff --git a/consensus/src/rand/rand_gen/rand_store.rs b/consensus/src/rand/rand_gen/rand_store.rs
@@ -345,6 +345,8 @@ mod tests {
                 dealer_epoch: 999,
                 dealer_validator_set: consensus_info_move_structs.clone(),
                 target_validator_set: consensus_info_move_structs.clone(),
+                block_dkg: false,
+                block_randomness: false,
             };
             let dkg_pub_params = DefaultDKG::new_public_params(&dkg_session_metadata);
             let input_secret = <DefaultDKG as DKGTrait>::InputSecret::generate_for_testing();
@@ -384,6 +386,7 @@ mod tests {
                 vuf_pub_params,
                 rand_keys,
                 weighted_config,
+                false,
             );
 
             Self {

diff --git a/consensus/src/rand/rand_gen/types.rs b/consensus/src/rand/rand_gen/types.rs
@@ -459,6 +459,8 @@ pub struct RandConfig {
     pub keys: Arc<RandKeys>,
     // weighted config for weighted VUF
     pub wconfig: WeightedConfig,
+    /// If set, do not perform dequeue (so execution is blocked). This is test-only.
+    pub block_randomness: bool,
 }
 
 impl Debug for RandConfig {
@@ -479,6 +481,7 @@ impl RandConfig {
         vuf_pp: WvufPP,
         keys: RandKeys,
         wconfig: WeightedConfig,
+        block_randomness: bool,
     ) -> Self {
         Self {
             author,
@@ -487,6 +490,7 @@ impl RandConfig {
             vuf_pp,
             keys: Arc::new(keys),
             wconfig,
+            block_randomness,
         }
     }
 

diff --git a/dkg/src/dkg_manager/mod.rs b/dkg/src/dkg_manager/mod.rs
@@ -292,6 +292,9 @@ impl<DKG: DKGTrait> DKGManager<DKG> {
             matches!(&self.state, InnerState::NotStarted),
             "transcript already dealt"
         );
+        if dkg_session_metadata.block_dkg {
+            bail!("DKG aborted per config");
+        }
         let dkg_start_time = Duration::from_micros(start_time_us);
         let deal_start = duration_since_epoch();
         let secs_since_dkg_start = deal_start.as_secs_f64() - dkg_start_time.as_secs_f64();

diff --git a/dkg/src/dkg_manager/tests.rs b/dkg/src/dkg_manager/tests.rs
@@ -86,6 +86,8 @@ async fn test_dkg_state_transition() {
             dealer_epoch: 999,
             dealer_validator_set: validator_consensus_info_move_structs.clone(),
             target_validator_set: validator_consensus_info_move_structs.clone(),
+            block_dkg: false,
+            block_randomness: false,
         },
         start_time_us: start_time_1.as_micros() as u64,
     };
-Original file line number
+Diff line change
@@ Expand Up / @@ -972,6 +972,7 @@ impl<P: OnChainConfigProvider> EpochManager<P> { @@
                 vuf_pp,
                 keys,
                 dkg_pub_params.pvss_config.wconfig.clone(),
+                dkg_session.metadata.block_randomness,
             );
             Ok(rand_config)
@@ Expand Down @@