add tests for current safe mode code (MystenLabs#10423)

## Description 

Add a simple safe mode test using similar trick as
`sui_system_injection` to inject epoch change failure.

## Test Plan 

Added a test.

---
If your changes are not user-facing and not a breaking change, you can
skip the following section. Otherwise, please indicate what changed, and
then add to the Release Notes section as highlighted during the release
process.

### Type of Change (Check all that apply)

- [ ] user-visible impact
- [ ] breaking change for a client SDKs
- [ ] breaking change for FNs (FN binary must upgrade)
- [ ] breaking change for validators or node operators (must upgrade
binaries)
- [ ] breaking change for on-chain data layout
- [ ] necessitate either a data wipe or data migration

### Release notes

crates/sui-adapter/src/execution_engine.rs

@@ -46,6 +46,9 @@ use sui_types::{
 
 use sui_types::temporary_store::TemporaryStore;
 
+#[cfg(msim)]
+use self::advance_epoch_result_injection::maybe_modify_result;
+
 checked_arithmetic! {
 
 #[instrument(name = "tx_execute_to_effects", level = "debug", skip_all)]
@@ -563,6 +566,9 @@ fn advance_epoch<S: ObjectStore + BackingPackageStore + ParentSync + ChildObject
         advance_epoch_pt,
     );
 
+    #[cfg(msim)]
+    let result = maybe_modify_result(result);
+
     if result.is_err() {
         tracing::error!(
             "Failed to execute advance epoch transaction. Switching to safe mode. Error: {:?}. Input objects: {:?}. Tx data: {:?}",
@@ -671,3 +677,30 @@ fn setup_consensus_commit<S: BackingPackageStore + ParentSync + ChildObjectResol
 }
 
 }
+
+#[cfg(msim)]
+pub mod advance_epoch_result_injection {
+    use std::cell::RefCell;
+    use sui_types::error::{ExecutionError, ExecutionErrorKind};
+
+    thread_local! {
+        static OVERRIDE: RefCell<bool>  = RefCell::new(false);
+    }
+
+    pub fn set_override(value: bool) {
+        OVERRIDE.with(|o| *o.borrow_mut() = value);
+    }
+
+    /// This function is used to modify the result of advance_epoch transaction for testing.
+    /// If the override is set, the result will be an execution error, otherwise the original result will be returned.
+    pub fn maybe_modify_result(result: Result<(), ExecutionError>) -> Result<(), ExecutionError> {
+        if OVERRIDE.with(|o| *o.borrow()) {
+            Err::<(), ExecutionError>(ExecutionError::new(
+                ExecutionErrorKind::FunctionNotFound,
+                None,
+            ))
+        } else {
+            result
+        }
+    }
+}

crates/sui-benchmark/src/workloads/delegation.rs

@@ -51,7 +51,7 @@ impl Payload for DelegationTestPayload {
                 coin,
                 self.validator,
                 self.sender,
-                &self.keypair,
+                self.keypair.as_ref(),
                 Some(
                     self.system_state_observer
                         .state

crates/sui/Cargo.toml

@@ -28,6 +28,7 @@ rand = "0.8.5"
 tap = "1.0"
 inquire = "0.6.0"
 
+sui-adapter = { path = "../sui-adapter" }
 sui-core = { path = "../sui-core" }
 sui-framework = { path = "../sui-framework" }
 sui-framework-build = { path = "../sui-framework-build" }

crates/sui/tests/reconfiguration_tests.rs

@@ -1059,6 +1059,74 @@ async fn test_reconfig_with_committee_change_stress() {
     assert_eq!(epoch, 3);
 }
 
+#[cfg(msim)]
+#[sim_test]
+async fn safe_mode_reconfig_test() {
+    use sui_adapter::execution_engine::advance_epoch_result_injection;
+    use test_utils::messages::make_staking_transaction_with_wallet_context;
+
+    let mut test_cluster = TestClusterBuilder::new()
+        .with_epoch_duration_ms(5000)
+        .build()
+        .await
+        .unwrap();
+
+    let system_state = test_cluster
+        .sui_client()
+        .governance_api()
+        .get_latest_sui_system_state()
+        .await
+        .unwrap();
+
+    // On startup, we should be at V1.
+    assert_eq!(system_state.system_state_version, 1);
+    assert_eq!(system_state.epoch, 0);
+
+    // Wait for regular epoch change to happen once. Migration from V1 to V2 should happen here.
+    let system_state = test_cluster.wait_for_epoch(Some(1)).await;
+    assert!(!system_state.safe_mode());
+    assert_eq!(system_state.epoch(), 1);
+    assert_eq!(system_state.system_state_version(), 2);
+
+    let prev_epoch_start_timestamp = system_state.epoch_start_timestamp_ms();
+
+    // We are going to enter safe mode so set the expectation right.
+    test_cluster.set_safe_mode_expected(true);
+
+    // Now inject an error into epoch change txn execution.
+    advance_epoch_result_injection::set_override(true);
+
+    // Reconfig again and check that we are in safe mode now.
+    let system_state = test_cluster.wait_for_epoch(Some(2)).await;
+    assert!(system_state.safe_mode());
+    assert_eq!(system_state.epoch(), 2);
+    // Check that time is properly set even in safe mode.
+    assert!(system_state.epoch_start_timestamp_ms() >= prev_epoch_start_timestamp + 5000);
+
+    // Try a staking transaction.
+    let validator_address = system_state
+        .into_sui_system_state_summary()
+        .active_validators[0]
+        .sui_address;
+    let txn =
+        make_staking_transaction_with_wallet_context(test_cluster.wallet_mut(), validator_address)
+            .await;
+    let response = test_cluster
+        .execute_transaction(txn)
+        .await
+        .expect("Staking txn failed");
+    assert!(response.status_ok().unwrap());
+
+    // Now remove the override and check that in the next epoch we are no longer in safe mode.
+    test_cluster.set_safe_mode_expected(false);
+    advance_epoch_result_injection::set_override(false);
+
+    let system_state = test_cluster.wait_for_epoch(Some(3)).await;
+    assert!(!system_state.safe_mode());
+    assert_eq!(system_state.epoch(), 3);
+    assert_eq!(system_state.system_state_version(), 2);
+}
+
 async fn execute_add_validator_candidate_tx(
     authorities: &[SuiNodeHandle],
     gas_object: ObjectRef,

crates/test-utils/src/messages.rs

@@ -18,7 +18,7 @@ use sui_types::base_types::{ObjectID, SequenceNumber};
 use sui_types::committee::Committee;
 use sui_types::crypto::{
     deterministic_random_account_key, get_key_pair, AccountKeyPair, AuthorityKeyPair,
-    AuthorityPublicKeyBytes, AuthoritySignInfo, KeypairTraits,
+    AuthorityPublicKeyBytes, AuthoritySignInfo, KeypairTraits, Signature, Signer,
 };
 use sui_types::gas_coin::GasCoin;
 use sui_types::messages::SignedTransactionEffects;
@@ -215,6 +215,39 @@ pub async fn make_transactions_with_wallet_context_and_budget(
     res
 }
 
+pub async fn make_staking_transaction_with_wallet_context(
+    context: &mut WalletContext,
+    validator_address: SuiAddress,
+) -> VerifiedTransaction {
+    let accounts_and_objs = get_account_and_gas_objects(context).await;
+    let sender = accounts_and_objs[0].0;
+    let gas_object = accounts_and_objs[0].1[0]
+        .clone()
+        .into_object()
+        .expect("Gas coin could not be converted to object ref.")
+        .object_ref();
+    let stake_object = accounts_and_objs[0].1[1]
+        .clone()
+        .into_object()
+        .expect("Stake coin could not be converted to object ref.")
+        .object_ref();
+    let client = context.get_client().await.unwrap();
+    let gas_price = client
+        .governance_api()
+        .get_reference_gas_price()
+        .await
+        .unwrap();
+
+    make_staking_transaction(
+        gas_object,
+        stake_object,
+        validator_address,
+        sender,
+        context.config.keystore.get_key(&sender).unwrap(),
+        Some(gas_price),
+    )
+}
+
 pub async fn make_counter_increment_transaction_with_wallet_context(
     context: &WalletContext,
     sender: SuiAddress,
@@ -440,7 +473,7 @@ pub fn make_staking_transaction(
     coin: ObjectRef,
     validator: SuiAddress,
     sender: SuiAddress,
-    keypair: &AccountKeyPair,
+    keypair: &dyn Signer<Signature>,
     gas_price: Option<u64>,
 ) -> VerifiedTransaction {
     let data = TransactionData::new_move_call(

crates/test-utils/src/network.rs

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use futures::future::join_all;
-use move_core_types::ident_str;
 use std::num::NonZeroUsize;
 use std::sync::Arc;
 use std::time::Duration;
@@ -16,7 +15,6 @@ use tokio::time::timeout;
 use tokio::{task::JoinHandle, time::sleep};
 use tracing::info;
 
-use crate::messages::get_sui_gas_object_with_wallet_context;
 use mysten_metrics::RegistryService;
 use sui::config::SuiEnv;
 use sui::{client_commands::WalletContext, config::SuiClientConfig};
@@ -25,8 +23,6 @@ use sui_config::genesis_config::GenesisConfig;
 use sui_config::node::DBCheckpointConfig;
 use sui_config::{Config, SUI_CLIENT_CONFIG, SUI_NETWORK_CONFIG};
 use sui_config::{FullnodeConfigBuilder, NodeConfig, PersistedConfig, SUI_KEYSTORE_FILENAME};
-use sui_core::test_utils::MAX_GAS;
-use sui_framework::{SuiSystem, SystemPackage};
 use sui_json_rpc_types::{SuiTransactionBlockResponse, SuiTransactionBlockResponseOptions};
 use sui_keys::keystore::{AccountKeystore, FileBasedKeystore, Keystore};
 use sui_node::SuiNode;
@@ -37,17 +33,12 @@ use sui_sdk::{SuiClient, SuiClientBuilder};
 use sui_swarm::memory::{Swarm, SwarmBuilder};
 use sui_types::base_types::{AuthorityName, SuiAddress};
 use sui_types::committee::EpochId;
+use sui_types::crypto::KeypairTraits;
 use sui_types::crypto::SuiKeyPair;
-use sui_types::crypto::{AccountKeyPair, KeypairTraits};
-use sui_types::messages::CallArg;
-use sui_types::messages::ObjectArg;
-use sui_types::messages::TransactionData;
+use sui_types::messages::VerifiedTransaction;
 use sui_types::object::Object;
 use sui_types::sui_system_state::SuiSystemState;
 use sui_types::sui_system_state::SuiSystemStateTrait;
-use sui_types::utils::to_sender_signed_transaction;
-use sui_types::SUI_SYSTEM_STATE_OBJECT_ID;
-use sui_types::SUI_SYSTEM_STATE_OBJECT_SHARED_VERSION;
 
 const NUM_VALIDAOTR: usize = 4;
 
@@ -174,38 +165,18 @@ impl TestCluster {
         .expect("Timed out waiting for cluster to target epoch")
     }
 
-    pub async fn request_add_stake(
+    pub async fn execute_transaction(
         &self,
-        validator_address: SuiAddress,
-        sender: SuiAddress,
-        sender_keypair: &AccountKeyPair,
+        transaction: VerifiedTransaction,
     ) -> SuiRpcResult<SuiTransactionBlockResponse> {
-        let objects = get_sui_gas_object_with_wallet_context(&self.wallet, &sender).await;
-        let stake_tx_data = TransactionData::new_move_call_with_dummy_gas_price(
-            sender,
-            SuiSystem::ID,
-            ident_str!("sui_system").to_owned(),
-            ident_str!("request_add_stake").to_owned(),
-            vec![],
-            objects[0].1,
-            vec![
-                CallArg::Object(ObjectArg::SharedObject {
-                    id: SUI_SYSTEM_STATE_OBJECT_ID,
-                    initial_shared_version: SUI_SYSTEM_STATE_OBJECT_SHARED_VERSION,
-                    mutable: true,
-                }),
-                CallArg::Object(ObjectArg::ImmOrOwnedObject(objects[1].1)),
-                CallArg::Pure(bcs::to_bytes(&validator_address).unwrap()),
-            ],
-            MAX_GAS,
-        )
-        .unwrap();
-        let transaction = to_sender_signed_transaction(stake_tx_data, sender_keypair);
-
         self.fullnode_handle
             .sui_client
             .quorum_driver()
-            .execute_transaction_block(transaction, SuiTransactionBlockResponseOptions::new(), None)
+            .execute_transaction_block(
+                transaction,
+                SuiTransactionBlockResponseOptions::new().with_effects(),
+                None,
+            )
             .await
     }