Manual Para Lock (#5451)

* remove para lock check for now

* fmt

* manual para lock

* expose schedule_code_upgrade and set_current_head

* extrinsics and benchmarks

* use zero

* add weights

* fix variable name

* add and fix comments

* fix weights

* add back default lock

Co-authored-by: parity-processbot <>

polkadot/runtime/common/src/paras_registrar.rs

@@ -60,6 +60,8 @@ pub trait WeightInfo {
 	fn force_register() -> Weight;
 	fn deregister() -> Weight;
 	fn swap() -> Weight;
+	fn schedule_code_upgrade(b: u32) -> Weight;
+	fn set_current_head(b: u32) -> Weight;
 }
 
 pub struct TestWeightInfo;
@@ -79,6 +81,12 @@ impl WeightInfo for TestWeightInfo {
 	fn swap() -> Weight {
 		Weight::zero()
 	}
+	fn schedule_code_upgrade(_b: u32) -> Weight {
+		Weight::zero()
+	}
+	fn set_current_head(_b: u32) -> Weight {
+		Weight::zero()
+	}
 }
 
 #[frame_support::pallet]
@@ -318,11 +326,11 @@ pub mod pallet {
 		/// Remove a manager lock from a para. This will allow the manager of a
 		/// previously locked para to deregister or swap a para without using governance.
 		///
-		/// Can only be called by the Root origin.
+		/// Can only be called by the Root origin or the parachain.
 		#[pallet::weight(T::DbWeight::get().reads_writes(1, 1))]
-		pub fn force_remove_lock(origin: OriginFor<T>, para: ParaId) -> DispatchResult {
-			ensure_root(origin)?;
-			Self::remove_lock(para);
+		pub fn remove_lock(origin: OriginFor<T>, para: ParaId) -> DispatchResult {
+			Self::ensure_root_or_para(origin, para)?;
+			<Self as Registrar>::remove_lock(para);
 			Ok(())
 		}
 
@@ -348,6 +356,45 @@ pub mod pallet {
 			NextFreeParaId::<T>::set(id + 1);
 			Ok(())
 		}
+
+		/// Add a manager lock from a para. This will prevent the manager of a
+		/// para to deregister or swap a para.
+		///
+		/// Can be called by Root, the parachain, or the parachain manager if the parachain is unlocked.
+		#[pallet::weight(T::DbWeight::get().reads_writes(1, 1))]
+		pub fn add_lock(origin: OriginFor<T>, para: ParaId) -> DispatchResult {
+			Self::ensure_root_para_or_owner(origin, para)?;
+			<Self as Registrar>::apply_lock(para);
+			Ok(())
+		}
+
+		/// Schedule a parachain upgrade.
+		///
+		/// Can be called by Root, the parachain, or the parachain manager if the parachain is unlocked.
+		#[pallet::weight(<T as Config>::WeightInfo::schedule_code_upgrade(new_code.0.len() as u32))]
+		pub fn schedule_code_upgrade(
+			origin: OriginFor<T>,
+			para: ParaId,
+			new_code: ValidationCode,
+		) -> DispatchResult {
+			Self::ensure_root_para_or_owner(origin, para)?;
+			runtime_parachains::schedule_code_upgrade::<T>(para, new_code)?;
+			Ok(())
+		}
+
+		/// Set the parachain's current head.
+		///
+		/// Can be called by Root, the parachain, or the parachain manager if the parachain is unlocked.
+		#[pallet::weight(<T as Config>::WeightInfo::set_current_head(new_head.0.len() as u32))]
+		pub fn set_current_head(
+			origin: OriginFor<T>,
+			para: ParaId,
+			new_head: HeadData,
+		) -> DispatchResult {
+			Self::ensure_root_para_or_owner(origin, para)?;
+			runtime_parachains::set_current_head::<T>(para, new_head);
+			Ok(())
+		}
 	}
 }
 
@@ -379,7 +426,7 @@ impl<T: Config> Registrar for Pallet<T> {
 		Paras::<T>::mutate(id, |x| x.as_mut().map(|mut info| info.locked = true));
 	}
 
-	// Apply a lock to the parachain.
+	// Remove a lock from the parachain.
 	fn remove_lock(id: ParaId) {
 		Paras::<T>::mutate(id, |x| x.as_mut().map(|mut info| info.locked = false));
 	}
@@ -467,17 +514,23 @@ impl<T: Config> Pallet<T> {
 				ensure!(para_info.manager == who, Error::<T>::NotOwner);
 				Ok(())
 			})
