Look at the upgrade go-ahead and restriction signals

pallets/parachain-system/src/lib.rs

@@ -54,6 +54,7 @@ use sp_runtime::{
 };
 use sp_std::{cmp, collections::btree_map::BTreeMap, prelude::*};
 
+mod migration;
 mod relay_state_snapshot;
 #[macro_use]
 pub mod validate_block;
@@ -94,6 +95,7 @@ pub mod pallet {
 	use frame_system::pallet_prelude::*;
 
 	#[pallet::pallet]
+	#[pallet::storage_version(migration::STORAGE_VERSION)]
 	pub struct Pallet<T>(_);
 
 	#[pallet::config]
@@ -129,8 +131,13 @@ pub mod pallet {
 
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+		fn on_runtime_upgrade() -> Weight {
+			migration::on_runtime_upgrade::<T>()
+		}
+
 		fn on_finalize(_: T::BlockNumber) {
 			<DidSetValidationCode<T>>::kill();
+			<UpgradeRestrictionSignal<T>>::kill();
 
 			assert!(
 				<ValidationData<T>>::exists(),
@@ -268,26 +275,6 @@ pub mod pallet {
 
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {
-		/// 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.
-		#[pallet::weight((0, DispatchClass::Operational))]
-		pub fn set_upgrade_block(
-			origin: OriginFor<T>,
-			relay_chain_block: RelayChainBlockNumber,
-		) -> DispatchResult {
-			ensure_root(origin)?;
-			if <PendingRelayChainBlockNumber<T>>::get().is_some() {
-				<PendingRelayChainBlockNumber<T>>::put(relay_chain_block);
-				Ok(())
-			} else {
-				Err(Error::<T>::NotScheduled.into())
-			}
-		}
-
 		/// Set the current validation data.
 		///
 		/// This should be invoked exactly once per block. It will panic at the finalization
@@ -318,25 +305,41 @@ pub mod pallet {
 
 			Self::validate_validation_data(&vfp);
 
+			let relay_state_proof = RelayChainStateProof::new(
+				T::SelfParaId::get(),
+				vfp.relay_parent_storage_root,
+				relay_chain_state,
+			)
+			.expect("Invalid relay chain state proof");
+
 			// 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) = <PendingRelayChainBlockNumber<T>>::get() {
-				if vfp.relay_parent_number >= apply_block {
-					<PendingRelayChainBlockNumber<T>>::kill();
+			let upgrade_go_ahead_signal = relay_state_proof
+				.read_upgrade_go_ahead_signal()
+				.expect("Invalid upgrade go ahead signal");
+			match upgrade_go_ahead_signal {
+				Some(relay_chain::v1::UpgradeGoAhead::GoAhead) => {
+					assert!(
+						<PendingValidationCode<T>>::exists(),
+						"No new validation function found in storage, GoAhead signal is not expected",
+					);
 					let validation_code = <PendingValidationCode<T>>::take();
-					<LastUpgrade<T>>::put(&apply_block);
+
 					Self::put_parachain_code(&validation_code);
 					Self::deposit_event(Event::ValidationFunctionApplied(vfp.relay_parent_number));
-				}
+				},
+				Some(relay_chain::v1::UpgradeGoAhead::Abort) => {
+					<PendingValidationCode<T>>::kill();
+					Self::deposit_event(Event::ValidationFunctionDiscarded);
+				},
+				None => {},
 			}
-
-			let relay_state_proof = RelayChainStateProof::new(
-				T::SelfParaId::get(),
-				vfp.relay_parent_storage_root,
-				relay_chain_state,
-			)
-			.expect("Invalid relay chain state proof");
+			<UpgradeRestrictionSignal<T>>::put(
+				relay_state_proof
+					.read_upgrade_restriction_signal()
+					.expect("Invalid upgrade restriction signal"),
+			);
 
 			let host_config = relay_state_proof
 				.read_abridged_host_configuration()
@@ -401,11 +404,12 @@ pub mod pallet {
 	#[pallet::event]
 	#[pallet::generate_deposit(pub(super) fn deposit_event)]
 	pub enum Event<T: Config> {
-		/// The validation function has been scheduled to apply as of the contained relay chain
-		/// block number.
-		ValidationFunctionStored(RelayChainBlockNumber),
+		/// The validation function has been scheduled to apply.
+		ValidationFunctionStored,
 		/// The validation function was applied as of the contained relay chain block number.
 		ValidationFunctionApplied(RelayChainBlockNumber),
+		/// The relay-chain aborted the upgrade process.
+		ValidationFunctionDiscarded,
 		/// An upgrade has been authorized.
 		UpgradeAuthorized(T::Hash),
 		/// Some downward messages have been received and will be processed.
@@ -437,22 +441,27 @@ pub mod pallet {
 		Unauthorized,
 	}
 
-	/// We need to store the new validation function for the span between
-	/// setting it and applying it. If it has a
-	/// value, then [`PendingValidationCode`] must have a real value, and
-	/// together will coordinate the block number where the upgrade will happen.
-	#[pallet::storage]
-	pub(super) type PendingRelayChainBlockNumber<T: Config> =
-		StorageValue<_, 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.
+	/// In case of a scheduled upgrade, this storage field contains the validation code to be applied.
+	///
+	/// As soon as the relay chain gives us the go-ahead signal, we will overwrite the [`:code`][well_known_keys::CODE]
+	/// which will result the next block process with the new validation code. This concludes the upgrade process.
+	///
+	/// [well_known_keys::CODE]: sp_core::storage::well_known_keys::CODE
 	#[pallet::storage]
 	#[pallet::getter(fn new_validation_function)]
 	pub(super) type PendingValidationCode<T: Config> = StorageValue<_, Vec<u8>, ValueQuery>;
 
+	/// Validation code that is set by the parachain and is to be communicated to collator and
+	/// consequently the relay-chain.
+	///
+	/// This will be cleared in `on_initialize` of each new block if no other pallet already set
+	/// the value.
+	#[pallet::storage]
+	pub(super) type NewValidationCode<T: Config> = StorageValue<_, Vec<u8>, OptionQuery>;
+
 	/// The [`PersistedValidationData`] set for this block.
+	/// This value is expected to be set only once per block and it's never stored
+	/// in the trie.
 	#[pallet::storage]
 	#[pallet::getter(fn validation_data)]
 	pub(super) type ValidationData<T: Config> = StorageValue<_, PersistedValidationData>;
@@ -461,9 +470,16 @@ pub mod pallet {
 	#[pallet::storage]
 	pub(super) type DidSetValidationCode<T: Config> = StorageValue<_, bool, ValueQuery>;
 
-	/// The last relay parent block number at which we signalled the code upgrade.
+	/// An option which indicates if the relay-chain restricts signalling a validation code upgrade.
+	/// In other words, if this is `Some` and [`NewValidationCode`] is `Some` then the produced
+	/// candidate will be invalid.
+	///
+	/// This storage item is a mirror of the corresponding value for the current parachain from the
+	/// relay-chain. This value is ephemeral which means it doesn't hit the storage. This value is
+	/// set after the inherent.
 	#[pallet::storage]
-	pub(super) type LastUpgrade<T: Config> = StorageValue<_, relay_chain::BlockNumber, ValueQuery>;
+	pub(super) type UpgradeRestrictionSignal<T: Config> =
+		StorageValue<_, Option<relay_chain::v1::UpgradeRestriction>, ValueQuery>;
 
 	/// The snapshot of some state related to messaging relevant to the current parachain as per
 	/// the relay parent.
@@ -507,13 +523,6 @@ pub mod pallet {
 	#[pallet::storage]
 	pub(super) type ProcessedDownwardMessages<T: Config> = StorageValue<_, u32, ValueQuery>;
 
-	/// New validation code that was set in a block.
-	///
-	/// This will be cleared in `on_initialize` of each new block if no other pallet already set
-	/// the value.
-	#[pallet::storage]
-	pub(super) type NewValidationCode<T: Config> = StorageValue<_, Vec<u8>, OptionQuery>;
-
 	/// HRMP watermark that was set in a block.
 	///
 	/// This will be cleared in `on_initialize` of each new block.
@@ -860,48 +869,27 @@ impl<T: Config> Pallet<T> {
 		<HostConfiguration<T>>::get().map(|cfg| cfg.max_code_size)
 	}
 
-	/// Returns if a PVF/runtime upgrade could be signalled at the current block, and if so
-	/// when the new code will take the effect.
-	fn code_upgrade_allowed(
-		vfp: &PersistedValidationData,
-		cfg: &AbridgedHostConfiguration,
-	) -> Option<relay_chain::BlockNumber> {
-		if <PendingRelayChainBlockNumber<T>>::get().is_some() {
-			// There is already upgrade scheduled. Upgrade is not allowed.
-			return None
-		}
-
-		let relay_blocks_since_last_upgrade =
-			vfp.relay_parent_number.saturating_sub(<LastUpgrade<T>>::get());
-
-		if relay_blocks_since_last_upgrade <= cfg.validation_upgrade_frequency {
-			// The cooldown after the last upgrade hasn't elapsed yet. Upgrade is not allowed.
-			return None
-		}
-
-		Some(vfp.relay_parent_number + cfg.validation_upgrade_delay)
-	}
-
 	/// The implementation of the runtime upgrade functionality for parachains.
 	fn set_code_impl(validation_function: Vec<u8>) -> DispatchResult {
+		// Ensure that `ValidationData` exists. We do not care about the validation data per se,
+		// but we do care about the [`UpgradeRestrictionSignal`] which arrives with the same inherent.
+		ensure!(<ValidationData<T>>::exists(), Error::<T>::ValidationDataNotAvailable,);
+		ensure!(<UpgradeRestrictionSignal<T>>::get().is_none(), Error::<T>::ProhibitedByPolkadot);
+
 		ensure!(!<PendingValidationCode<T>>::exists(), Error::<T>::OverlappingUpgrades);
-		let vfp = Self::validation_data().ok_or(Error::<T>::ValidationDataNotAvailable)?;
 		let cfg = Self::host_configuration().ok_or(Error::<T>::HostConfigurationNotAvailable)?;
 		ensure!(validation_function.len() <= cfg.max_code_size as usize, Error::<T>::TooBig);
-		let apply_block =
-			Self::code_upgrade_allowed(&vfp, &cfg).ok_or(Error::<T>::ProhibitedByPolkadot)?;
 
 		// When a code upgrade is scheduled, it has to be applied in two
 		// places, synchronized: both polkadot and the individual parachain
 		// have to upgrade on the same relay chain block.
 		//
 		// `notify_polkadot_of_pending_upgrade` notifies polkadot; the `PendingValidationCode`
 		// storage keeps track locally for the parachain upgrade, which will
-		// be applied later.
+		// be applied later: when the relay-chain communicates go-ahead signal to us.
 		Self::notify_polkadot_of_pending_upgrade(&validation_function);
-		<PendingRelayChainBlockNumber<T>>::put(apply_block);
 		<PendingValidationCode<T>>::put(validation_function);
-		Self::deposit_event(Event::ValidationFunctionStored(apply_block));
+		Self::deposit_event(Event::ValidationFunctionStored);
 
 		Ok(())
 	}

pallets/parachain-system/src/migration.rs

@@ -0,0 +1,50 @@
+// Copyright 2021 Parity Technologies (UK) Ltd.
+// This file is part of Cumulus.
+
+// Cumulus is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Cumulus is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Cumulus.  If not, see <http://www.gnu.org/licenses/>.
+
+use crate::{Config, Pallet};
+use frame_support::{
+	traits::{Get, StorageVersion},
+	weights::Weight,
+};
+
+/// The current storage version.
+pub const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
+
+/// Call this during the next runtime upgrade for this module.
+pub fn on_runtime_upgrade<T: Config>() -> Weight {
+	let mut weight: Weight = 0;
+
+	if StorageVersion::get::<Pallet<T>>() == 0 {
+		weight = weight
+			.saturating_add(v1::migrate::<T>())
+			.saturating_add(T::DbWeight::get().writes(1));
+		StorageVersion::new(1).put::<Pallet<T>>();
+	}
+
+	weight
+}
+
+/// V1: `LastUpgrade` block number is removed from the storage since the upgrade
+/// mechanism now uses signals instead of block offsets.
+mod v1 {
+	use crate::{Config, Pallet};
+	use frame_support::{migration::remove_storage_prefix, pallet_prelude::*};
+
+	pub fn migrate<T: Config>() -> Weight {
+		remove_storage_prefix(<Pallet<T>>::name().as_bytes(), b"LastUpgrade", b"");
+		T::DbWeight::get().writes(1)
+	}
+}

pallets/parachain-system/src/relay_state_snapshot.rs

@@ -65,6 +65,10 @@ pub enum Error {
 	RootMismatch,
 	/// The slot cannot be extracted.
 	Slot(ReadEntryErr),
+	/// The upgrade go-ahead signal cannot be read.
+	UpgradeGoAhead(ReadEntryErr),
+	/// The upgrade restriction signal cannot be read.
+	UpgradeRestriction(ReadEntryErr),
 	/// The host configuration cannot be extracted.
 	Config(ReadEntryErr),
 	/// The DMQ MQC head cannot be extracted.
@@ -109,6 +113,23 @@ where
 		.ok_or(ReadEntryErr::Absent)
 }
 
+/// Read an optional entry given by the key and try to decode it.
+/// Returns `None` if the value specified by the key according to the proof is empty.
+///
+/// Returns `Err` in case the backend can't return the value under the specific key (likely due to
+/// a malformed proof) or if the value couldn't be decoded.
+fn read_optional_entry<T, B>(backend: &B, key: &[u8]) -> Result<Option<T>, ReadEntryErr>
+where
+	T: Decode,
+	B: Backend<HashFor<relay_chain::Block>>,
+{
+	match read_entry(backend, key, None) {
+		Ok(v) => Ok(Some(v)),
+		Err(ReadEntryErr::Absent) => Ok(None),
+		Err(err) => Err(err),
+	}
+}
+
 /// A state proof extracted from the relay chain.
 ///
 /// This state proof is extracted from the relay chain block we are building on top of.
@@ -219,4 +240,37 @@ impl RelayChainStateProof {
 		read_entry(&self.trie_backend, relay_chain::well_known_keys::CURRENT_SLOT, None)
 			.map_err(Error::Slot)
 	}
+
+	/// Read the go-ahead signal for the upgrade from the relay chain state proof.
+	///
+	/// The go-ahead specifies whether the parachain can apply the upgrade or should abort it. If
+	/// the value is absent then there is either no judgment by the relay chain yet or no upgrade
+	/// is pending.
+	///
+	/// Returns an error if anything failed at reading or decoding.
+	pub fn read_upgrade_go_ahead_signal(
+		&self,
+	) -> Result<Option<relay_chain::v1::UpgradeGoAhead>, Error> {
+		read_optional_entry(
+			&self.trie_backend,
+			&relay_chain::well_known_keys::upgrade_go_ahead_signal(self.para_id),
+		)
+		.map_err(Error::UpgradeGoAhead)
+	}
+
+	/// Read the upgrade restriction signal for the upgrade from the relay chain state proof.
+	///
+	/// If the upgrade restriction is not `None`, then the parachain cannot signal an upgrade at
+	/// this block.
+	///
+	/// Returns an error if anything failed at reading or decoding.
+	pub fn read_upgrade_restriction_signal(
+		&self,
+	) -> Result<Option<relay_chain::v1::UpgradeRestriction>, Error> {
+		read_optional_entry(
+			&self.trie_backend,
+			&relay_chain::well_known_keys::upgrade_restriction_signal(self.para_id),
+		)
+		.map_err(Error::UpgradeRestriction)
+	}
 }