[AHM] Add Polkadot call filtering

integration-tests/ahm/Cargo.toml

@@ -20,6 +20,8 @@ grandpa = { workspace = true }
 log = { workspace = true, default-features = true }
 pallet-rc-migrator = { workspace = true, default-features = true }
 pallet-ah-migrator = { workspace = true, default-features = true }
+pallet-indices = { workspace = true, default-features = true }
+pallet-message-queue = { workspace = true, default-features = true }
 parachains-common = { workspace = true, default-features = true }
 polkadot-parachain-primitives = { workspace = true, default-features = true }
 polkadot-primitives = { workspace = true, default-features = true }

pallets/rc-migrator/src/lib.rs

@@ -46,7 +46,7 @@ use frame_support::{
 	traits::{
 		fungible::{Inspect, InspectFreeze, Mutate, MutateFreeze, MutateHold},
 		tokens::{Fortitude, Precision, Preservation},
-		Defensive, LockableCurrency, ReservableCurrency,
+		Contains, Defensive, LockableCurrency, ReservableCurrency,
 	},
 	weights::WeightMeter,
 };
@@ -96,7 +96,19 @@ impl<T: Config> From<OutOfWeightError> for Error<T> {
 	}
 }
 
-#[derive(Encode, Decode, Clone, PartialEq, Eq, Default, RuntimeDebug, TypeInfo, MaxEncodedLen)]
+#[derive(
+	Encode,
+	Decode,
+	Clone,
+	PartialEq,
+	Eq,
+	Default,
+	RuntimeDebug,
+	TypeInfo,
+	MaxEncodedLen,
+	PartialOrd,
+	Ord,
+)]
 pub enum MigrationStage<AccountId> {
 	/// The migration has not yet started but will start in the next block.
 	#[default]
@@ -196,6 +208,12 @@ pub mod pallet {
 		type AhWeightInfo: AhWeightInfo;
 		/// The existential deposit on the Asset Hub.
 		type AhExistentialDeposit: Get<<Self as pallet_balances::Config>::Balance>;
+		/// Contains all calls that are allowed during the migration.
+		///
+		/// The calls in here will be available again after the migration.
+		type RcCallEnabledDuringMigration: Contains<<Self as frame_system::Config>::RuntimeCall>;
+		/// Contains all calls that are allowed after the migration finished.
+		type RcCallEnabledAfterMigration: Contains<<Self as frame_system::Config>::RuntimeCall>;
 	}
 
 	#[pallet::error]
@@ -250,16 +268,6 @@ pub mod pallet {
 	#[pallet::pallet]
 	pub struct Pallet<T>(_);
 
-	#[pallet::call]
-	impl<T: Config> Pallet<T> {
-		/// TODO
-		#[pallet::call_index(0)]
-		#[pallet::weight({1})]
-		pub fn do_something(_origin: OriginFor<T>) -> DispatchResultWithPostInfo {
-			Ok(().into())
-		}
-	}
-
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		fn on_initialize(_: BlockNumberFor<T>) -> Weight {
@@ -593,3 +601,25 @@ pub mod pallet {
 		}
 	}
 }
+
+impl<T: Config> Contains<<T as frame_system::Config>::RuntimeCall> for Pallet<T> {
+	fn contains(call: &<T as frame_system::Config>::RuntimeCall) -> bool {
+		let stage = RcMigrationStage::<T>::get();
+		let is_finished = stage >= MigrationStage::MigrationDone;
+		let is_ongoing = stage > MigrationStage::Pending && !is_finished;
+
+		// We have to return whether the call is allowed:
+		const ALLOWED: bool = true;
+		const FORBIDDEN: bool = false;
+
+		if is_finished && !T::RcCallEnabledAfterMigration::contains(call) {
+			return FORBIDDEN;
+		}
+
+		if is_ongoing && !T::RcCallEnabledDuringMigration::contains(call) {
+			return FORBIDDEN;
+		}
+
+		ALLOWED
+	}
+}

relay/polkadot/src/lib.rs

@@ -67,7 +67,7 @@ use frame_support::{
 	traits::{
 		fungible::HoldConsideration,
 		tokens::{imbalance::ResolveTo, UnityOrOuterConversion},
-		ConstU32, ConstU8, EitherOf, EitherOfDiverse, Everything, FromContains, Get,
+		ConstU32, ConstU8, Contains, EitherOf, EitherOfDiverse, Everything, FromContains, Get,
 		InstanceFilter, KeyOwnerProofSystem, LinearStoragePrice, OnRuntimeUpgrade, PrivilegeCmp,
 		ProcessMessage, ProcessMessageError, WithdrawReasons,
 	},
@@ -190,7 +190,7 @@ parameter_types! {
 }
 
 impl frame_system::Config for Runtime {
-	type BaseCallFilter = Everything;
+	type BaseCallFilter = RcMigrator;
 	type BlockWeights = BlockWeights;
 	type BlockLength = BlockLength;
 	type RuntimeOrigin = RuntimeOrigin;
@@ -1576,6 +1576,113 @@ impl pallet_rc_migrator::Config for Runtime {
 	type AhExistentialDeposit = AhExistentialDeposit;
 	type RcWeightInfo = ();
 	type AhWeightInfo = ();
+	type RcCallEnabledDuringMigration = CallsEnabledDuringMigration;
+	type RcCallEnabledAfterMigration = CallsEnabledAfterMigration;
+}
+
+/// Contains all calls that are enabled during the migration.
+pub struct CallsEnabledDuringMigration;
+impl Contains<<Runtime as frame_system::Config>::RuntimeCall> for CallsEnabledDuringMigration {
+	fn contains(call: &<Runtime as frame_system::Config>::RuntimeCall) -> bool {
+		let (during, _after) = call_disable_status(call);
+		!during
+	}
+}
+
+/// Contains all calls that are enabled after the migration.
+pub struct CallsEnabledAfterMigration;
+impl Contains<<Runtime as frame_system::Config>::RuntimeCall> for CallsEnabledAfterMigration {
+	fn contains(call: &<Runtime as frame_system::Config>::RuntimeCall) -> bool {
+		let (_during, after) = call_disable_status(call);
+		!after
+	}
+}
+
+/// Return whether a call should be disabled during and/or after the migration.
+///
+/// Time line looks like this:
+///
+/// --------|-----------|--------->
+///       Start        End
+///
+/// We now define 2 periods:
+///
+/// 1. During the migration: [Start, End]
+/// 2. After the migration: (End, ∞)
+///
+/// Visually:
+///
+///         |-----1-----|
+///                      |---2---->
+/// --------|-----------|--------->
+///       Start        End
+///
+/// This call returns a 2-tuple to indicate whether a call is disabled during these periods.
+pub fn call_disable_status(call: &<Runtime as frame_system::Config>::RuntimeCall) -> (bool, bool) {
+	use RuntimeCall::*;
+	const ON: bool = false;
+	const OFF: bool = true;
+
+	match call {
+		System(..) => (OFF, ON),
+		Scheduler(..) => (OFF, OFF),
+		Preimage(..) => (OFF, OFF),
+		Babe(..) => todo!("FAIL-CI"),
+		Timestamp(..) => (OFF, OFF),
+		Indices(..) => (OFF, OFF),
+		Balances(..) => (OFF, ON),
+		// TransactionPayment has no calls
+		// Authorship has no calls
+		Staking(..) => (OFF, OFF),
+		// Offences has no calls
+		// Historical has no calls
+		Session(..) => (OFF, OFF),
+		Grandpa(..) => todo!("FAIL-CI"),
+		// AuthorityDiscovery has no calls
+		Treasury(..) => (OFF, OFF),
+		ConvictionVoting(..) => (OFF, OFF),
+		Referenda(..) => (OFF, OFF),
+		// Origins has no calls
+		Whitelist(..) => (OFF, OFF),
+		Claims(..) => (OFF, OFF),
+		Vesting(..) => (OFF, OFF),
+		Utility(..) => (OFF, ON), // batching etc
+		Proxy(..) => (OFF, ON),
+		Multisig(..) => (OFF, ON),
+		Bounties(..) => (OFF, OFF),
+		ChildBounties(..) => (OFF, OFF),
+		ElectionProviderMultiPhase(..) => (OFF, OFF),
+		VoterList(..) => (OFF, OFF),
+		NominationPools(..) => (OFF, OFF),
+		FastUnstake(..) => (OFF, OFF),
+		// DelegatedStaking has on calls
+		// ParachainsOrigin has no calls
+		Configuration(..) => (OFF, ON),
+		ParasShared(..) => (OFF, OFF), /* Has no calls but a call enum https://github.com/paritytech/polkadot-sdk/blob/ee803b74056fac5101c06ec5998586fa6eaac470/polkadot/runtime/parachains/src/shared.rs#L185-L186 */
+		ParaInclusion(..) => (OFF, OFF), /* Has no calls but a call enum https://github.com/paritytech/polkadot-sdk/blob/74ec1ee226ace087748f38dfeffc869cd5534ac8/polkadot/runtime/parachains/src/inclusion/mod.rs#L352-L353 */
+		ParaInherent(..) => (ON, ON),    // only inherents
+		// ParaScheduler has no calls
+		Paras(..) => (OFF, ON), // TODO only root so could think about keeping it on
+		Initializer(..) => (OFF, ON), // only root so could think about keeping it on
+		// Dmp // has no calls and deprecated
+		Hrmp(..) => (OFF, OFF),
+		// ParaSessionInfo has no calls
+		ParasDisputes(..) => (OFF, ON), // TODO check with security
+		ParasSlashing(..) => (OFF, ON), // TODO check with security
+		OnDemand(..) => (OFF, ON),
+		// CoretimeAssignmentProvider has no calls
+		Registrar(..) => (OFF, ON),
+		Slots(..) => (OFF, ON), // TODO not sure
+		Auctions(..) => (OFF, OFF),
+		Crowdloan(..) => (OFF, ON), // TODO maybe only payouts
+		Coretime(..) => (OFF, ON),
+		StateTrieMigration(..) => (OFF, OFF), // Deprecated
+		XcmPallet(..) => (OFF, ON),
+		MessageQueue(..) => (ON, ON), // TODO think about this
+		AssetRate(..) => (OFF, OFF),  // TODO @muharem
+		Beefy(..) => (OFF, ON),       /* TODO @claravanstaden @bkontur
+		                                * RcMigrator has no calls currently */
+	}
 }
 
 construct_runtime! {

relay/polkadot/tests/ahm/mod.rs

@@ -18,12 +18,13 @@
 
 mod accounts;
 
-// Runtime specific imports
-// Use general aliases for the imports to make it easier to copy&paste the tests for other runtimes.
-use polkadot_runtime::{Block, RcMigrator, Runtime as T, System, *};
-
-// General imports
+use frame_support::{sp_runtime::traits::Dispatchable, traits::Contains};
+use pallet_rc_migrator::*;
+use polkadot_primitives::Id as ParaId;
+use polkadot_runtime::{Block, BuildStorage, RcMigrator, Runtime as T, RuntimeOrigin, System, *};
 use remote_externalities::{Builder, Mode, OfflineConfig, RemoteExternalities};
+use runtime_parachains::inclusion::AggregateMessageOrigin;
+use sp_runtime::AccountId32;
 
 /// Create externalities that have their state initialized from a snapshot.
 ///
@@ -48,3 +49,101 @@ async fn remote_ext_test_setup() -> Option<RemoteExternalities<Block>> {
 
 	Some(ext)
 }
+
+#[test]
+fn call_filter_works() {
+	let mut t: sp_io::TestExternalities =
+		frame_system::GenesisConfig::<T>::default().build_storage().unwrap().into();
+
+	// MQ calls are never filtered:
+	let mq_call =
+		polkadot_runtime::RuntimeCall::MessageQueue(pallet_message_queue::Call::<T>::reap_page {
+			message_origin: AggregateMessageOrigin::Ump(
+				runtime_parachains::inclusion::UmpQueueId::Para(ParaId::from(1000)),
+			),
+			page_index: 0,
+		});
+	// System calls are filtered during the migration:
+	let system_call =
+		polkadot_runtime::RuntimeCall::System(frame_system::Call::<T>::remark { remark: vec![] });
+	// Indices calls are filtered during and after the migration:
+	let indices_call =
+		polkadot_runtime::RuntimeCall::Indices(pallet_indices::Call::<T>::claim { index: 0 });
+
+	let is_allowed = |call: &polkadot_runtime::RuntimeCall| Pallet::<T>::contains(call);
+
+	// Try the BaseCallFilter
+	t.execute_with(|| {
+		// Before the migration starts
+		{
+			RcMigrationStage::<T>::put(MigrationStage::Pending);
+
+			assert!(is_allowed(&mq_call));
+			assert!(is_allowed(&system_call));
+			assert!(is_allowed(&indices_call));
+		}
+
+		// During the migration
+		{
+			RcMigrationStage::<T>::put(MigrationStage::ProxyMigrationInit);
+
+			assert!(is_allowed(&mq_call));
+			assert!(!is_allowed(&system_call));
+			assert!(!is_allowed(&indices_call));
+		}
+
+		// After the migration
+		{
+			RcMigrationStage::<T>::put(MigrationStage::MigrationDone);
+
+			assert!(is_allowed(&mq_call));
+			assert!(is_allowed(&system_call));
+			assert!(!is_allowed(&indices_call));
+		}
+	});
+
+	// Try to actually dispatch the calls
+	t.execute_with(|| {
+		let alice = AccountId32::from([0; 32]);
+		<pallet_balances::Pallet<T> as frame_support::traits::Currency<_>>::deposit_creating(
+			&alice,
+			u64::MAX.into(),
+		);
+
+		// Before the migration starts
+		{
+			RcMigrationStage::<T>::put(MigrationStage::Pending);
+
+			assert!(!is_forbidden(&mq_call));
+			assert!(!is_forbidden(&system_call));
+			assert!(!is_forbidden(&indices_call));
+		}
+
+		// During the migration
+		{
+			RcMigrationStage::<T>::put(MigrationStage::ProxyMigrationInit);
+
+			assert!(!is_forbidden(&mq_call));
+			assert!(is_forbidden(&system_call));
+			assert!(is_forbidden(&indices_call));
+		}
+
+		// After the migration
+		{
+			RcMigrationStage::<T>::put(MigrationStage::MigrationDone);
+
+			assert!(!is_forbidden(&mq_call));
+			assert!(!is_forbidden(&system_call));
+			assert!(is_forbidden(&indices_call));
+		}
+	});
+}
+
+fn is_forbidden(call: &polkadot_runtime::RuntimeCall) -> bool {
+	let Err(err) = call.clone().dispatch(RuntimeOrigin::signed(AccountId32::from([0; 32]))) else {
+		return false;
+	};
+
+	let runtime_err: sp_runtime::DispatchError = frame_system::Error::<T>::CallFiltered.into();
+	err.error == runtime_err
+}