-			.or_else(|_| -> DispatchResult {
-				// Else check if para origin...
-				let caller_id =
-					ensure_parachain(<T as Config>::RuntimeOrigin::from(origin.clone()))?;
-				ensure!(caller_id == id, Error::<T>::NotOwner);
-				Ok(())
-			})
-			.or_else(|_| -> DispatchResult {
-				// Check if root...
-				ensure_root(origin.clone()).map_err(|e| e.into())
-			})
+			.or_else(|_| -> DispatchResult { Self::ensure_root_or_para(origin, id) })
+	}
+
+	/// Ensure the origin is one of Root or the `para` itself.
+	fn ensure_root_or_para(
+		origin: <T as frame_system::Config>::RuntimeOrigin,
+		id: ParaId,
+	) -> DispatchResult {
+		if let Ok(caller_id) = ensure_parachain(<T as Config>::RuntimeOrigin::from(origin.clone()))
+		{
+			// Check if matching para id...
+			ensure!(caller_id == id, Error::<T>::NotOwner);
+		} else {
+			// Check if root...
+			ensure_root(origin.clone())?;
+		}
+		Ok(())
 	}
 
 	fn do_reserve(
@@ -1087,21 +1140,20 @@ mod tests {
 				vec![1, 2, 3].into(),
 			));
 
-			// Owner can call swap
-			assert_ok!(Registrar::swap(RuntimeOrigin::signed(1), para_id, para_id + 1));
-
-			// 2 session changes to fully onboard.
-			run_to_session(2);
-			assert_eq!(Parachains::lifecycle(para_id), Some(ParaLifecycle::Parathread));
-
+			assert_noop!(Registrar::add_lock(RuntimeOrigin::signed(2), para_id), BadOrigin);
 			// Once they begin onboarding, we lock them in.
-			assert_ok!(Registrar::make_parachain(para_id));
-
-			// Owner cannot call swap anymore
+			assert_ok!(Registrar::add_lock(RuntimeOrigin::signed(1), para_id));
+			// Owner cannot pass origin check when checking lock
 			assert_noop!(
-				Registrar::swap(RuntimeOrigin::signed(1), para_id, para_id + 2),
+				Registrar::ensure_root_para_or_owner(RuntimeOrigin::signed(1), para_id),
 				BadOrigin
 			);
+			// Owner cannot remove lock.
+			assert_noop!(Registrar::remove_lock(RuntimeOrigin::signed(1), para_id), BadOrigin);
+			// Para can.
+			assert_ok!(Registrar::remove_lock(para_origin(para_id), para_id));
+			// Owner can pass origin check again
+			assert_ok!(Registrar::ensure_root_para_or_owner(RuntimeOrigin::signed(1), para_id));
 		});
 	}
 
