Enclave update [4-1]: SGX-Runtime migration (#1303)

* runtime-migration first commit

* minor update

* remove comment out line

---------

Co-authored-by: ericzhang <jingleizhang@users.noreply.github.com>

Author: zhizming-zhong

tee-worker/Cargo.lock

@@ -31,6 +31,8 @@
 #[cfg(feature = "evm")]
 mod evm;
 
+mod migration;
+
 #[cfg(feature = "evm")]
 pub use evm::{
 	AddressMapping, EnsureAddressTruncated, EvmCall, FeeCalculator, FixedGasPrice,
@@ -105,13 +107,15 @@ pub type UncheckedExtrinsic =
 	generic::UncheckedExtrinsic<Address, RuntimeCall, Signature, SignedExtra>;
 /// Extrinsic type that has already been checked.
 pub type CheckedExtrinsic = generic::CheckedExtrinsic<AccountId, RuntimeCall, SignedExtra>;
+
 /// Executive: handles dispatch to the various modules.
 pub type Executive = frame_executive::Executive<
 	Runtime,
 	Block,
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPalletsWithSystem,
+	migration::Upgrade,
 >;
 
 /// Opaque types. These are used by the CLI to instantiate machinery that don't need to know

tee-worker/app-libs/sgx-runtime/src/lib.rs

@@ -0,0 +1,14 @@
+use super::*;
+
+use crate::{pallet_imt::migrations, Runtime};
+use frame_support::{traits::OnRuntimeUpgrade, weights::Weight};
+
+// This is just an example of how to write a custom migration
+// It has no effect on tee-worker running
+// For more details, see: https://docs.substrate.io/maintain/runtime-upgrades/#storage-migration
+pub struct Upgrade;
+impl OnRuntimeUpgrade for Upgrade {
+	fn on_runtime_upgrade() -> Weight {
+		migrations::migrate_to_v1::<Runtime, IdentityManagement>()
+	}
+}

tee-worker/app-libs/sgx-runtime/src/migration.rs

@@ -21,12 +21,14 @@ use crate::test_genesis::test_genesis_setup;
 use crate::{helpers::enclave_signer_account, Stf, StfError, ENCLAVE_ACCOUNT_KEY};
 use codec::Encode;
 use frame_support::traits::{OriginTrait, UnfilteredDispatchable};
+use ita_sgx_runtime::Executive;
 use itp_node_api::metadata::{
 	pallet_imp::IMPCallIndexes, pallet_teerex::TeerexCallIndexes, provider::AccessNodeMetadata,
 };
 use itp_sgx_externalities::SgxExternalitiesTrait;
 use itp_stf_interface::{
 	parentchain_pallet::ParentchainPalletInterface,
+	runtime_upgrade::RuntimeUpgradeInterface,
 	sudo_pallet::SudoPalletInterface,
 	system_pallet::{SystemPalletAccountInterface, SystemPalletEventInterface},
 	ExecuteCall, ExecuteGetter, InitState, StateCallInterface, StateGetterInterface, UpdateState,
@@ -247,6 +249,42 @@ where
 	}
 }
 
+impl<Call, Getter, State, Runtime> RuntimeUpgradeInterface<State>
+	for Stf<Call, Getter, State, Runtime>
+where
+	State: SgxExternalitiesTrait,
+	Runtime: frame_system::Config,
+{
+	type Error = StfError;
+
+	fn on_runtime_upgrade(state: &mut State) -> Result<(), Self::Error> {
+		// Returns if the runtime was upgraded since the last time this function was called.
+		let runtime_upgraded =
+			|| -> bool {
+				let last = frame_system::LastRuntimeUpgrade::<Runtime>::get();
+				let current = <<Runtime as frame_system::Config>::Version as frame_support::traits::Get<_>>::get();
+
+				if last.map(|v| v.was_upgraded(&current)).unwrap_or(true) {
+					frame_system::LastRuntimeUpgrade::<Runtime>::put(
+						frame_system::LastRuntimeUpgradeInfo::from(current),
+					);
+					debug!("Do some migraions");
+					true
+				} else {
+					debug!("No need to migrate");
+					false
+				}
+			};
+
+		state.execute_with(|| {
+			if runtime_upgraded() {
+				Executive::execute_on_runtime_upgrade();
+			}
+		});
+		Ok(())
+	}
+}
+
 pub fn storage_hashes_to_update_per_shard(_shard: &ShardIdentifier) -> Vec<Vec<u8>> {
 	Vec::new()
 }

tee-worker/app-libs/stf/src/stf_sgx.rs

@@ -11,7 +11,6 @@ sgx_tstd = { branch = "master", git = "https://github.com/apache/teaclave-sgx-sd
 sgx_types = { branch = "master", git = "https://github.com/apache/teaclave-sgx-sdk.git" }
 
 # local dependencies
-ita-sgx-runtime = { path = "../../app-libs/sgx-runtime", default-features = false }
 ita-stf = { path = "../../app-libs/stf", default-features = false }
 itp-node-api = { path = "../node-api", default-features = false }
 itp-ocall-api = { path = "../ocall-api", default-features = false }
@@ -57,7 +56,6 @@ mocks = []
 std = [
     # local
     "ita-stf/std",
-    "ita-sgx-runtime/std",
     "itp-node-api/std",
     "itp-ocall-api/std",
     "itp-sgx-crypto/std",

tee-worker/core-primitives/stf-executor/Cargo.toml

@@ -33,7 +33,8 @@ use itp_node_api::metadata::{
 use itp_ocall_api::{EnclaveAttestationOCallApi, EnclaveOnChainOCallApi};
 use itp_sgx_externalities::{SgxExternalitiesTrait, StateHash};
 use itp_stf_interface::{
-	parentchain_pallet::ParentchainPalletInterface, ExecuteCall, StateCallInterface, UpdateState,
+	parentchain_pallet::ParentchainPalletInterface, runtime_upgrade::RuntimeUpgradeInterface,
+	ExecuteCall, StateCallInterface, UpdateState,
 };
 use itp_stf_primitives::types::ShardIdentifier;
 use itp_stf_state_handler::{handle_state::HandleState, query_shard_state::QueryShardState};
@@ -244,12 +245,13 @@ where
 	Stf: UpdateState<
 			StateHandler::StateT,
 			<StateHandler::StateT as SgxExternalitiesTrait>::SgxExternalitiesDiffType,
-		> + StateCallInterface<TrustedCallSigned, StateHandler::StateT, NodeMetadataRepository>,
+		> + StateCallInterface<TrustedCallSigned, StateHandler::StateT, NodeMetadataRepository> + RuntimeUpgradeInterface<StateHandler::StateT>,
 	<StateHandler::StateT as SgxExternalitiesTrait>::SgxExternalitiesDiffType:
 		IntoIterator<Item = (Vec<u8>, Option<Vec<u8>>)>,
 	<StateHandler::StateT as SgxExternalitiesTrait>::SgxExternalitiesDiffType:
 		From<BTreeMap<Vec<u8>, Option<Vec<u8>>>>,
 	<Stf as StateCallInterface<TrustedCallSigned, StateHandler::StateT, NodeMetadataRepository>>::Error: Debug,
+	<Stf as RuntimeUpgradeInterface<StateHandler::StateT>>::Error: Debug,
 {
 	type Externalities = StateHandler::StateT;
 
@@ -273,6 +275,9 @@ where
 		let mut state = prepare_state_function(state);
 		let mut executed_and_failed_calls = Vec::<ExecutedOperation>::new();
 
+		// TODO: maybe we can move it to `prepare_state_function`. It seems more reasonable.
+		let _ = Stf::on_runtime_upgrade(&mut state);
+
 		// Iterate through all calls until time is over.
 		for trusted_call_signed in trusted_calls.into_iter() {
 			// Break if allowed time window is over.

tee-worker/core-primitives/stf-executor/src/executor.rs

@@ -30,6 +30,7 @@ use itp_types::{OpaqueCall, ShardIdentifier};
 #[cfg(feature = "mocks")]
 pub mod mocks;
 pub mod parentchain_pallet;
+pub mod runtime_upgrade;
 pub mod sudo_pallet;
 pub mod system_pallet;
 

tee-worker/core-primitives/stf-interface/src/lib.rs

@@ -0,0 +1,21 @@
+// Copyright 2020-2023 Litentry Technologies GmbH.
+// This file is part of Litentry.
+//
+// Litentry is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// Litentry is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with Litentry.  If not, see <https://www.gnu.org/licenses/>.
+
+pub trait RuntimeUpgradeInterface<State> {
+	type Error;
+
+	fn on_runtime_upgrade(state: &mut State) -> Result<(), Self::Error>;
+}

tee-worker/core-primitives/stf-interface/src/runtime_upgrade.rs

@@ -409,8 +409,8 @@ fn test_create_state_diff() {
 		get_from_state_diff(&state_diff, &account_key_hash::<AccountId>(&receiver.into()));
 
 	// state diff should consist of the following updates:
-	// (last_hash, sidechain block_number, sender_funds, receiver_funds, [no clear, after polkadot_v0.9.26 update], events)
-	assert_eq!(state_diff.len(), 6);
+	// (last_hash, sidechain block_number, sender_funds, receiver_funds, [no clear, after polkadot_v0.9.26 update], events, frame_system::LastRuntimeUpgradeInfo, pallet_identity_management_tee::migrations)
+	assert_eq!(state_diff.len(), 8);
 	assert_eq!(receiver_acc_info.data.free, 1000);
 	assert_eq!(sender_acc_info.data.free, 1000);
 }