[AHM] Preimage migration

integration-tests/ahm/src/tests.rs

@@ -33,7 +33,7 @@
 
 use cumulus_primitives_core::{AggregateMessageOrigin, ParaId};
 use frame_support::{pallet_prelude::*, traits::*, weights::WeightMeter};
-use pallet_rc_migrator::{MigrationStage, RcMigrationStage};
+use pallet_rc_migrator::{types::PalletMigrationChecks, MigrationStage, RcMigrationStage};
 use polkadot_primitives::InboundDownwardMessage;
 use remote_externalities::RemoteExternalities;
 
@@ -47,8 +47,10 @@ async fn account_migration_works() {
 	let para_id = ParaId::from(1000);
 
 	// Simulate relay blocks and grab the DMP messages
-	let dmp_messages = rc.execute_with(|| {
+	let (dmp_messages, pre_check_payload) = rc.execute_with(|| {
 		let mut dmps = Vec::new();
+		let pre_check_payload =
+			pallet_rc_migrator::preimage::PreimageChunkMigrator::<Polkadot>::pre_check();
 
 		// Loop until no more DMPs are added and we had at least 1
 		loop {
@@ -59,10 +61,10 @@ async fn account_migration_works() {
 			dmps.extend(new_dmps);
 
 			if RcMigrationStage::<Polkadot>::get() ==
-				pallet_rc_migrator::MigrationStage::ProxyMigrationDone
+				pallet_rc_migrator::MigrationStage::PreimageMigrationDone
 			{
 				log::info!("Migration done");
-				break dmps;
+				break (dmps, pre_check_payload);
 			}
 		}
 	});
@@ -88,6 +90,10 @@ async fn account_migration_works() {
 			log::debug!("AH DMP messages left: {}", fp.storage.count);
 			next_block_ah();
 		}
+
+		pallet_rc_migrator::preimage::PreimageChunkMigrator::<Polkadot>::post_check(
+			pre_check_payload,
+		);
 		// NOTE that the DMP queue is probably not empty because the snapshot that we use contains
 		// some overweight ones.
 	});

pallets/ah-migrator/src/lib.rs

@@ -33,6 +33,7 @@
 
 pub mod account;
 pub mod multisig;
+pub mod preimage;
 pub mod proxy;
 pub mod types;
 
@@ -48,8 +49,9 @@ use frame_support::{
 };
 use frame_system::pallet_prelude::*;
 use pallet_balances::{AccountData, Reasons as LockReasons};
-use pallet_rc_migrator::{accounts::Account as RcAccount, multisig::*, proxy::*};
+use pallet_rc_migrator::{accounts::Account as RcAccount, multisig::*, preimage::*, proxy::*};
 use sp_application_crypto::Ss58Codec;
+use sp_core::H256;
 use sp_runtime::{
 	traits::{Convert, TryConvert},
 	AccountId32,
@@ -71,6 +73,7 @@ pub mod pallet {
 		+ pallet_balances::Config<Balance = u128>
 		+ pallet_multisig::Config
 		+ pallet_proxy::Config
+		+ pallet_preimage::Config<Hash = H256>
 	{
 		/// The overarching event type.
 		type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
@@ -148,6 +151,18 @@ pub mod pallet {
 			/// How many proxy announcements failed to integrate.
 			count_bad: u32,
 		},
+		/// Received a batch of `RcPreimageChunk` that are going to be integrated.
+		PreimageChunkBatchReceived {
+			/// How many preimage chunks are in the batch.
+			count: u32,
+		},
+		/// We processed a batch of `RcPreimageChunk` that we received.
+		PreimageChunkBatchProcessed {
+			/// How many preimage chunks were successfully integrated.
+			count_good: u32,
+			/// How many preimage chunks failed to integrate.
+			count_bad: u32,
+		},
 	}
 
 	#[pallet::pallet]
@@ -207,8 +222,19 @@ pub mod pallet {
 			announcements: Vec<RcProxyAnnouncementOf<T>>,
 		) -> DispatchResult {
 			ensure_root(origin)?;
+
 			Self::do_receive_proxy_announcements(announcements).map_err(Into::into)
 		}
+
+		#[pallet::call_index(4)]
+		pub fn receive_preimage_chunks(
+			origin: OriginFor<T>,
+			chunks: Vec<RcPreimageChunk>,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+
+			Self::do_receive_preimage_chunks(chunks).map_err(Into::into)
+		}
 	}
 
 	#[pallet::hooks]

pallets/ah-migrator/src/preimage.rs

@@ -0,0 +1,95 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::{types::*, *};
+use pallet_rc_migrator::preimage::alias;
+
+impl<T: Config> Pallet<T> {
+	pub fn do_receive_preimage_chunks(chunks: Vec<RcPreimageChunk>) -> Result<(), Error<T>> {
+		Self::deposit_event(Event::PreimageChunkBatchReceived { count: chunks.len() as u32 });
+		let (mut count_good, mut count_bad) = (0, 0);
+		log::info!(target: LOG_TARGET, "Integrating {} preimage chunks", chunks.len());
+
+		for chunk in chunks {
+			match Self::do_receive_preimage_chunk(chunk) {
+				Ok(()) => count_good += 1,
+				Err(e) => {
+					count_bad += 1;
+					log::error!(target: LOG_TARGET, "Error while integrating preimage chunk: {:?}", e);
+				},
+			}
+		}
+		Self::deposit_event(Event::PreimageChunkBatchProcessed { count_good, count_bad });
+
+		Ok(())
+	}
+
+	pub fn do_receive_preimage_chunk(chunk: RcPreimageChunk) -> Result<(), Error<T>> {
+		log::debug!(target: LOG_TARGET, "Integrating preimage chunk {} offset {}/{}", chunk.preimage_hash, chunk.chunk_byte_offset + chunk.chunk_bytes.len() as u32, chunk.preimage_len);
+		let key = (chunk.preimage_hash, chunk.preimage_len);
+
+		// First check that we did not miss a chunk
+		let preimage = match alias::PreimageFor::<T>::get(&key) {
+			Some(mut preimage) => {
+				if preimage.len() != chunk.chunk_byte_offset as usize {
+					defensive!("Preimage chunk missing");
+					return Err(Error::<T>::TODO);
+				}
+
+				match preimage.try_mutate(|p| {
+					p.extend(chunk.chunk_bytes.clone());
+				}) {
+					Some(preimage) => {
+						alias::PreimageFor::<T>::insert(&key, &preimage);
+						preimage
+					},
+					None => {
+						defensive!("Preimage too big");
+						return Err(Error::<T>::TODO);
+					},
+				}
+			},
+			None => {
+				if chunk.chunk_byte_offset != 0 {
+					defensive!("Preimage chunk missing");
+					return Err(Error::<T>::TODO);
+				}
+
+				let preimage: BoundedVec<
+					u8,
+					ConstU32<{ pallet_rc_migrator::preimage::CHUNK_SIZE }>,
+				> = chunk.chunk_bytes;
+				debug_assert!(
+					pallet_rc_migrator::preimage::CHUNK_SIZE <=
+						pallet_rc_migrator::preimage::alias::MAX_SIZE
+				);
+				let bounded_preimage: BoundedVec<
+					u8,
+					ConstU32<{ pallet_rc_migrator::preimage::alias::MAX_SIZE }>,
+				> = preimage.into_inner().try_into().expect("Asserted");
+				alias::PreimageFor::<T>::insert(key, &bounded_preimage);
+				bounded_preimage
+			},
+		};
+
+		if preimage.len() == chunk.preimage_len as usize + chunk.chunk_byte_offset as usize {
+			log::debug!(target: LOG_TARGET, "Preimage complete: {}", chunk.preimage_hash);
+		}
+
+		Ok(())
+	}
+}

pallets/rc-migrator/Cargo.toml

@@ -23,6 +23,7 @@ pallet-balances = { workspace = true }
 pallet-staking = { workspace = true }
 pallet-proxy = { workspace = true }
 pallet-multisig = { workspace = true }
+pallet-preimage = { workspace = true }
 polkadot-runtime-common = { workspace = true }
 runtime-parachains = { workspace = true }
 polkadot-parachain-primitives = { workspace = true }
@@ -51,6 +52,7 @@ std = [
 	"sp-runtime/std",
 	"sp-std/std",
 	"xcm/std",
+	"pallet-preimage/std"
 ]
 runtime-benchmarks = [
 	"frame-benchmarking/runtime-benchmarks",
@@ -64,6 +66,7 @@ runtime-benchmarks = [
 	"polkadot-runtime-common/runtime-benchmarks",
 	"runtime-parachains/runtime-benchmarks",
 	"sp-runtime/runtime-benchmarks",
+	"pallet-preimage/runtime-benchmarks"
 ]
 try-runtime = [
 	"frame-support/try-runtime",
@@ -75,4 +78,5 @@ try-runtime = [
 	"polkadot-runtime-common/try-runtime",
 	"runtime-parachains/try-runtime",
 	"sp-runtime/try-runtime",
+	"pallet-preimage/try-runtime"
 ]

pallets/rc-migrator/src/lib.rs

@@ -33,6 +33,7 @@
 
 pub mod accounts;
 pub mod multisig;
+pub mod preimage;
 pub mod proxy;
 pub mod types;
 mod weights;
@@ -54,7 +55,7 @@ use pallet_balances::AccountData;
 use polkadot_parachain_primitives::primitives::Id as ParaId;
 use polkadot_runtime_common::paras_registrar;
 use runtime_parachains::hrmp;
-use sp_core::crypto::Ss58Codec;
+use sp_core::{crypto::Ss58Codec, H256};
 use sp_runtime::{traits::TryConvert, AccountId32};
 use sp_std::prelude::*;
 use storage::TransactionOutcome;
@@ -63,6 +64,7 @@ use weights::WeightInfo;
 use xcm::prelude::*;
 
 use multisig::MultisigMigrator;
+use preimage::PreimageChunkMigrator;
 use proxy::*;
 use types::PalletMigration;
 
@@ -103,6 +105,13 @@ pub enum MigrationStage<AccountId> {
 		last_key: Option<AccountId>,
 	},
 	ProxyMigrationDone,
+	PreimageMigrationInit,
+	PreimageMigrationChunksOngoing {
+		// TODO type
+		last_key: Option<(Option<(H256, u32)>, u32)>,
+	},
+	PreimageMigrationChunksDone,
+	PreimageMigrationDone,
 }
 
 type AccountInfoFor<T> =
@@ -125,6 +134,7 @@ pub mod pallet {
 		+ paras_registrar::Config
 		+ pallet_multisig::Config
 		+ pallet_proxy::Config
+		+ pallet_preimage::Config<Hash = H256>
 	{
 		/// The overarching event type.
 		type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
@@ -220,8 +230,9 @@ pub mod pallet {
 					// TODO: not complete
 
 					let _ = Self::obtain_rc_accounts();
-					Self::transition(MigrationStage::AccountsMigrationInit);
-					//Self::transition(MigrationStage::ProxyMigrationInit);
+					// Swap this out if you want to skip some migrations for faster debugging
+					//Self::transition(MigrationStage::AccountsMigrationInit);
+					Self::transition(MigrationStage::PreimageMigrationInit);
 				},
 				MigrationStage::AccountsMigrationInit => {
 					let first_acc = match Self::first_account(&mut weight_counter).defensive() {
@@ -359,7 +370,42 @@ pub mod pallet {
 					}
 				},
 				MigrationStage::ProxyMigrationDone => {
-					todo!();
+					Self::transition(MigrationStage::PreimageMigrationInit);
+				},
+				MigrationStage::PreimageMigrationInit => {
+					Self::transition(MigrationStage::PreimageMigrationChunksOngoing {
+						last_key: None,
+					});
+				},
+				MigrationStage::PreimageMigrationChunksOngoing { last_key } => {
+					let res = with_transaction_opaque_err::<Option<_>, Error<T>, _>(|| {
+						TransactionOutcome::Commit(PreimageChunkMigrator::<T>::migrate_many(
+							last_key,
+							&mut weight_counter,
+						))
+					})
+					.expect("Always returning Ok; qed");
+
+					match res {
+						Ok(None) => {
+							Self::transition(MigrationStage::PreimageMigrationDone);
+						},
+						Ok(Some(last_key)) => {
+							Self::transition(MigrationStage::PreimageMigrationChunksOngoing {
+								last_key: Some(last_key),
+							});
+						},
+						e => {
+							log::error!(target: LOG_TARGET, "Error while migrating preimages: {:?}", e);
+							defensive!("Error while migrating preimages");
+						},
+					}
+				},
+				MigrationStage::PreimageMigrationChunksDone => {
+					Self::transition(MigrationStage::PreimageMigrationDone);
+				},
+				MigrationStage::PreimageMigrationDone => {
+					unimplemented!()
 				},
 			};
 
@@ -405,7 +451,7 @@ pub mod pallet {
 					batch.push(items.pop().unwrap()); // FAIL-CI no unwrap
 				}
 
-				log::info!(target: LOG_TARGET, "Sending batch of {} proxies", batch.len());
+				log::info!(target: LOG_TARGET, "Sending XCM batch of {} items", batch.len());
 				let call = types::AssetHubPalletConfig::<T>::AhmController(create_call(batch));
 
 				let message = Xcm(vec![

pallets/rc-migrator/src/preimage.md

@@ -23,7 +23,7 @@ Deprecated. Will not be migrated but funds will be unreserved.
 
 For anyone who has registered a preimage:
 - If the preimage was in the new RequestStatusFor: Some unlocked funds 😎. We cannot calculate a list of affected accounts in advance since users can still influence this.
-- If the preimage was in the old StatusFor: will be removed and funds unlocked. Exhaustive list of all 166 Polkadot accounts that are affected by this and will have **UP TO** these funds unlocked (not a legally binding statement):
+- If the preimage was in the old StatusFor: will be removed and funds unlocked. [Exhaustive list](https://github.com/ggwpez/substrate-scripts/blob/master/ahm-preimage-statusfor-accounts.py) of all 166 Polkadot accounts that are affected by this and will have **UP TO** these funds unlocked (not a legally binding statement):
 
 - `16LKv69ct6xDzSiUjuz154vCg62dkyysektHFCeJe85xb6X`: 1256.897 DOT
 - `15ynbcMgPf7HbQErRz66RDLMuBVdcWVuURhR4SLPiqa6B8jx`: 633.121 DOT