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();