protect asset_registry migration

centrifuge · Oct 17, 2023 · cfad34e · cfad34e
1 parent 5e80739
commit cfad34e
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 57 deletions.
diff --git a/runtime/altair/src/migrations.rs b/runtime/altair/src/migrations.rs
@@ -62,31 +62,9 @@ mod asset_registry {
 	use sp_std::{vec, vec::Vec};
 	use xcm::{v3::prelude::*, VersionedMultiLocation};
 
-	pub const ALTAIR_ASSET_LOC_COUNT: u32 = 5;
-	pub const ALTAIR_ASSET_METADATA_COUNT: u32 = 5;
-
 	pub struct AltairAssets;
 	impl runtime_common::migrations::asset_registry_xcmv3::AssetsToMigrate for AltairAssets {
-		fn get_assets_to_migrate(
-			loc_count: u32,
-			meta_count: u32,
-		) -> Vec<(
-			CurrencyId,
-			orml_asset_registry::AssetMetadata<Balance, CustomMetadata>,
-		)> {
-			match (loc_count, meta_count) {
-				(loc, meta)
-					if (loc, meta) == (ALTAIR_ASSET_LOC_COUNT, ALTAIR_ASSET_METADATA_COUNT) =>
-				{
-					Self::get_altair_assets()
-				}
-				_ => vec![],
-			}
-		}
-	}
-
-	impl AltairAssets {
-		pub fn get_altair_assets() -> Vec<(
+		fn get_assets_to_migrate() -> Vec<(
 			CurrencyId,
 			orml_asset_registry::AssetMetadata<Balance, CustomMetadata>,
 		)> {

diff --git a/runtime/common/Cargo.toml b/runtime/common/Cargo.toml
@@ -76,7 +76,7 @@ log = "0.4"
 sp-io = { git = "https://github.com/paritytech/substrate", default-features = false, branch = "polkadot-v0.9.38" }
 
 [dev-dependencies]
-cfg-mocks = { path = "../../libs/mocks", features = ["runtime-benchmarks", "std"] }
+cfg-mocks = { path = "../../libs/mocks" }
 pallet-collective = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.38" }
 sp-io = { git = "https://github.com/paritytech/substrate", default-features = true, branch = "polkadot-v0.9.38" }
 
@@ -157,6 +157,7 @@ runtime-benchmarks = [
   "xcm-executor/runtime-benchmarks",
   "orml-asset-registry/runtime-benchmarks",
   "polkadot-parachain/runtime-benchmarks",
+  "cfg-mocks/runtime-benchmarks",
 ]
 
 on-chain-release-build = [
@@ -190,4 +191,5 @@ try-runtime = [
   "pallet-investments/try-runtime",
   "parachain-info/try-runtime",
   "orml-asset-registry/try-runtime",
+  "cfg-mocks/try-runtime",
 ]
diff --git a/runtime/common/src/migrations/asset_registry_xcmv3.rs b/runtime/common/src/migrations/asset_registry_xcmv3.rs
@@ -65,6 +65,10 @@ impl<
 		// Complexity: 2 reads
 		let (loc_count, meta_count) = Self::get_key_counts();
 
+		if Self::check_key_counts(loc_count, meta_count).is_err() {
+			return RocksDbWeight::get().reads(loc_count.saturating_add(meta_count).into());
+		}
+
 		// Complexity: O(loc_count) writes
 		let result = orml_asset_registry::LocationToAssetId::<T>::clear(loc_count, None);
 		match result.maybe_cursor {
@@ -96,28 +100,26 @@ impl<
             result.loops,
         );
 
+		let assets = Assets::get_assets_to_migrate();
 		log::info!(
 			"💎 AssetRegistryMultilocationToXCMV3: Starting migration of {:?} assets",
-			Assets::get_assets_to_migrate(loc_count, meta_count)
-				.iter()
-				.len()
+			assets.iter().len()
 		);
+
 		// Complexity: O(meta_count + loc_count) writes
-		Assets::get_assets_to_migrate(loc_count, meta_count)
-			.into_iter()
-			.for_each(|(asset_id, asset_metadata)| {
-				log::debug!("Migrating asset: {:?}", asset_id);
-				orml_asset_registry::Pallet::<T>::do_register_asset_without_asset_processor(
-					asset_metadata.into(),
-					asset_id.into(),
-				)
-				.map_err(|e| log::error!("Failed to register asset id: {:?}", e))
-				.ok();
-			});
+		assets.into_iter().for_each(|(asset_id, asset_metadata)| {
+			log::debug!("Migrating asset: {:?}", asset_id);
+			orml_asset_registry::Pallet::<T>::do_register_asset_without_asset_processor(
+				asset_metadata.into(),
+				asset_id.into(),
+			)
+			.map_err(|e| log::error!("Failed to register asset id: {:?}", e))
+			.ok();
+		});
 
 		log::info!("💎 AssetRegistryMultilocationToXCMV3: on_runtime_upgrade: completed!");
 		RocksDbWeight::get().reads_writes(
-			2,
+			loc_count.saturating_add(meta_count).into(),
 			loc_count
 				.saturating_add(meta_count)
 				.saturating_mul(2)
@@ -130,19 +132,7 @@ impl<
 		log::info!("💎 AssetRegistryMultilocationToXCMV3: pre-upgrade: started");
 		let (loc_count, meta_count) = Self::get_key_counts();
 
-		match (loc_count, meta_count) {
-			(loc, meta)
-				if (loc, meta) == (EXPECTED_MAINNET_LOC_COUNT, EXPECTED_MAINNET_META_COUNT) =>
-			{
-				Ok(())
-			}
-			(loc, meta)
-				if (loc, meta) == (EXPECTED_TESTNET_LOC_COUNT, EXPECTED_TESTNET_META_COUNT) =>
-			{
-				Ok(())
-			}
-			_ => Err("💎 AssetRegistryMultilocationToXCMV3: Unexpected counters"),
-		}?;
+		Self::check_key_counts(loc_count, meta_count)?;
 
 		log::info!("💎 AssetRegistryMultilocationToXCMV3: pre-upgrade: done");
 		Ok((loc_count, meta_count).encode())
@@ -185,7 +175,7 @@ impl<
 {
 	fn get_key_counts() -> (u32, u32) {
 		// let loc_count =
-		// orml_asset_registry::LocationToAssetId::<T>::iter_keys().count() as u32;
+		//  orml_asset_registry::LocationToAssetId::<T>::iter_keys().count() as u32;
 		// let meta_count = orml_asset_registry::Metadata::<T>::iter_keys().count() as
 		// u32;
 		let loc_count = Self::count_storage_keys(
@@ -206,6 +196,22 @@ impl<
 		(loc_count, meta_count)
 	}
 
+	fn check_key_counts(loc_count: u32, meta_count: u32) -> Result<(), &'static str> {
+		match (loc_count, meta_count) {
+			(loc, meta)
+				if (loc, meta) == (EXPECTED_MAINNET_LOC_COUNT, EXPECTED_MAINNET_META_COUNT) =>
+			{
+				Ok(())
+			}
+			(loc, meta)
+				if (loc, meta) == (EXPECTED_TESTNET_LOC_COUNT, EXPECTED_TESTNET_META_COUNT) =>
+			{
+				Ok(())
+			}
+			_ => Err("💎 AssetRegistryMultilocationToXCMV3: Unexpected counters"),
+		}
+	}
+
 	pub fn count_storage_keys(prefix: &[u8]) -> u32 {
 		let mut count = 0;
 		let mut next_key = prefix.to_vec();
@@ -225,10 +231,7 @@ impl<
 }
 
 pub trait AssetsToMigrate {
-	fn get_assets_to_migrate(
-		loc_count: u32,
-		meta_count: u32,
-	) -> Vec<(
+	fn get_assets_to_migrate() -> Vec<(
 		CurrencyId,
 		orml_asset_registry::AssetMetadata<Balance, CustomMetadata>,
 	)>;