Add Paras authorize_code_hash + apply_authorized_code feature

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

@@ -249,6 +249,30 @@ mod benchmarks {
 		}
 	}
 
+	#[benchmark]
+	fn authorize_force_set_current_code_hash() {
+		let para_id = ParaId::from(1000);
+		let code_hash = [0; 32].into();
+
+		#[extrinsic_call]
+		_(RawOrigin::Root, para_id, code_hash);
+
+		assert_last_event::<T>(Event::CodeAuthorized { para_id, code_hash }.into());
+	}
+
+	#[benchmark]
+	fn apply_authorized_force_set_current_code(c: Linear<MIN_CODE_SIZE, MAX_CODE_SIZE>) {
+		let new_code = ValidationCode(vec![0; c as usize]);
+		let para_id = ParaId::from(c as u32);
+		AuthorizedCodeHash::<T>::insert(&para_id, new_code.hash());
+		generate_disordered_pruning::<T>();
+
+		#[extrinsic_call]
+		_(RawOrigin::Root, para_id, new_code);
+
+		assert_last_event::<T>(Event::CurrentCodeUpdated(para_id).into());
+	}
+
 	impl_benchmark_test_suite!(
 		Pallet,
 		crate::mock::new_test_ext(Default::default()),

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

@@ -117,7 +117,10 @@ use alloc::{collections::btree_set::BTreeSet, vec::Vec};
 use bitvec::{order::Lsb0 as BitOrderLsb0, vec::BitVec};
 use codec::{Decode, Encode};
 use core::{cmp, mem};
-use frame_support::{pallet_prelude::*, traits::EstimateNextSessionRotation, DefaultNoBound};
+use frame_support::{
+	dispatch::PostDispatchInfo, pallet_prelude::*, traits::EstimateNextSessionRotation,
+	DefaultNoBound,
+};
 use frame_system::pallet_prelude::*;
 use polkadot_primitives::{
 	ConsensusLog, HeadData, Id as ParaId, PvfCheckStatement, SessionIndex, UpgradeGoAhead,
@@ -551,6 +554,8 @@ pub trait WeightInfo {
 	fn include_pvf_check_statement_finalize_onboarding_accept() -> Weight;
 	fn include_pvf_check_statement_finalize_onboarding_reject() -> Weight;
 	fn include_pvf_check_statement() -> Weight;
+	fn authorize_force_set_current_code_hash() -> Weight;
+	fn apply_authorized_force_set_current_code(c: u32) -> Weight;
 }
 
 pub struct TestWeightInfo;
@@ -596,6 +601,12 @@ impl WeightInfo for TestWeightInfo {
 		// This special value is to distinguish from the finalizing variants above in tests.
 		Weight::MAX - Weight::from_parts(1, 1)
 	}
+	fn authorize_force_set_current_code_hash() -> Weight {
+		Weight::MAX
+	}
+	fn apply_authorized_force_set_current_code(_c: u32) -> Weight {
+		Weight::MAX
+	}
 }
 
 #[frame_support::pallet]
@@ -649,6 +660,8 @@ pub mod pallet {
 	pub enum Event {
 		/// Current code has been updated for a Para. `para_id`
 		CurrentCodeUpdated(ParaId),
+		/// New code hash has been authorized for a Para.
+		CodeAuthorized { code_hash: ValidationCodeHash, para_id: ParaId },
 		/// Current head has been updated for a Para. `para_id`
 		CurrentHeadUpdated(ParaId),
 		/// A code upgrade has been scheduled for a Para. `para_id`
@@ -791,6 +804,11 @@ pub mod pallet {
 	#[pallet::storage]
 	pub type FutureCodeHash<T: Config> = StorageMap<_, Twox64Concat, ParaId, ValidationCodeHash>;
 
+	/// The code hash of a para which is authorized.
+	#[pallet::storage]
+	pub(super) type AuthorizedCodeHash<T: Config> =
+		StorageMap<_, Twox64Concat, ParaId, ValidationCodeHash>;
+
 	/// This is used by the relay-chain to communicate to a parachain a go-ahead with in the upgrade
 	/// procedure.
 	///
@@ -895,10 +913,7 @@ pub mod pallet {
 			new_code: ValidationCode,
 		) -> DispatchResult {
 			ensure_root(origin)?;
-			let new_code_hash = new_code.hash();
-			Self::increase_code_ref(&new_code_hash, &new_code);
-			Self::set_current_code(para, new_code_hash, frame_system::Pallet::<T>::block_number());
-			Self::deposit_event(Event::CurrentCodeUpdated(para));
+			Self::do_force_set_current_code_update(para, new_code);
 			Ok(())
 		}
 
@@ -1149,6 +1164,78 @@ pub mod pallet {
 			MostRecentContext::<T>::insert(&para, context);
 			Ok(())
 		}
+
+		/// Sets the storage for the authorized current code hash of the parachain.
+		///
+		/// This can be useful when we want to trigger `Paras::force_set_current_code(para, code)`
+		/// from a different chain than the one where the `Paras` pallet is deployed.
+		///
+		/// The main reason is to avoid transferring the entire `new_code` wasm blob between chains.
+		/// Instead, we authorize `new_code_hash` with `root`, which can later be applied by
+		/// `Paras::apply_authorized_force_set_current_code(para, new_code)` by anyone.
+		#[pallet::call_index(9)]
+		#[pallet::weight(<T as Config>::WeightInfo::authorize_force_set_current_code_hash())]
+		pub fn authorize_force_set_current_code_hash(
+			origin: OriginFor<T>,
+			para: ParaId,
+			new_code_hash: ValidationCodeHash,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+
+			// if one is already authorized, means it has not been applied yet, so we just replace
+			// it.
+			if let Some(already_authorized) = AuthorizedCodeHash::<T>::take(para) {
+				log::warn!(
+					target: LOG_TARGET,
+					"Already authorized code hash: {:?} found for para {:?}, just removing it!",
+					already_authorized, para
+				);
+			}
+
+			// TODO: FAIL-CI - more validations?
+			// do we need to check against `FutureCodeHash`, `CodeHashRef`,
+			// `PastCodeHash`,... code hashes?
+
+			// insert authorized code hash.
+			AuthorizedCodeHash::<T>::insert(&para, new_code_hash);
+			Self::deposit_event(Event::CodeAuthorized { para_id: para, code_hash: new_code_hash });
+
+			Ok(())
+		}
+
+		/// Applies the already authorized current code for the parachain,
+		/// triggering the same functionality as `force_set_current_code`.
+		#[pallet::call_index(10)]
+		#[pallet::weight(<T as Config>::WeightInfo::apply_authorized_force_set_current_code(new_code.0.len() as u32))]
+		pub fn apply_authorized_force_set_current_code(
+			_origin: OriginFor<T>,
+			para: ParaId,
+			new_code: ValidationCode,
+		) -> DispatchResultWithPostInfo {
+			// no need to ensure, anybody can do this
+
+			// check `new_code` is authorized
+			let Some(authorized_code_hash) = AuthorizedCodeHash::<T>::take(&para) else {
+				log::error!(target: LOG_TARGET, "No authorized code hash found for para {:?}!", para);
+				return Err(Error::<T>::CannotUpgradeCode.into())
+			};
+			let new_code_hash = new_code.hash();
+			ensure!(new_code_hash == authorized_code_hash, Error::<T>::CannotUpgradeCode);
+
+			// TODO: FAIL-CI - more validations?
+
+			// set current code
+			Self::do_force_set_current_code_update(para, new_code);
+
+			// if ok, then allows "for free"
+			let post = PostDispatchInfo {
+				// consume the rest of the block to prevent further transactions
+				actual_weight: Some(T::BlockWeights::get().max_block),
+				// no fee for valid upgrade
+				pays_fee: Pays::No,
+			};
+			Ok(post)
+		}
 	}
 
 	#[pallet::validate_unsigned]
@@ -2159,6 +2246,14 @@ impl<T: Config> Pallet<T> {
 		weight + T::DbWeight::get().writes(1)
 	}
 
+	/// Force set the current code for the given parachain.
+	fn do_force_set_current_code_update(para: ParaId, new_code: ValidationCode) {
+		let new_code_hash = new_code.hash();
+		Self::increase_code_ref(&new_code_hash, &new_code);
+		Self::set_current_code(para, new_code_hash, frame_system::Pallet::<T>::block_number());
+		Self::deposit_event(Event::CurrentCodeUpdated(para));
+	}
+
 	/// Returns the list of PVFs (aka validation code) that require casting a vote by a validator in
 	/// the active validator set.
 	pub(crate) fn pvfs_require_precheck() -> Vec<ValidationCodeHash> {

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

@@ -2013,3 +2013,152 @@ fn parachains_cache_preserves_order() {
 		assert_eq!(Parachains::<Test>::get(), vec![a, c]);
 	});
 }
+
+#[test]
+fn authorize_and_apply_set_current_code_works() {
+	new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+		let para_a = ParaId::from(111);
+		let code_1 = ValidationCode(vec![1]);
+		let code_2 = ValidationCode(vec![2]);
+		let code_3 = ValidationCode(vec![3]);
+		let code_1_hash = code_1.hash();
+		let code_2_hash = code_2.hash();
+		let code_3_hash = code_3.hash();
+		let root = RuntimeOrigin::root();
+		let non_root = RuntimeOrigin::signed(1);
+
+		// check before
+		assert!(AuthorizedCodeHash::<Test>::get(para_a).is_none());
+		assert!(CurrentCodeHash::<Test>::get(para_a).is_none());
+		check_code_is_not_stored(&code_1);
+		check_code_is_not_stored(&code_2);
+		check_code_is_not_stored(&code_3);
+
+		// cannot apply non-authorized code
+		assert_err!(
+			Paras::apply_authorized_force_set_current_code(
+				non_root.clone(),
+				para_a,
+				code_1.clone()
+			),
+			Error::<Test>::CannotUpgradeCode,
+		);
+
+		// non-root user cannot authorize
+		assert_err!(
+			Paras::authorize_force_set_current_code_hash(non_root.clone(), para_a, code_1_hash),
+			DispatchError::BadOrigin,
+		);
+		// root can authorize
+		assert_ok!(Paras::authorize_force_set_current_code_hash(root.clone(), para_a, code_1_hash));
+
+		// check authorized code hash stored
+		assert_eq!(AuthorizedCodeHash::<Test>::get(para_a), Some(code_1_hash));
+		assert!(CurrentCodeHash::<Test>::get(para_a).is_none());
+		check_code_is_not_stored(&code_1);
+		check_code_is_not_stored(&code_2);
+		check_code_is_not_stored(&code_3);
+
+		// non-root cannot apply unauthorized code
+		assert_err!(
+			Paras::apply_authorized_force_set_current_code(
+				non_root.clone(),
+				para_a,
+				code_2.clone()
+			),
+			Error::<Test>::CannotUpgradeCode,
+		);
+		assert_eq!(AuthorizedCodeHash::<Test>::get(para_a), Some(code_1_hash));
+		assert!(CurrentCodeHash::<Test>::get(para_a).is_none());
+		check_code_is_not_stored(&code_1);
+		check_code_is_not_stored(&code_2);
+		check_code_is_not_stored(&code_3);
+
+		// non-root can apply authorized code
+		assert_ok!(Paras::apply_authorized_force_set_current_code(
+			non_root.clone(),
+			para_a,
+			code_1.clone()
+		));
+
+		// check authorized code was applied
+		assert!(AuthorizedCodeHash::<Test>::get(para_a).is_none());
+		assert_eq!(CurrentCodeHash::<Test>::get(para_a), Some(code_1_hash));
+		check_code_is_stored(&code_1);
+		check_code_is_not_stored(&code_2);
+		check_code_is_not_stored(&code_3);
+
+		// authorize multiple without apply:
+		// authorize code_2_hash
+		assert_ok!(Paras::authorize_force_set_current_code_hash(root.clone(), para_a, code_2_hash));
+		assert_eq!(AuthorizedCodeHash::<Test>::get(para_a), Some(code_2_hash));
+		assert_eq!(CurrentCodeHash::<Test>::get(para_a), Some(code_1_hash));
+		check_code_is_stored(&code_1);
+		check_code_is_not_stored(&code_2);
+		check_code_is_not_stored(&code_3);
+		// authorize code_3_hash
+		assert_ok!(Paras::authorize_force_set_current_code_hash(root, para_a, code_3_hash));
+		assert_eq!(AuthorizedCodeHash::<Test>::get(para_a), Some(code_3_hash));
+		assert_eq!(CurrentCodeHash::<Test>::get(para_a), Some(code_1_hash));
+		check_code_is_stored(&code_1);
+		check_code_is_not_stored(&code_2);
+		check_code_is_not_stored(&code_3);
+
+		// cannot apply older ones
+		assert_err!(
+			Paras::apply_authorized_force_set_current_code(
+				non_root.clone(),
+				para_a,
+				code_1.clone()
+			),
+			Error::<Test>::CannotUpgradeCode,
+		);
+		assert_err!(
+			Paras::apply_authorized_force_set_current_code(
+				non_root.clone(),
+				para_a,
+				code_2.clone()
+			),
+			Error::<Test>::CannotUpgradeCode,
+		);
+
+		// apply just authorized
+		assert_ok!(Paras::apply_authorized_force_set_current_code(
+			non_root.clone(),
+			para_a,
+			code_3.clone()
+		));
+		assert!(AuthorizedCodeHash::<Test>::get(para_a).is_none());
+		assert_eq!(CurrentCodeHash::<Test>::get(para_a), Some(code_3_hash));
+		check_code_is_stored(&code_1);
+		check_code_is_not_stored(&code_2);
+		check_code_is_stored(&code_3);
+	})
+}
+
+#[test]
+fn force_set_current_code_works() {
+	new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+		let para_a = ParaId::from(111);
+		let code_1 = ValidationCode(vec![1]);
+		let code_1_hash = code_1.hash();
+		let root = RuntimeOrigin::root();
+		let non_root = RuntimeOrigin::signed(1);
+
+		// check before
+		assert!(CurrentCodeHash::<Test>::get(para_a).is_none());
+		check_code_is_not_stored(&code_1);
+
+		// non-root user cannot execute
+		assert_err!(
+			Paras::force_set_current_code(non_root, para_a, code_1.clone()),
+			DispatchError::BadOrigin,
+		);
+		// root can execute
+		assert_ok!(Paras::force_set_current_code(root, para_a, code_1.clone()));
+
+		// check after
+		assert_eq!(CurrentCodeHash::<Test>::get(para_a), Some(code_1_hash));
+		check_code_is_stored(&code_1);
+	})
+}

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

