Look at the upgrade go-ahead and restriction signals (#517)

* Look at the upgrade go-ahead and restriction signals

* Update Cargo.toml

* Drop old docs for validation code

* Update tests

* Fix typo

* Add doc-comments for read_optional_entry

* Add a note about ValidationData

* Introduce migration for removing unused storage entry

* Fix indentation

* Use intra-doc link syntax

* Double-check that GoAhead signal is not spurious

* fmt

* Drop commented code

* Fix typos

Co-authored-by: Alexander Popiak <alexander.popiak@parity.io>

* Add a weight for StorageVersion write

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

Co-authored-by: Chris Sosnin <chris125_@live.com>
Co-authored-by: Chris Sosnin <48099298+slumber@users.noreply.github.com>
Co-authored-by: Alexander Popiak <alexander.popiak@parity.io>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

cumulus/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,19 +305,6 @@ pub mod pallet {
 
 			Self::validate_validation_data(&vfp);
 
-			// 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 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));
-				}
-			}
-
 			let relay_state_proof = RelayChainStateProof::new(
 				T::SelfParaId::get(),
 				vfp.relay_parent_storage_root,
@@ -338,6 +312,35 @@ pub mod pallet {
 			)
 			.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.
+			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();
+
+					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 => {},
+			}
+			<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()
 				.expect("Invalid host configuration in relay chain state proof");
@@ -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,36 +869,16 @@ 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
@@ -897,11 +886,10 @@ impl<T: Config> Pallet<T> {
 		//
 		// `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(())
 	}

cumulus/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)
+	}
+}

cumulus/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)
+	}
 }

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

