Improve shard migration (#2866)

* removing force-migrate-shard and improving migrate-shard

* improving mrenclave command

* fixing ci

* putting back comment

* cleaning up mock

* update cli doc

---------

Co-authored-by: Kailai Wang <Kailai.Wang@hotmail.com>

bitacross-worker/core-primitives/enclave-api/ffi/src/lib.rs

@@ -235,14 +235,6 @@ extern "C" {
 
 	// litentry
 	pub fn migrate_shard(
-		eid: sgx_enclave_id_t,
-		retval: *mut sgx_status_t,
-		old_shard: *const u8,
-		new_shard: *const u8,
-		shard_size: u32,
-	) -> sgx_status_t;
-
-	pub fn force_migrate_shard(
 		eid: sgx_enclave_id_t,
 		retval: *mut sgx_status_t,
 		new_shard: *const u8,

bitacross-worker/core-primitives/enclave-api/src/enclave_base.rs

@@ -79,9 +79,7 @@ pub trait EnclaveBase: Send + Sync + 'static {
 	fn get_fingerprint(&self) -> EnclaveResult<EnclaveFingerprint>;
 
 	// litentry
-	fn migrate_shard(&self, old_shard: Vec<u8>, new_shard: Vec<u8>) -> EnclaveResult<()>;
-
-	fn force_migrate_shard(&self, new_shard: Vec<u8>) -> EnclaveResult<()>;
+	fn migrate_shard(&self, new_shard: Vec<u8>) -> EnclaveResult<()>;
 
 	/// Publish generated wallets on parachain
 	fn publish_wallets(&self) -> EnclaveResult<()>;
@@ -393,30 +391,11 @@ mod impl_ffi {
 			Ok(mr_enclave.into())
 		}
 
-		fn migrate_shard(&self, old_shard: Vec<u8>, new_shard: Vec<u8>) -> EnclaveResult<()> {
+		fn migrate_shard(&self, new_shard: Vec<u8>) -> EnclaveResult<()> {
 			let mut retval = sgx_status_t::SGX_SUCCESS;
 
 			let result = unsafe {
 				ffi::migrate_shard(
-					self.eid,
-					&mut retval,
-					old_shard.as_ptr(),
-					new_shard.as_ptr(),
-					old_shard.len() as u32,
-				)
-			};
-
-			ensure!(result == sgx_status_t::SGX_SUCCESS, Error::Sgx(result));
-			ensure!(retval == sgx_status_t::SGX_SUCCESS, Error::Sgx(retval));
-
-			Ok(())
-		}
-
-		fn force_migrate_shard(&self, new_shard: Vec<u8>) -> EnclaveResult<()> {
-			let mut retval = sgx_status_t::SGX_SUCCESS;
-
-			let result = unsafe {
-				ffi::force_migrate_shard(
 					self.eid,
 					&mut retval,
 					new_shard.as_ptr(),

bitacross-worker/core-primitives/stf-state-handler/src/handle_state.rs

@@ -74,12 +74,5 @@ pub trait HandleState {
 	fn reset(&self, state: Self::StateT, shard: &ShardIdentifier) -> Result<Self::HashType>;
 
 	// litentry
-	/// Migrate state from old shard to new shard
-	fn migrate_shard(
-		&self,
-		old_shard: ShardIdentifier,
-		new_shard: ShardIdentifier,
-	) -> Result<Self::HashType>;
-
-	fn force_migrate_shard(&self, new_shard: ShardIdentifier) -> Result<Self::HashType>;
+	fn migrate_shard(&self, new_shard: ShardIdentifier) -> Result<Self::HashType>;
 }

bitacross-worker/core-primitives/stf-state-handler/src/state_handler.rs

@@ -228,16 +228,7 @@ where
 		self.write_after_mutation(state, state_write_lock, shard)
 	}
 
-	fn migrate_shard(
-		&self,
-		old_shard: ShardIdentifier,
-		new_shard: ShardIdentifier,
-	) -> Result<Self::HashType> {
-		let (state, _) = self.load_cloned(&old_shard)?;
-		self.reset(state, &new_shard)
-	}
-
-	fn force_migrate_shard(&self, new_shard: ShardIdentifier) -> Result<Self::HashType> {
+	fn migrate_shard(&self, new_shard: ShardIdentifier) -> Result<Self::HashType> {
 		if self.shard_exists(&new_shard)? {
 			let (_, state_hash) = self.load_cloned(&new_shard)?;
 			return Ok(state_hash)
@@ -246,10 +237,11 @@ where
 			shards if shards.len() == 1 => shards[0],
 			_ =>
 				return Err(Error::Other(
-					"Cannot force migrate shard if there is more than one shard".into(),
+					"Cannot migrate shard if there is more than one shard".into(),
 				)),
 		};
-		self.migrate_shard(old_shard, new_shard)
+		let (state, _) = self.load_cloned(&old_shard)?;
+		self.reset(state, &new_shard)
 	}
 }
 

bitacross-worker/core-primitives/test/src/mock/handle_state_mock.rs

@@ -56,20 +56,12 @@ impl HandleState for HandleStateMock {
 		self.reset(StfState::default(), &shard)
 	}
 
-	fn migrate_shard(
-		&self,
-		old_shard: ShardIdentifier,
-		new_shard: ShardIdentifier,
-	) -> Result<Self::HashType> {
+	fn migrate_shard(&self, new_shard: ShardIdentifier) -> Result<Self::HashType> {
+		let old_shard = *self.state_map.read().unwrap().keys().next().unwrap();
 		let (state, _) = self.load_cloned(&old_shard)?;
 		self.reset(state, &new_shard)
 	}
 
-	fn force_migrate_shard(&self, new_shard: ShardIdentifier) -> Result<Self::HashType> {
-		let old_shard = *self.state_map.read().unwrap().keys().next().unwrap();
-		self.migrate_shard(old_shard, new_shard)
-	}
-
 	fn execute_on_current<E, R>(&self, shard: &ShardIdentifier, executing_function: E) -> Result<R>
 	where
 		E: FnOnce(&Self::StateT, Self::HashType) -> R,
@@ -233,27 +225,7 @@ pub mod tests {
 		state.insert(key.encode(), value.encode());
 		state_handler.write_after_mutation(state, lock, &old_shard).unwrap();
 
-		state_handler.migrate_shard(old_shard, new_shard).unwrap();
-		let (new_state, _) = state_handler.load_cloned(&new_shard).unwrap();
-		let inserted_value =
-			new_state.get(key.encode().as_slice()).expect("value for key should exist");
-		assert_eq!(*inserted_value, value.encode());
-	}
-
-	pub fn force_migrate_shard_works() {
-		let state_handler = HandleStateMock::default();
-		let old_shard = ShardIdentifier::default();
-		let bytes = hex::decode("91de6f606be264f089b155256385470f5395969386894ffba38775442f508ee2")
-			.unwrap();
-		let new_shard = ShardIdentifier::from_slice(&bytes);
-		state_handler.initialize_shard(old_shard).unwrap();
-
-		let (lock, mut state) = state_handler.load_for_mutation(&old_shard).unwrap();
-		let (key, value) = ("my_key", "my_value");
-		state.insert(key.encode(), value.encode());
-		state_handler.write_after_mutation(state, lock, &old_shard).unwrap();
-
-		state_handler.force_migrate_shard(new_shard).unwrap();
+		state_handler.migrate_shard(new_shard).unwrap();
 		let (new_state, _) = state_handler.load_cloned(&new_shard).unwrap();
 		let inserted_value =
 			new_state.get(key.encode().as_slice()).expect("value for key should exist");

bitacross-worker/enclave-runtime/Enclave.edl

@@ -173,12 +173,6 @@ enclave {
 		public size_t test_main_entrance();
 
 		public sgx_status_t migrate_shard(
-			[in, size=shard_size] uint8_t* old_shard,
-			[in, size=shard_size] uint8_t* new_shard,
-			uint32_t shard_size
-		);
-
-		public sgx_status_t force_migrate_shard(
 			[in, size=shard_size] uint8_t* new_shard,
 			uint32_t shard_size
 		);

bitacross-worker/enclave-runtime/src/initialization/mod.rs

@@ -453,18 +453,9 @@ pub(crate) fn init_shard(shard: ShardIdentifier) -> EnclaveResult<()> {
 	Ok(())
 }
 
-pub(crate) fn migrate_shard(
-	old_shard: ShardIdentifier,
-	new_shard: ShardIdentifier,
-) -> EnclaveResult<()> {
-	let state_handler = GLOBAL_STATE_HANDLER_COMPONENT.get()?;
-	let _ = state_handler.migrate_shard(old_shard, new_shard)?;
-	Ok(())
-}
-
-pub(crate) fn force_migrate_shard(new_shard: ShardIdentifier) -> EnclaveResult<()> {
+pub(crate) fn migrate_shard(new_shard: ShardIdentifier) -> EnclaveResult<()> {
 	let state_handler = GLOBAL_STATE_HANDLER_COMPONENT.get()?;
-	let _ = state_handler.force_migrate_shard(new_shard)?;
+	let _ = state_handler.migrate_shard(new_shard)?;
 	Ok(())
 }
 

bitacross-worker/enclave-runtime/src/lib.rs

@@ -505,35 +505,12 @@ pub unsafe extern "C" fn init_shard(shard: *const u8, shard_size: u32) -> sgx_st
 }
 
 #[no_mangle]
-pub unsafe extern "C" fn migrate_shard(
-	old_shard: *const u8,
-	new_shard: *const u8,
-	shard_size: u32,
-) -> sgx_status_t {
-	let old_shard_identifier =
-		ShardIdentifier::from_slice(slice::from_raw_parts(old_shard, shard_size as usize));
-
-	let new_shard_identifier =
-		ShardIdentifier::from_slice(slice::from_raw_parts(new_shard, shard_size as usize));
-
-	if let Err(e) = initialization::migrate_shard(old_shard_identifier, new_shard_identifier) {
-		error!("Failed to initialize shard ({:?}): {:?}", old_shard_identifier, e);
-		return sgx_status_t::SGX_ERROR_UNEXPECTED
-	}
-
-	sgx_status_t::SGX_SUCCESS
-}
-
-#[no_mangle]
-pub unsafe extern "C" fn force_migrate_shard(
-	new_shard: *const u8,
-	shard_size: u32,
-) -> sgx_status_t {
+pub unsafe extern "C" fn migrate_shard(new_shard: *const u8, shard_size: u32) -> sgx_status_t {
 	let shard_identifier =
 		ShardIdentifier::from_slice(slice::from_raw_parts(new_shard, shard_size as usize));
 
-	if let Err(e) = initialization::force_migrate_shard(shard_identifier) {
-		error!("Failed to (force) migrate shard ({:?}): {:?}", shard_identifier, e);
+	if let Err(e) = initialization::migrate_shard(shard_identifier) {
+		error!("Failed to migrate shard ({:?}): {:?}", shard_identifier, e);
 		return sgx_status_t::SGX_ERROR_UNEXPECTED
 	}
 

bitacross-worker/service/src/cli.yml

@@ -134,10 +134,6 @@ subcommands:
                 long: request-state
                 short: r
                 help: Run the worker and request key and state provisioning from another worker.
-            - force-migrate-shard:
-                long: force-migrate-shard
-                help: Force migrate the shard before starting the worker
-                required: false
     - request-state:
         about: (Deprecated - TODO) join a shard by requesting key provisioning from another worker
         args:
@@ -169,16 +165,7 @@ subcommands:
                 index: 1
                 help: shard identifier base58 encoded
     - migrate-shard:
-        about: Migrate shard
-        args:
-            - old-shard:
-                long: old-shard
-                help: shard identifier hex encoded
-                takes_value: true
-            - new-shard:
-                long: new-shard
-                help: shard identifier hex encoded
-                takes_value: true
+        about: Migrate state from old shards to the new(current) shard, which is identical to mrenclave
     - test:
           about: Run tests involving the enclave
           takes_value: true

bitacross-worker/service/src/config.rs

@@ -288,8 +288,6 @@ pub struct RunConfig {
 	marblerun_base_url: Option<String>,
 	/// parentchain which should be used for shielding/unshielding the stf's native token
 	pub shielding_target: Option<ParentchainId>,
-	/// Whether to migrate the shard before initializing the enclave
-	pub force_migrate_shard: bool,
 }
 
 impl RunConfig {
@@ -334,9 +332,8 @@ impl From<&ArgMatches<'_>> for RunConfig {
 				i
 			),
 		});
-		let force_migrate_shard = m.is_present("force-migrate-shard");
 
-		Self { skip_ra, dev, shard, marblerun_base_url, shielding_target, force_migrate_shard }
+		Self { skip_ra, dev, shard, marblerun_base_url, shielding_target }
 	}
 }
 

bitacross-worker/service/src/main_impl.rs

@@ -164,10 +164,6 @@ pub(crate) fn main() {
 
 		if clean_reset {
 			setup::initialize_shard_and_keys(enclave.as_ref(), &shard).unwrap();
-		} else if run_config.force_migrate_shard {
-			setup::force_migrate_shard(enclave.as_ref(), &shard);
-			let new_shard_name = shard.encode().to_base58();
-			setup::remove_old_shards(config.data_dir(), &new_shard_name);
 		}
 
 		let node_api =
@@ -212,7 +208,10 @@ pub(crate) fn main() {
 			enclave.dump_dcap_ra_cert_to_disk().unwrap();
 		}
 	} else if matches.is_present("mrenclave") {
-		println!("{}", enclave.get_fingerprint().unwrap().encode().to_base58());
+		let mrenclave = enclave.get_fingerprint().unwrap();
+		let hex_value = hex::encode(mrenclave);
+		println!("MRENCLAVE hex: {}", hex_value);
+		println!("MRENCLAVE base58: {}", mrenclave.encode().to_base58());
 	} else if let Some(sub_matches) = matches.subcommand_matches("init-shard") {
 		setup::init_shard(
 			enclave.as_ref(),
@@ -246,32 +245,10 @@ pub(crate) fn main() {
 			tests::run_enclave_tests(sub_matches);
 		}
 	} else if let Some(sub_matches) = matches.subcommand_matches("migrate-shard") {
-		// This subcommand `migrate-shard` is only used for manual testing. Maybe deleted later.
-		let old_shard = sub_matches
-			.value_of("old-shard")
-			.map(|value| {
-				let mut shard = [0u8; 32];
-				hex::decode_to_slice(value, &mut shard)
-					.expect("shard must be hex encoded without 0x");
-				ShardIdentifier::from_slice(&shard)
-			})
-			.unwrap();
-
-		let new_shard: ShardIdentifier = sub_matches
-			.value_of("new-shard")
-			.map(|value| {
-				let mut shard = [0u8; 32];
-				hex::decode_to_slice(value, &mut shard)
-					.expect("shard must be hex encoded without 0x");
-				ShardIdentifier::from_slice(&shard)
-			})
-			.unwrap();
-
-		if old_shard == new_shard {
-			println!("old_shard should not be the same as new_shard");
-		} else {
-			setup::migrate_shard(enclave.as_ref(), &old_shard, &new_shard);
-		}
+		let new_shard = extract_shard(None, enclave.as_ref());
+		setup::migrate_shard(enclave.as_ref(), &new_shard);
+		let new_shard_name = new_shard.encode().to_base58();
+		setup::remove_old_shards(config.data_dir(), &new_shard_name);
 	} else if let Some(sub_matches) = matches.subcommand_matches("wallet") {
 		println!("Bitcoin wallet:");
 		let bitcoin_keypair = enclave.get_bitcoin_wallet_pair().unwrap();

bitacross-worker/service/src/setup.rs

@@ -26,8 +26,8 @@ use std::{fs, path::Path};
 
 #[cfg(feature = "link-binary")]
 pub(crate) use needs_enclave::{
-	force_migrate_shard, generate_shielding_key_file, generate_signing_key_file, init_shard,
-	initialize_shard_and_keys, migrate_shard,
+	generate_shielding_key_file, generate_signing_key_file, init_shard, initialize_shard_and_keys,
+	migrate_shard,
 };
 
 #[cfg(feature = "link-binary")]
@@ -76,34 +76,13 @@ mod needs_enclave {
 		}
 	}
 
-	pub(crate) fn migrate_shard(
-		enclave: &Enclave,
-		old_shard: &ShardIdentifier,
-		new_shard: &ShardIdentifier,
-	) {
-		match enclave.migrate_shard(old_shard.encode(), new_shard.encode()) {
-			Err(e) => {
-				println!(
-					"Failed to migrate old shard {:?} to new shard{:?}. {:?}",
-					old_shard, new_shard, e
-				);
-			},
-			Ok(_) => {
-				println!(
-					"Successfully migrate old shard {:?} to new shard{:?}",
-					old_shard, new_shard
-				);
-			},
-		}
-	}
-
-	pub(crate) fn force_migrate_shard(enclave: &Enclave, &new_shard: &ShardIdentifier) {
-		match enclave.force_migrate_shard(new_shard.encode()) {
+	pub(crate) fn migrate_shard(enclave: &Enclave, &new_shard: &ShardIdentifier) {
+		match enclave.migrate_shard(new_shard.encode()) {
 			Err(e) => {
-				panic!("Failed to force migrate shard {:?}. {:?}", new_shard, e);
+				panic!("Failed to migrate shard {:?}. {:?}", new_shard, e);
 			},
 			Ok(_) => {
-				println!("Successfully force migrate shard {:?}", new_shard);
+				println!("Shard {:?} migrated Successfully", new_shard);
 			},
 		}
 	}