Adds support for returning a custom header from validate_block (#825)

* Adds support for returning a custom header from `validate_block`

This adds support for returning a custom header from `validate_block`. Before this, we always
returned the header of the block that was validated (and still do it by default). However, after
this pr it is possible to set a custom header or better custom head data that will be returned
instead from `validate_block`. This can be for example when a chain wants to fork.

* FMT

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

@@ -236,8 +236,9 @@ pub mod pallet {
 			HrmpWatermark::<T>::kill();
 			UpwardMessages::<T>::kill();
 			HrmpOutboundMessages::<T>::kill();
+			CustomValidationHeadData::<T>::kill();
 
-			weight += T::DbWeight::get().writes(5);
+			weight += T::DbWeight::get().writes(6);
 
 			// Here, in `on_initialize` we must report the weight for both `on_initialize` and
 			// `on_finalize`.
@@ -567,6 +568,12 @@ pub mod pallet {
 	#[pallet::storage]
 	pub(super) type AuthorizedUpgrade<T: Config> = StorageValue<_, T::Hash>;
 
+	/// A custom head data that should be returned as result of `validate_block`.
+	///
+	/// See [`Pallet::set_custom_validation_head_data`] for more information.
+	#[pallet::storage]
+	pub(super) type CustomValidationHeadData<T: Config> = StorageValue<_, Vec<u8>, OptionQuery>;
+
 	#[pallet::inherent]
 	impl<T: Config> ProvideInherent for Pallet<T> {
 		type Call = Call<T>;
@@ -692,6 +699,9 @@ impl<T: Config> Pallet<T> {
 	/// import, this is a no-op.
 	///
 	/// # Panics
+	///
+	/// Panics while validating the `PoV` on the relay chain if the [`PersistedValidationData`]
+	/// passed by the block author was incorrect.
 	fn validate_validation_data(validation_data: &PersistedValidationData) {
 		validate_block::with_validation_params(|params| {
 			assert_eq!(
@@ -714,8 +724,10 @@ impl<T: Config> Pallet<T> {
 	/// Checks if the sequence of the messages is valid, dispatches them and communicates the
 	/// number of processed messages to the collator via a storage update.
 	///
-	/// **Panics** if it turns out that after processing all messages the Message Queue Chain
-	///            hash doesn't match the expected.
+	/// # Panics
+	///
+	/// If it turns out that after processing all messages the Message Queue Chain
+	/// hash doesn't match the expected.
 	fn process_inbound_downward_messages(
 		expected_dmq_mqc_head: relay_chain::Hash,
 		downward_messages: Vec<InboundDownwardMessage>,
@@ -907,6 +919,22 @@ impl<T: Config> Pallet<T> {
 			new_validation_code: NewValidationCode::<T>::get().map(Into::into),
 		}
 	}
+
+	/// Set a custom head data that should be returned as result of `validate_block`.
+	///
+	/// This will overwrite the head data that is returned as result of `validate_block` while
+	/// validating a `PoV` on the relay chain. Normally the head data that is being returned
+	/// by `validate_block` is the header of the block that is validated, thus it can be
+	/// enacted as the new best block. However, for features like forking it can be useful
+	/// to overwrite the head data with a custom header.
+	///
+	/// # Attention
+	///
+	/// This should only be used when you are sure what you are doing as this can brick
+	/// your Parachain.
+	pub fn set_custom_validation_head_data(head_data: Vec<u8>) {
+		CustomValidationHeadData::<T>::put(head_data);
+	}
 }
 
 pub struct ParachainSetCode<T>(sp_std::marker::PhantomData<T>);

cumulus/pallets/parachain-system/src/validate_block/implementation.rs

@@ -161,6 +161,13 @@ where
 			let horizontal_messages = crate::HrmpOutboundMessages::<PSC>::get();
 			let hrmp_watermark = crate::HrmpWatermark::<PSC>::get();
 
+			let head_data =
+				if let Some(custom_head_data) = crate::CustomValidationHeadData::<PSC>::get() {
+					HeadData(custom_head_data)
+				} else {
+					head_data
+				};
+
 			ValidationResult {
 				head_data,
 				new_validation_code: new_validation_code.map(Into::into),

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

@@ -17,7 +17,8 @@
 use codec::{Decode, Encode};
 use cumulus_primitives_core::{ParachainBlockData, PersistedValidationData};
 use cumulus_test_client::{
-	runtime::{Block, Hash, Header, UncheckedExtrinsic, WASM_BINARY},
+	generate_extrinsic,
+	runtime::{Block, Hash, Header, TestPalletCall, UncheckedExtrinsic, WASM_BINARY},
 	transfer, BlockData, BuildParachainBlockData, Client, DefaultTestClientBuilderExt, HeadData,
 	InitBlockBuilder, TestClientBuilder, TestClientBuilderExt, ValidationParams,
 };
@@ -26,11 +27,11 @@ use sp_keyring::AccountKeyring::*;
 use sp_runtime::{generic::BlockId, traits::Header as HeaderT};
 use std::{env, process::Command};
 
-fn call_validate_block(
+fn call_validate_block_encoded_header(
 	parent_head: Header,
 	block_data: ParachainBlockData<Block>,
 	relay_parent_storage_root: Hash,
-) -> cumulus_test_client::ExecutorResult<Header> {
+) -> cumulus_test_client::ExecutorResult<Vec<u8>> {
 	cumulus_test_client::validate_block(
 		ValidationParams {
 			block_data: BlockData(block_data.encode()),
@@ -40,7 +41,16 @@ fn call_validate_block(
 		},
 		&WASM_BINARY.expect("You need to build the WASM binaries to run the tests!"),
 	)
-	.map(|v| Header::decode(&mut &v.head_data.0[..]).expect("Decodes `Header`."))
+	.map(|v| v.head_data.0)
+}
+
+fn call_validate_block(
+	parent_head: Header,
+	block_data: ParachainBlockData<Block>,
+	relay_parent_storage_root: Hash,
+) -> cumulus_test_client::ExecutorResult<Header> {
+	call_validate_block_encoded_header(parent_head, block_data, relay_parent_storage_root)
+		.map(|v| Header::decode(&mut &v[..]).expect("Decodes `Header`."))
 }
 
 fn create_test_client() -> (Client, Header) {
@@ -126,6 +136,43 @@ fn validate_block_with_extra_extrinsics() {
 	assert_eq!(header, res_header);
 }
 
+#[test]
+fn validate_block_returns_custom_head_data() {
+	sp_tracing::try_init_simple();
+
+	let expected_header = vec![1, 3, 3, 7, 4, 5, 6];
+
+	let (client, parent_head) = create_test_client();
+	let extra_extrinsics = vec![
+		transfer(&client, Alice, Bob, 69),
+		generate_extrinsic(
+			&client,
+			Charlie,
+			TestPalletCall::set_custom_validation_head_data {
+				custom_header: expected_header.clone(),
+			},
+		),
+		transfer(&client, Bob, Charlie, 100),
+	];
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		extra_extrinsics,
+		parent_head.clone(),
+		Default::default(),
+	);
+	let header = block.header().clone();
+	assert_ne!(expected_header, header.encode());
+
+	let res_header = call_validate_block_encoded_header(
+		parent_head,
+		block,
+		validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block`");
+	assert_eq!(expected_header, res_header);
+}
+
 #[test]
 fn validate_block_invalid_parent_hash() {
 	sp_tracing::try_init_simple();

cumulus/test/client/src/lib.rs

@@ -122,7 +122,7 @@ fn genesis_config() -> GenesisConfig {
 pub fn generate_extrinsic(
 	client: &Client,
 	origin: sp_keyring::AccountKeyring,
-	function: Call,
+	function: impl Into<Call>,
 ) -> UncheckedExtrinsic {
 	let current_block_hash = client.info().best_hash;
 	let current_block = client.info().best_number.saturated_into();
@@ -139,6 +139,9 @@ pub fn generate_extrinsic(
 		frame_system::CheckWeight::<Runtime>::new(),
 		pallet_transaction_payment::ChargeTransactionPayment::<Runtime>::from(tip),
 	);
+
+	let function = function.into();
+
 	let raw_payload = SignedPayload::from_raw(
 		function.clone(),
 		extra.clone(),

cumulus/test/runtime/src/lib.rs

@@ -27,6 +27,8 @@ pub mod wasm_spec_version_incremented {
 	include!(concat!(env!("OUT_DIR"), "/wasm_binary_spec_version_incremented.rs"));
 }
 
+mod test_pallet;
+
 use frame_support::traits::OnRuntimeUpgrade;
 use sp_api::{decl_runtime_apis, impl_runtime_apis};
 use sp_core::OpaqueMetadata;
@@ -58,6 +60,7 @@ pub use pallet_timestamp::Call as TimestampCall;
 #[cfg(any(feature = "std", test))]
 pub use sp_runtime::BuildStorage;
 pub use sp_runtime::{Perbill, Permill};
+pub use test_pallet::Call as TestPalletCall;
 
 pub type SessionHandlers = ();
 
@@ -265,20 +268,21 @@ parameter_types! {
 	pub storage ParachainId: cumulus_primitives_core::ParaId = 100.into();
 }
 
+impl test_pallet::Config for Runtime {}
+
 construct_runtime! {
 	pub enum Runtime where
 		Block = Block,
 		NodeBlock = NodeBlock,
 		UncheckedExtrinsic = UncheckedExtrinsic,
 	{
-		System: frame_system::{Pallet, Call, Storage, Config, Event<T>},
-		ParachainSystem: cumulus_pallet_parachain_system::{
-			Pallet, Call, Config, Storage, Inherent, Event<T>, ValidateUnsigned,
-		},
-		Timestamp: pallet_timestamp::{Pallet, Call, Storage, Inherent},
-		Balances: pallet_balances::{Pallet, Call, Storage, Config<T>, Event<T>},
-		Sudo: pallet_sudo::{Pallet, Call, Storage, Config<T>, Event<T>},
-		TransactionPayment: pallet_transaction_payment::{Pallet, Storage},
+		System: frame_system,
+		ParachainSystem: cumulus_pallet_parachain_system,
+		Timestamp: pallet_timestamp,
+		Balances: pallet_balances,
+		Sudo: pallet_sudo,
+		TransactionPayment: pallet_transaction_payment,
+		TestPallet: test_pallet,
 	}
 }
 

cumulus/test/runtime/src/test_pallet.rs

@@ -0,0 +1,49 @@
+// 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/>.
+
+/// A special pallet that exposes dispatchables that are only useful for testing.
+pub use pallet::*;
+
+#[frame_support::pallet]
+pub mod pallet {
+	use frame_support::pallet_prelude::*;
+	use frame_system::pallet_prelude::*;
+
+	#[pallet::pallet]
+	#[pallet::generate_store(pub(super) trait Store)]
+	pub struct Pallet<T>(_);
+
+	#[pallet::config]
+	pub trait Config: frame_system::Config + cumulus_pallet_parachain_system::Config {}
+
+	#[pallet::hooks]
+	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {}
+
+	#[pallet::call]
+	impl<T: Config> Pallet<T> {
+		/// A test dispatchable for setting a custom head data in `validate_block`.
+		#[pallet::weight(0)]
+		pub fn set_custom_validation_head_data(
+			_: OriginFor<T>,
+			custom_header: sp_std::vec::Vec<u8>,
+		) -> DispatchResult {
+			cumulus_pallet_parachain_system::Pallet::<T>::set_custom_validation_head_data(
+				custom_header,
+			);
+			Ok(())
+		}
+	}
+}

cumulus/test/service/benches/transaction_throughput.rs

@@ -18,7 +18,7 @@
 
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
 use cumulus_test_runtime::{AccountId, BalancesCall, SudoCall};
-use futures::{future, join, StreamExt};
+use futures::{future, StreamExt};
 use polkadot_service::polkadot_runtime::constants::currency::DOLLARS;
 use sc_transaction_pool_api::{TransactionPool as _, TransactionSource, TransactionStatus};
 use sp_core::{crypto::Pair, sr25519};

cumulus/test/service/src/chain_spec.rs

@@ -128,5 +128,6 @@ fn testnet_genesis(
 			balances: endowed_accounts.iter().cloned().map(|k| (k, 1 << 60)).collect(),
 		},
 		sudo: cumulus_test_runtime::SudoConfig { key: root_key },
+		transaction_payment: Default::default(),
 	}
 }