@@ -384,8 +384,10 @@ fn block_tests_run_on_drop() {
 #[test]
 fn events() {
 	BlockTests::new()
-		.with_relay_sproof_builder(|_, _, builder| {
-			builder.host_config.validation_upgrade_delay = 1000;
+		.with_relay_sproof_builder(|_, block_number, builder| {
+			if block_number > 123 {
+				builder.upgrade_go_ahead = Some(relay_chain::v1::UpgradeGoAhead::GoAhead);
+			}
 		})
 		.add_with_post_test(
 			123,
@@ -396,7 +398,7 @@ fn events() {
 				let events = System::events();
 				assert_eq!(
 					events[0].event,
-					Event::ParachainSystem(crate::Event::ValidationFunctionStored(1123).into())
+					Event::ParachainSystem(crate::Event::ValidationFunctionStored.into())
 				);
 			},
 		)
@@ -433,6 +435,11 @@ fn non_overlapping() {
 #[test]
 fn manipulates_storage() {
 	BlockTests::new()
+		.with_relay_sproof_builder(|_, block_number, builder| {
+			if block_number > 123 {
+				builder.upgrade_go_ahead = Some(relay_chain::v1::UpgradeGoAhead::GoAhead);
+			}
+		})
 		.add(123, || {
 			assert!(
 				!<PendingValidationCode<Test>>::exists(),
@@ -453,6 +460,34 @@ fn manipulates_storage() {
 		);
 }
 
+#[test]
+fn aborted_upgrade() {
+	BlockTests::new()
+		.with_relay_sproof_builder(|_, block_number, builder| {
+			if block_number > 123 {
+				builder.upgrade_go_ahead = Some(relay_chain::v1::UpgradeGoAhead::Abort);
+			}
+		})
+		.add(123, || {
+			assert_ok!(System::set_code(RawOrigin::Root.into(), Default::default()));
+		})
+		.add_with_post_test(
+			1234,
+			|| {},
+			|| {
+				assert!(
+					!<PendingValidationCode<Test>>::exists(),
+					"validation function must have been unset"
+				);
+				let events = System::events();
+				assert_eq!(
+					events[0].event,
+					Event::ParachainSystem(crate::Event::ValidationFunctionDiscarded.into())
+				);
+			},
+		);
+}
+
 #[test]
 fn checks_size() {
 	BlockTests::new()

cumulus/primitives/parachain-inherent/src/client_side.rs

@@ -172,6 +172,8 @@ fn collect_relay_storage_proof(
 	relevant_keys.push(relay_well_known_keys::relay_dispatch_queue_size(para_id));
 	relevant_keys.push(relay_well_known_keys::hrmp_ingress_channel_index(para_id));
 	relevant_keys.push(relay_well_known_keys::hrmp_egress_channel_index(para_id));
+	relevant_keys.push(relay_well_known_keys::upgrade_go_ahead_signal(para_id));
+	relevant_keys.push(relay_well_known_keys::upgrade_restriction_signal(para_id));
 	relevant_keys.extend(ingress_channels.into_iter().map(|sender| {
 		relay_well_known_keys::hrmp_channels(HrmpChannelId { sender, recipient: para_id })
 	}));

cumulus/test/relay-sproof-builder/src/lib.rs

@@ -17,12 +17,12 @@
 use cumulus_primitives_core::{
 	relay_chain, AbridgedHostConfiguration, AbridgedHrmpChannel, ParaId,
 };
+use polkadot_primitives::v1::UpgradeGoAhead;
 use sp_runtime::traits::HashFor;
 use sp_state_machine::MemoryDB;
 use sp_std::collections::btree_map::BTreeMap;
 
 /// Builds a sproof (portmanteau of 'spoof' and 'proof') of the relay chain state.
-#[derive(Clone)]
 pub struct RelayStateSproofBuilder {
 	/// The para id of the current parachain.
 	///
@@ -36,6 +36,7 @@ pub struct RelayStateSproofBuilder {
 
 	pub host_config: AbridgedHostConfiguration,
 	pub dmq_mqc_head: Option<relay_chain::Hash>,
+	pub upgrade_go_ahead: Option<UpgradeGoAhead>,
 	pub relay_dispatch_queue_size: Option<(u32, u32)>,
 	pub hrmp_ingress_channel_index: Option<Vec<ParaId>>,
 	pub hrmp_egress_channel_index: Option<Vec<ParaId>>,
@@ -59,6 +60,7 @@ impl Default for RelayStateSproofBuilder {
 				validation_upgrade_delay: 6,
 			},
 			dmq_mqc_head: None,
+			upgrade_go_ahead: None,
 			relay_dispatch_queue_size: None,
 			hrmp_ingress_channel_index: None,
 			hrmp_egress_channel_index: None,
@@ -68,6 +70,26 @@ impl Default for RelayStateSproofBuilder {
 	}
 }
 
+// TODO: derive `Copy` and `Clone` for `UpgradeGoAhead` to avoid manual implementation.
+impl Clone for RelayStateSproofBuilder {
+	fn clone(&self) -> Self {
+		RelayStateSproofBuilder {
+			para_id: self.para_id,
+			host_config: self.host_config.clone(),
+			dmq_mqc_head: self.dmq_mqc_head.clone(),
+			upgrade_go_ahead: self.upgrade_go_ahead.as_ref().map(|u| match u {
+				UpgradeGoAhead::Abort => UpgradeGoAhead::Abort,
+				UpgradeGoAhead::GoAhead => UpgradeGoAhead::GoAhead,
+			}),
+			relay_dispatch_queue_size: self.relay_dispatch_queue_size,
+			hrmp_ingress_channel_index: self.hrmp_ingress_channel_index.clone(),
+			hrmp_egress_channel_index: self.hrmp_egress_channel_index.clone(),
+			hrmp_channels: self.hrmp_channels.clone(),
+			current_slot: self.current_slot.clone(),
+		}
+	}
+}
+
 impl RelayStateSproofBuilder {
 	/// Returns a mutable reference to HRMP channel metadata for a channel (`sender`, `self.para_id`).
 	///
@@ -120,6 +142,12 @@ impl RelayStateSproofBuilder {
 					relay_dispatch_queue_size.encode(),
 				);
 			}
+			if let Some(upgrade_go_ahead) = self.upgrade_go_ahead {
+				insert(
+					relay_chain::well_known_keys::upgrade_go_ahead_signal(self.para_id),
+					upgrade_go_ahead.encode(),
+				);
+			}
 			if let Some(hrmp_ingress_channel_index) = self.hrmp_ingress_channel_index {
 				let mut sorted = hrmp_ingress_channel_index.clone();
 				sorted.sort();