@@ -1227,6 +1279,7 @@ mod benchmarking {
 	use crate::traits::Registrar as RegistrarT;
 	use frame_support::assert_ok;
 	use frame_system::RawOrigin;
+	use primitives::v2::{MAX_CODE_SIZE, MAX_HEAD_DATA_SIZE};
 	use runtime_parachains::{paras, shared, Origin as ParaOrigin};
 	use sp_runtime::traits::Bounded;
 
@@ -1343,6 +1396,18 @@ mod benchmarking {
 			assert_eq!(paras::Pallet::<T>::lifecycle(parathread), Some(ParaLifecycle::Parachain));
 		}
 
+		schedule_code_upgrade {
+			let b in 1 .. MAX_CODE_SIZE;
+			let new_code = ValidationCode(vec![0; b as usize]);
+			let para_id = ParaId::from(1000);
+		}: _(RawOrigin::Root, para_id, new_code)
+
+		set_current_head {
+			let b in 1 .. MAX_HEAD_DATA_SIZE;
+			let new_head = HeadData(vec![0; b as usize]);
+			let para_id = ParaId::from(1000);
+		}: _(RawOrigin::Root, para_id, new_head)
+
 		impl_benchmark_test_suite!(
 			Registrar,
 			crate::integration_tests::new_test_ext(),

polkadot/runtime/kusama/src/weights/runtime_common_paras_registrar.rs

@@ -103,4 +103,28 @@ impl<T: frame_system::Config> runtime_common::paras_registrar::WeightInfo for We
 			.saturating_add(T::DbWeight::get().reads(10 as u64))
 			.saturating_add(T::DbWeight::get().writes(8 as u64))
 	}
+	// Storage: Paras FutureCodeHash (r:1 w:1)
+	// Storage: Paras CurrentCodeHash (r:1 w:0)
+	// Storage: Paras UpgradeCooldowns (r:1 w:1)
+	// Storage: Paras PvfActiveVoteMap (r:1 w:0)
+	// Storage: Paras CodeByHash (r:1 w:1)
+	// Storage: Paras UpcomingUpgrades (r:1 w:1)
+	// Storage: System Digest (r:1 w:1)
+	// Storage: Paras CodeByHashRefs (r:1 w:1)
+	// Storage: Paras FutureCodeUpgrades (r:0 w:1)
+	// Storage: Paras UpgradeRestrictionSignal (r:0 w:1)
+	fn schedule_code_upgrade(b: u32, ) -> Weight {
+		Weight::from_ref_time(0 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(3_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().reads(8 as u64))
+			.saturating_add(T::DbWeight::get().writes(8 as u64))
+	}
+	// Storage: Paras Heads (r:0 w:1)
+	fn set_current_head(b: u32, ) -> Weight {
+		Weight::from_ref_time(5_494_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(1_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().writes(1 as u64))
+	}
 }

polkadot/runtime/parachains/src/lib.rs

@@ -50,7 +50,8 @@ mod mock;
 
 pub use origin::{ensure_parachain, Origin};
 pub use paras::ParaLifecycle;
-use primitives::v2::Id as ParaId;
+use primitives::v2::{HeadData, Id as ParaId, ValidationCode};
+use sp_runtime::DispatchResult;
 
 /// Schedule a para to be initialized at the start of the next session with the given genesis data.
 ///
@@ -78,3 +79,20 @@ pub fn schedule_parathread_upgrade<T: paras::Config>(id: ParaId) -> Result<(), (
 pub fn schedule_parachain_downgrade<T: paras::Config>(id: ParaId) -> Result<(), ()> {
 	paras::Pallet::<T>::schedule_parachain_downgrade(id).map_err(|_| ())
 }
+
+/// Schedules a validation code upgrade to a parachain with the given id.
+///
+/// This simply calls [`crate::paras::Pallet::schedule_code_upgrade_external`].
+pub fn schedule_code_upgrade<T: paras::Config>(
+	id: ParaId,
+	new_code: ValidationCode,
+) -> DispatchResult {
+	paras::Pallet::<T>::schedule_code_upgrade_external(id, new_code)
+}
+
+/// Sets the current parachain head with the given id.
+///
+/// This simply calls [`crate::paras::Pallet::set_current_head`].
+pub fn set_current_head<T: paras::Config>(id: ParaId, new_head: HeadData) {
+	paras::Pallet::<T>::set_current_head(id, new_head)
+}

polkadot/runtime/parachains/src/paras/mod.rs

@@ -118,7 +118,7 @@ use primitives::v2::{
 use scale_info::TypeInfo;
 use sp_core::RuntimeDebug;
 use sp_runtime::{
-	traits::{AppVerify, One},
+	traits::{AppVerify, One, Saturating},
 	DispatchResult, SaturatedConversion,
 };
 use sp_std::{cmp, mem, prelude::*};
@@ -535,6 +535,8 @@ pub mod pallet {
 		/// The PVF pre-checking statement cannot be included since the PVF pre-checking mechanism
 		/// is disabled.
 		PvfCheckDisabled,
+		/// Parachain cannot currently schedule a code upgrade.
+		CannotUpgradeCode,
 	}
 
 	/// All currently active PVF pre-checking votes.
@@ -752,8 +754,7 @@ pub mod pallet {
 			new_head: HeadData,
 		) -> DispatchResult {
 			ensure_root(origin)?;
-			<Self as Store>::Heads::insert(&para, new_head);
-			Self::deposit_event(Event::CurrentHeadUpdated(para));
+			Self::set_current_head(para, new_head);
 			Ok(())
 		}
 
@@ -1050,6 +1051,31 @@ const INVALID_TX_DOUBLE_VOTE: u8 = 3;
 const INVALID_TX_PVF_CHECK_DISABLED: u8 = 4;
 
 impl<T: Config> Pallet<T> {
+	/// This is a call to schedule code upgrades for parachains which is safe to be called
+	/// outside of this module. That means this function does all checks necessary to ensure
+	/// that some external code is allowed to trigger a code upgrade. We do not do auth checks,
+	/// that should be handled by whomever calls this function.
+	pub(crate) fn schedule_code_upgrade_external(
+		id: ParaId,
+		new_code: ValidationCode,
+	) -> DispatchResult {
+		// Check that we can schedule an upgrade at all.
+		ensure!(Self::can_upgrade_validation_code(id), Error::<T>::CannotUpgradeCode);
+		let config = configuration::Pallet::<T>::config();
+		let current_block = frame_system::Pallet::<T>::block_number();
+		// Schedule the upgrade with a delay just like if a parachain triggered the upgrade.
+		let upgrade_block = current_block.saturating_add(config.validation_upgrade_delay);
+		Self::schedule_code_upgrade(id, new_code, upgrade_block, &config);
+		Self::deposit_event(Event::CodeUpgradeScheduled(id));
+		Ok(())
+	}
+
+	/// Set the current head of a parachain.
+	pub(crate) fn set_current_head(para: ParaId, new_head: HeadData) {
+		<Self as Store>::Heads::insert(&para, new_head);
+		Self::deposit_event(Event::CurrentHeadUpdated(para));
+	}
+
 	/// Called by the initializer to initialize the paras pallet.
 	pub(crate) fn initializer_initialize(now: T::BlockNumber) -> Weight {
 		let weight = Self::prune_old_code(now);

polkadot/runtime/polkadot/src/weights/runtime_common_paras_registrar.rs

@@ -105,4 +105,28 @@ impl<T: frame_system::Config> runtime_common::paras_registrar::WeightInfo for We
 			.saturating_add(T::DbWeight::get().reads(10 as u64))
 			.saturating_add(T::DbWeight::get().writes(8 as u64))
 	}
+	// Storage: Paras FutureCodeHash (r:1 w:1)
+	// Storage: Paras CurrentCodeHash (r:1 w:0)
+	// Storage: Paras UpgradeCooldowns (r:1 w:1)
+	// Storage: Paras PvfActiveVoteMap (r:1 w:0)
+	// Storage: Paras CodeByHash (r:1 w:1)
+	// Storage: Paras UpcomingUpgrades (r:1 w:1)
+	// Storage: System Digest (r:1 w:1)
+	// Storage: Paras CodeByHashRefs (r:1 w:1)
+	// Storage: Paras FutureCodeUpgrades (r:0 w:1)
+	// Storage: Paras UpgradeRestrictionSignal (r:0 w:1)
+	fn schedule_code_upgrade(b: u32, ) -> Weight {
+		Weight::from_ref_time(0 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(3_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().reads(8 as u64))
+			.saturating_add(T::DbWeight::get().writes(8 as u64))
+	}
+	// Storage: Paras Heads (r:0 w:1)
+	fn set_current_head(b: u32, ) -> Weight {
+		Weight::from_ref_time(5_494_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(1_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().writes(1 as u64))
+	}
 }

polkadot/runtime/rococo/src/weights/runtime_common_paras_registrar.rs

@@ -105,4 +105,28 @@ impl<T: frame_system::Config> runtime_common::paras_registrar::WeightInfo for We
 			.saturating_add(T::DbWeight::get().reads(10 as u64))
 			.saturating_add(T::DbWeight::get().writes(8 as u64))
 	}
+	// Storage: Paras FutureCodeHash (r:1 w:1)
+	// Storage: Paras CurrentCodeHash (r:1 w:0)
+	// Storage: Paras UpgradeCooldowns (r:1 w:1)
+	// Storage: Paras PvfActiveVoteMap (r:1 w:0)
+	// Storage: Paras CodeByHash (r:1 w:1)
+	// Storage: Paras UpcomingUpgrades (r:1 w:1)
+	// Storage: System Digest (r:1 w:1)
+	// Storage: Paras CodeByHashRefs (r:1 w:1)
+	// Storage: Paras FutureCodeUpgrades (r:0 w:1)
+	// Storage: Paras UpgradeRestrictionSignal (r:0 w:1)
+	fn schedule_code_upgrade(b: u32, ) -> Weight {
+		Weight::from_ref_time(0 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(3_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().reads(8 as u64))
+			.saturating_add(T::DbWeight::get().writes(8 as u64))
+	}
+	// Storage: Paras Heads (r:0 w:1)
+	fn set_current_head(b: u32, ) -> Weight {
+		Weight::from_ref_time(5_494_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(1_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().writes(1 as u64))
+	}
 }

polkadot/runtime/westend/src/weights/runtime_common_paras_registrar.rs

@@ -103,4 +103,28 @@ impl<T: frame_system::Config> runtime_common::paras_registrar::WeightInfo for We
 			.saturating_add(T::DbWeight::get().reads(10 as u64))
 			.saturating_add(T::DbWeight::get().writes(8 as u64))
 	}
+	// Storage: Paras FutureCodeHash (r:1 w:1)
+	// Storage: Paras CurrentCodeHash (r:1 w:0)
+	// Storage: Paras UpgradeCooldowns (r:1 w:1)
+	// Storage: Paras PvfActiveVoteMap (r:1 w:0)
+	// Storage: Paras CodeByHash (r:1 w:1)
+	// Storage: Paras UpcomingUpgrades (r:1 w:1)
+	// Storage: System Digest (r:1 w:1)
+	// Storage: Paras CodeByHashRefs (r:1 w:1)
+	// Storage: Paras FutureCodeUpgrades (r:0 w:1)
+	// Storage: Paras UpgradeRestrictionSignal (r:0 w:1)
+	fn schedule_code_upgrade(b: u32, ) -> Weight {
+		Weight::from_ref_time(0 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(3_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().reads(8 as u64))
+			.saturating_add(T::DbWeight::get().writes(8 as u64))
+	}
+	// Storage: Paras Heads (r:0 w:1)
+	fn set_current_head(b: u32, ) -> Weight {
+		Weight::from_ref_time(5_494_000 as u64)
+			// Standard Error: 0
+			.saturating_add(Weight::from_ref_time(1_000 as u64).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().writes(1 as u64))
+	}
 }