@@ -295,4 +295,20 @@ impl<T: frame_system::Config> polkadot_runtime_parachains::paras::WeightInfo for
 			.saturating_add(T::DbWeight::get().reads(3))
 			.saturating_add(T::DbWeight::get().writes(1))
 	}
+	fn authorize_force_set_current_code_hash() -> Weight {
+		// TODO: FAIL-CI fresh
+		Weight::from_parts(120_274_000, 0)
+			.saturating_add(Weight::from_parts(0, 30147))
+			.saturating_add(T::DbWeight::get().reads(3))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+	fn apply_authorized_force_set_current_code(c: u32, ) -> Weight {
+		// TODO: FAIL-CI fresh
+		Weight::from_parts(33_700_000, 0)
+			.saturating_add(Weight::from_parts(0, 11774))
+			// Standard Error: 10
+			.saturating_add(Weight::from_parts(2_659, 0).saturating_mul(c.into()))
+			.saturating_add(T::DbWeight::get().reads(4))
+			.saturating_add(T::DbWeight::get().writes(6))
+	}
 }

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

@@ -286,4 +286,20 @@ impl<T: frame_system::Config> polkadot_runtime_parachains::paras::WeightInfo for
 			.saturating_add(T::DbWeight::get().reads(3))
 			.saturating_add(T::DbWeight::get().writes(1))
 	}
+	fn authorize_force_set_current_code_hash() -> Weight {
+		// TODO: FAIL-CI fresh
+		Weight::from_parts(120_274_000, 0)
+			.saturating_add(Weight::from_parts(0, 30147))
+			.saturating_add(T::DbWeight::get().reads(3))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+	fn apply_authorized_force_set_current_code(c: u32, ) -> Weight {
+		// TODO: FAIL-CI fresh
+		Weight::from_parts(33_700_000, 0)
+			.saturating_add(Weight::from_parts(0, 11774))
+			// Standard Error: 10
+			.saturating_add(Weight::from_parts(2_659, 0).saturating_mul(c.into()))
+			.saturating_add(T::DbWeight::get().reads(4))
+			.saturating_add(T::DbWeight::get().writes(6))
+	}
 }