Improve Storage and Add set_upgrade_block to Validation Function Upgrade (#383)

* set_upgrade_block

* Update lib.rs

* Use Two Storage Items for Validation Function Upgrade

* note issue #374

* fix docs nits

* Apply suggestions from code review

* Update pallets/parachain-system/src/lib.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

cumulus/pallets/parachain-system/src/lib.rs

@@ -77,10 +77,16 @@ pub trait Config: frame_system::Config {
 // This pallet's storage items.
 decl_storage! {
 	trait Store for Module<T: Config> as ParachainSystem {
-		// we need to store the new validation function for the span between
+		/// We need to store the new validation function for the span between
-		// setting it and applying it.
+		/// setting it and applying it. If it has a
-		PendingValidationFunction get(fn new_validation_function):
+		/// value, then [`PendingValidationFunction`] must have a real value, and
-			Option<(RelayChainBlockNumber, Vec<u8>)>;
+		/// together will coordinate the block number where the upgrade will happen.
+		PendingRelayChainBlockNumber: Option<RelayChainBlockNumber>;
+
+		/// The new validation function we will upgrade to when the relay chain
+		/// reaches [`PendingRelayChainBlockNumber`]. A real validation function must
+		/// exist here as long as [`PendingRelayChainBlockNumber`] is set.
+		PendingValidationFunction get(fn new_validation_function): Vec<u8>;
 
 		/// The [`PersistedValidationData`] set for this block.
 		ValidationData get(fn validation_data): Option<PersistedValidationData>;
@@ -139,6 +145,7 @@ decl_module! {
 		fn deposit_event() = default;
 
 		// TODO: figure out a better weight than this
+		// TODO: Bring back the correct validation checks: #374
 		#[weight = (0, DispatchClass::Operational)]
 		pub fn schedule_upgrade(origin, validation_function: Vec<u8>) {
 			ensure_root(origin)?;
@@ -155,6 +162,21 @@ decl_module! {
 			Self::schedule_upgrade_impl(validation_function)?;
 		}
 
+		/// Force an already scheduled validation function upgrade to happen on a particular block.
+		///
+		/// Note that coordinating this block for the upgrade has to happen independently on the relay
+		/// chain and this parachain. Synchronizing the block for the upgrade is sensitive, and this
+		/// bypasses all checks and and normal protocols. Very easy to brick your chain if done wrong.
+		#[weight = (0, DispatchClass::Operational)]
+		pub fn set_upgrade_block(origin, relay_chain_block: RelayChainBlockNumber) {
+			ensure_root(origin)?;
+			if let Some(_old_block) = PendingRelayChainBlockNumber::get() {
+				PendingRelayChainBlockNumber::put(relay_chain_block);
+			} else {
+				return Err(Error::<T>::NotScheduled.into())
+			}
+		}
+
 		/// Set the current validation data.
 		///
 		/// This should be invoked exactly once per block. It will panic at the finalization
@@ -184,9 +206,10 @@ decl_module! {
 			// initialization logic: we know that this runs exactly once every block,
 			// which means we can put the initialization logic here to remove the
 			// sequencing problem.
-			if let Some((apply_block, validation_function)) = PendingValidationFunction::get() {
+			if let Some(apply_block) = PendingRelayChainBlockNumber::get() {
 				if vfp.relay_parent_number >= apply_block {
-					PendingValidationFunction::kill();
+					PendingRelayChainBlockNumber::kill();
+					let validation_function = PendingValidationFunction::take();
 					LastUpgrade::put(&apply_block);
 					Self::put_parachain_code(&validation_function);
 					Self::deposit_event(Event::ValidationFunctionApplied(vfp.relay_parent_number));
@@ -615,7 +638,7 @@ impl<T: Config> Module<T> {
 		vfp: &PersistedValidationData,
 		cfg: &AbridgedHostConfiguration,
 	) -> Option<relay_chain::BlockNumber> {
-		if PendingValidationFunction::get().is_some() {
+		if PendingRelayChainBlockNumber::get().is_some() {
 			// There is already upgrade scheduled. Upgrade is not allowed.
 			return None;
 		}
@@ -654,7 +677,8 @@ impl<T: Config> Module<T> {
 		// storage keeps track locally for the parachain upgrade, which will
 		// be applied later.
 		Self::notify_polkadot_of_pending_upgrade(&validation_function);
-		PendingValidationFunction::put((apply_block, validation_function));
+		PendingRelayChainBlockNumber::put(apply_block);
+		PendingValidationFunction::put(validation_function);
 		Self::deposit_event(Event::ValidationFunctionStored(apply_block));
 
 		Ok(())
@@ -870,6 +894,8 @@ decl_error! {
 		/// That means that one or more channels had at least some of the submitted messages altered,
 		/// omitted or added illegaly.
 		HrmpMqcMismatch,
+		/// No validation function upgrade is currently scheduled.
+		NotScheduled,
 	}
 }