Refactoring Checkpoint: (WIP)

pezcumulus/pezpallets/teyrchain-system/src/validate_block/implementation.rs

@@ -0,0 +1,550 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! The actual implementation of the validate block functionality.
+
+use super::{trie_cache, trie_recorder, MemoryOptimizedValidationParams};
+use alloc::vec::Vec;
+use codec::{Decode, Encode};
+use pezcumulus_primitives_core::{
+	relay_chain::{BlockNumber as RNumber, Hash as RHash, UMPSignal, UMP_SEPARATOR},
+	ClaimQueueOffset, CoreSelector, PersistedValidationData, TeyrchainBlockData,
+};
+use pezframe_support::{
+	traits::{ExecuteBlock, Get, IsSubType},
+	BoundedVec,
+};
+use pezkuwi_teyrchain_primitives::primitives::{HeadData, ValidationResult};
+use pezsp_core::storage::{well_known_keys, ChildInfo, StateVersion};
+use pezsp_externalities::{set_and_run_with_externalities, Externalities};
+use pezsp_io::{hashing::blake2_128, KillStorageResult};
+use pezsp_runtime::traits::{
+	Block as BlockT, ExtrinsicCall, Hash as HashT, HashingFor, Header as HeaderT, LazyBlock,
+};
+use pezsp_state_machine::OverlayedChanges;
+use pezsp_trie::{HashDBT, ProofSizeProvider, EMPTY_PREFIX};
+use trie_recorder::{SeenNodes, SizeOnlyRecorderProvider};
+
+type Ext<'a, Block, Backend> = pezsp_state_machine::Ext<'a, HashingFor<Block>, Backend>;
+
+fn with_externalities<F: FnOnce(&mut dyn Externalities) -> R, R>(f: F) -> R {
+	pezsp_externalities::with_externalities(f).expect("Environmental externalities not set.")
+}
+
+// Recorder instance to be used during this validate_block call.
+environmental::environmental!(recorder: trait ProofSizeProvider);
+
+/// Validate the given teyrchain block.
+///
+/// This function is doing roughly the following:
+///
+/// 1. We decode the [`TeyrchainBlockData`] from the `block_data` in `params`.
+///
+/// 2. We are doing some security checks like checking that the `parent_head` in `params`
+/// is the parent of the block we are going to check. We also ensure that the `set_validation_data`
+/// inherent is present in the block and that the validation data matches the values in `params`.
+///
+/// 3. We construct the sparse in-memory database from the storage proof inside the block data and
+/// then ensure that the storage root matches the storage root in the `parent_head`.
+///
+/// 4. We replace all the storage related host functions with functions inside the wasm blob.
+/// This means instead of calling into the host, we will stay inside the wasm execution. This is
+/// very important as the relay chain validator hasn't the state required to verify the block. But
+/// we have the in-memory database that contains all the values from the state of the teyrchain
+/// that we require to verify the block.
+///
+/// 5. The last step is to execute the entire block in the machinery we just have setup. Executing
+/// the blocks include running all transactions in the block against our in-memory database and
+/// ensuring that the final storage root matches the storage root in the header of the block. In the
+/// end we return back the [`ValidationResult`] with all the required information for the validator.
+#[doc(hidden)]
+pub fn validate_block<B: BlockT, E: ExecuteBlock<B>, PSC: crate::Config>(
+	MemoryOptimizedValidationParams {
+		block_data,
+		parent_head: teyrchain_head,
+		relay_parent_number,
+		relay_parent_storage_root,
+	}: MemoryOptimizedValidationParams,
+) -> ValidationResult
+where
+	B::Extrinsic: ExtrinsicCall,
+	<B::Extrinsic as ExtrinsicCall>::Call: IsSubType<crate::Call<PSC>>,
+{
+	let _guard = (
+		// Replace storage calls with our own implementations
+		pezsp_io::storage::host_read.replace_implementation(host_storage_read),
+		pezsp_io::storage::host_set.replace_implementation(host_storage_set),
+		pezsp_io::storage::host_get.replace_implementation(host_storage_get),
+		pezsp_io::storage::host_exists.replace_implementation(host_storage_exists),
+		pezsp_io::storage::host_clear.replace_implementation(host_storage_clear),
+		pezsp_io::storage::host_root.replace_implementation(host_storage_root),
+		pezsp_io::storage::host_clear_prefix.replace_implementation(host_storage_clear_prefix),
+		pezsp_io::storage::host_append.replace_implementation(host_storage_append),
+		pezsp_io::storage::host_next_key.replace_implementation(host_storage_next_key),
+		pezsp_io::storage::host_start_transaction
+			.replace_implementation(host_storage_start_transaction),
+		pezsp_io::storage::host_rollback_transaction
+			.replace_implementation(host_storage_rollback_transaction),
+		pezsp_io::storage::host_commit_transaction
+			.replace_implementation(host_storage_commit_transaction),
+		pezsp_io::default_child_storage::host_get
+			.replace_implementation(host_default_child_storage_get),
+		pezsp_io::default_child_storage::host_read
+			.replace_implementation(host_default_child_storage_read),
+		pezsp_io::default_child_storage::host_set
+			.replace_implementation(host_default_child_storage_set),
+		pezsp_io::default_child_storage::host_clear
+			.replace_implementation(host_default_child_storage_clear),
+		pezsp_io::default_child_storage::host_storage_kill
+			.replace_implementation(host_default_child_storage_storage_kill),
+		pezsp_io::default_child_storage::host_exists
+			.replace_implementation(host_default_child_storage_exists),
+		pezsp_io::default_child_storage::host_clear_prefix
+			.replace_implementation(host_default_child_storage_clear_prefix),
+		pezsp_io::default_child_storage::host_root
+			.replace_implementation(host_default_child_storage_root),
+		pezsp_io::default_child_storage::host_next_key
+			.replace_implementation(host_default_child_storage_next_key),
+		pezsp_io::offchain_index::host_set.replace_implementation(host_offchain_index_set),
+		pezsp_io::offchain_index::host_clear.replace_implementation(host_offchain_index_clear),
+		pezcumulus_primitives_proof_size_hostfunction::storage_proof_size::host_storage_proof_size
+			.replace_implementation(host_storage_proof_size),
+	);
+
+	let block_data = codec::decode_from_bytes::<TeyrchainBlockData<B::LazyBlock>>(block_data)
+		.expect("Invalid teyrchain block data");
+
+	// Initialize hashmaps randomness.
+	pezsp_trie::add_extra_randomness(build_seed_from_head_data::<B>(
+		&block_data,
+		relay_parent_storage_root,
+	));
+
+	let mut parent_header =
+		codec::decode_from_bytes::<B::Header>(teyrchain_head.clone()).expect("Invalid parent head");
+
+	let (blocks, proof) = block_data.into_inner();
+
+	assert_eq!(
+		*blocks
+			.first()
+			.expect("BlockData should have at least one block")
+			.header()
+			.parent_hash(),
+		parent_header.hash(),
+		"Teyrchain head needs to be the parent of the first block"
+	);
+
+	blocks.iter().fold(parent_header.hash(), |p, b| {
+		assert_eq!(
+			p,
+			*b.header().parent_hash(),
+			"Not a valid chain of blocks :(; {:?} not a parent of {:?}?",
+			array_bytes::bytes2hex("0x", p.as_ref()),
+			array_bytes::bytes2hex("0x", b.header().parent_hash().as_ref()),
+		);
+		b.header().hash()
+	});
+
+	let mut processed_downward_messages = 0;
+	let mut upward_messages = BoundedVec::default();
+	let mut upward_message_signals = Vec::<Vec<_>>::new();
+	let mut horizontal_messages = BoundedVec::default();
+	let mut hrmp_watermark = Default::default();
+	let mut head_data = None;
+	let mut new_validation_code = None;
+	let num_blocks = blocks.len();
+
+	// Create the db
+	let mut db = match proof.to_memory_db(Some(parent_header.state_root())) {
+		Ok((db, _)) => db,
+		Err(_) => panic!("Compact proof decoding failure."),
+	};
+
+	core::mem::drop(proof);
+
+	let cache_provider = trie_cache::CacheProvider::new();
+	let seen_nodes = SeenNodes::<HashingFor<B>>::default();
+
+	for (block_index, mut block) in blocks.into_iter().enumerate() {
+		// We use the storage root of the `parent_head` to ensure that it is the correct root.
+		// This is already being done above while creating the in-memory db, but let's be paranoid!!
+		let backend = pezsp_state_machine::TrieBackendBuilder::new_with_cache(
+			&db,
+			*parent_header.state_root(),
+			&cache_provider,
+		)
+		.build();
+
+		// We use the same recorder when executing all blocks. So, each node only contributes once
+		// to the total size of the storage proof. This recorder should only be used for
+		// `execute_block`.
+		let mut execute_recorder = SizeOnlyRecorderProvider::with_seen_nodes(seen_nodes.clone());
+		// `backend` with the `execute_recorder`. As the `execute_recorder`, this should only be
+		// used for `execute_block`.
+		let execute_backend = pezsp_state_machine::TrieBackendBuilder::wrap(&backend)
+			.with_recorder(execute_recorder.clone())
+			.build();
+
+		// We let all blocks contribute to the same overlay. Data written by a previous block will
+		// be directly accessible without going to the db.
+		let mut overlay = OverlayedChanges::default();
+
+		parent_header = block.header().clone();
+
+		run_with_externalities_and_recorder::<B, _, _>(
+			&backend,
+			&mut Default::default(),
+			&mut Default::default(),
+			|| {
+				E::verify_and_remove_seal(&mut block);
+			},
+		);
+
+		run_with_externalities_and_recorder::<B, _, _>(
+			&execute_backend,
+			// Here is the only place where we want to use the recorder.
+			// We want to ensure that we not accidentally read something from the proof, that
+			// was not yet read and thus, alter the proof size. Otherwise, we end up with
+			// mismatches in later blocks.
+			&mut execute_recorder,
+			&mut overlay,
+			|| {
+				E::execute_verified_block(block);
+			},
+		);
+
+		if overlay.storage(well_known_keys::CODE).is_some() && num_blocks > 1 {
+			panic!("When applying a runtime upgrade, only one block per PoV is allowed. Received {num_blocks}.")
+		}
+
+		run_with_externalities_and_recorder::<B, _, _>(
+			&backend,
+			&mut Default::default(),
+			// We are only reading here, but need to know what the old block has written. Thus, we
+			// are passing here the overlay.
+			&mut overlay,
+			|| {
+				// Ensure the validation data is correct.
+				validate_validation_data(
+					crate::ValidationData::<PSC>::get()
+						.expect("`ValidationData` must be set after executing a block; qed"),
+					&teyrchain_head,
+					relay_parent_number,
+					relay_parent_storage_root,
+				);
+
+				new_validation_code =
+					new_validation_code.take().or(crate::NewValidationCode::<PSC>::get());
+
+				let mut found_separator = false;
+				crate::UpwardMessages::<PSC>::get()
+					.into_iter()
+					.filter_map(|m| {
+						// Filter out the `UMP_SEPARATOR` and the `UMPSignals`.
+						if m == UMP_SEPARATOR {
+							found_separator = true;
+							None
+						} else if found_separator {
+							if upward_message_signals.iter().all(|s| *s != m) {
+								upward_message_signals.push(m);
+							}
+							None
+						} else {
+							// No signal or separator
+							Some(m)
+						}
+					})
+					.for_each(|m| {
+						upward_messages.try_push(m)
+							.expect(
+								"Number of upward messages should not be greater than `MAX_UPWARD_MESSAGE_NUM`",
+							)
+					});
+
+				processed_downward_messages += crate::ProcessedDownwardMessages::<PSC>::get();
+				horizontal_messages.try_extend(crate::HrmpOutboundMessages::<PSC>::get().into_iter()).expect(
+					"Number of horizontal messages should not be greater than `MAX_HORIZONTAL_MESSAGE_NUM`",
+				);
+				hrmp_watermark = crate::HrmpWatermark::<PSC>::get();
+
+				if block_index + 1 == num_blocks {
+					head_data = Some(
+						crate::CustomValidationHeadData::<PSC>::get()
+							.map_or_else(|| HeadData(parent_header.encode()), HeadData),
+					);
+				}
+			},
+		);
+
+		if block_index + 1 != num_blocks {
+			let mut changes = overlay
+				.drain_storage_changes(
+					&backend,
+					<PSC as pezframe_system::Config>::Version::get().state_version(),
+				)
+				.expect("Failed to get drain storage changes from the overlay.");
+
+			drop(backend);
+
+			// We just forward the changes directly to our db.
+			changes.transaction.drain().into_iter().for_each(|(_, (value, count))| {
+				// We only care about inserts and not deletes.
+				if count > 0 {
+					db.insert(EMPTY_PREFIX, &value);
+
+					let hash = HashingFor::<B>::hash(&value);
+					seen_nodes.borrow_mut().insert(hash);
+				}
+			});
+		}
+	}
+
+	if !upward_message_signals.is_empty() {
+		let mut selected_core: Option<(CoreSelector, ClaimQueueOffset)> = None;
+		let mut approved_peer = None;
+
+		upward_message_signals.iter().for_each(|s| {
+			match UMPSignal::decode(&mut &s[..]).expect("Failed to decode `UMPSignal`") {
+				UMPSignal::SelectCore(selector, offset) => match &selected_core {
+					Some(selected_core) if *selected_core != (selector, offset) => {
+						panic!(
+							"All `SelectCore` signals need to select the same core: {selected_core:?} vs {:?}",
+							(selector, offset),
+						)
+					},
+					Some(_) => {},
+					None => {
+						selected_core = Some((selector, offset));
+					},
+				},
+				UMPSignal::ApprovedPeer(new_approved_peer) => match &approved_peer {
+					Some(approved_peer) if *approved_peer != new_approved_peer => {
+						panic!(
+							"All `ApprovedPeer` signals need to select the same peer_id: {new_approved_peer:?} vs {approved_peer:?}",
+						)
+					},
+					Some(_) => {},
+					None => {
+						approved_peer = Some(new_approved_peer);
+					},
+				},
+			}
+		});
+
+		upward_messages
+			.try_push(UMP_SEPARATOR)
+			.expect("UMPSignals does not fit in UMPMessages");
+		upward_messages
+			.try_extend(upward_message_signals.into_iter())
+			.expect("UMPSignals does not fit in UMPMessages");
+	}
+
+	ValidationResult {
+		head_data: head_data.expect("HeadData not set"),
+		new_validation_code: new_validation_code.map(Into::into),
+		upward_messages,
+		processed_downward_messages,
+		horizontal_messages,
+		hrmp_watermark,
+	}
+}
+
+/// Validates the given [`PersistedValidationData`] against the data from the relay chain.
+fn validate_validation_data(
+	validation_data: PersistedValidationData,
+	parent_header: &[u8],
+	relay_parent_number: RNumber,
+	relay_parent_storage_root: RHash,
+) {
+	assert_eq!(parent_header, &validation_data.parent_head.0, "Parent head doesn't match");
+	assert_eq!(
+		relay_parent_number, validation_data.relay_parent_number,
+		"Relay parent number doesn't match",
+	);
+	assert_eq!(
+		relay_parent_storage_root, validation_data.relay_parent_storage_root,
+		"Relay parent storage root doesn't match",
+	);
+}
+
+/// Build a seed from the head data of the teyrchain block.
+///
+/// Uses both the relay parent storage root and the hash of the blocks
+/// in the block data, to make sure the seed changes every block and that
+/// the user cannot find about it ahead of time.
+fn build_seed_from_head_data<B: BlockT>(
+	block_data: &TeyrchainBlockData<B::LazyBlock>,
+	relay_parent_storage_root: crate::relay_chain::Hash,
+) -> [u8; 16] {
+	let mut bytes_to_hash = Vec::with_capacity(
+		block_data.blocks().len() * size_of::<B::Hash>() + size_of::<crate::relay_chain::Hash>(),
+	);
+
+	bytes_to_hash.extend_from_slice(relay_parent_storage_root.as_ref());
+	block_data.blocks().iter().for_each(|block| {
+		bytes_to_hash.extend_from_slice(block.header().hash().as_ref());
+	});
+
+	blake2_128(&bytes_to_hash)
+}
+
+/// Run the given closure with the externalities and recorder set.
+fn run_with_externalities_and_recorder<Block: BlockT, R, F: FnOnce() -> R>(
+	backend: &impl pezsp_state_machine::Backend<HashingFor<Block>>,
+	recorder: &mut SizeOnlyRecorderProvider<HashingFor<Block>>,
+	overlay: &mut OverlayedChanges<HashingFor<Block>>,
+	execute: F,
+) -> R {
+	let mut ext = Ext::<Block, _>::new(overlay, backend);
+
+	recorder::using(recorder, || set_and_run_with_externalities(&mut ext, || execute()))
+}
+
+fn host_storage_read(key: &[u8], value_out: &mut [u8], value_offset: u32) -> Option<u32> {
+	match with_externalities(|ext| ext.storage(key)) {
+		Some(value) => {
+			let value_offset = value_offset as usize;
+			let data = &value[value_offset.min(value.len())..];
+			let written = core::cmp::min(data.len(), value_out.len());
+			value_out[..written].copy_from_slice(&data[..written]);
+			Some(value.len() as u32)
+		},
+		None => None,
+	}
+}
+
+fn host_storage_set(key: &[u8], value: &[u8]) {
+	with_externalities(|ext| ext.place_storage(key.to_vec(), Some(value.to_vec())))
+}
+
+fn host_storage_get(key: &[u8]) -> Option<bytes::Bytes> {
+	with_externalities(|ext| ext.storage(key).map(|value| value.into()))
+}
+
+fn host_storage_exists(key: &[u8]) -> bool {
+	with_externalities(|ext| ext.exists_storage(key))
+}
+
+fn host_storage_clear(key: &[u8]) {
+	with_externalities(|ext| ext.place_storage(key.to_vec(), None))
+}
+
+fn host_storage_proof_size() -> u64 {
+	recorder::with(|rec| rec.estimate_encoded_size()).expect("Recorder is always set; qed") as _
+}
+
+fn host_storage_root(version: StateVersion) -> Vec<u8> {
+	with_externalities(|ext| ext.storage_root(version))
+}
+
+fn host_storage_clear_prefix(prefix: &[u8], limit: Option<u32>) -> KillStorageResult {
+	with_externalities(|ext| ext.clear_prefix(prefix, limit, None).into())
+}
+
+fn host_storage_append(key: &[u8], value: Vec<u8>) {
+	with_externalities(|ext| ext.storage_append(key.to_vec(), value))
+}
+
+fn host_storage_next_key(key: &[u8]) -> Option<Vec<u8>> {
+	with_externalities(|ext| ext.next_storage_key(key))
+}
+
+fn host_storage_start_transaction() {
+	with_externalities(|ext| ext.storage_start_transaction())
+}
+
+fn host_storage_rollback_transaction() {
+	with_externalities(|ext| ext.storage_rollback_transaction().ok())
+		.expect("No open transaction that can be rolled back.");
+}
+
+fn host_storage_commit_transaction() {
+	with_externalities(|ext| ext.storage_commit_transaction().ok())
+		.expect("No open transaction that can be committed.");
+}
+
+fn host_default_child_storage_get(storage_key: &[u8], key: &[u8]) -> Option<Vec<u8>> {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.child_storage(&child_info, key))
+}
+
+fn host_default_child_storage_read(
+	storage_key: &[u8],
+	key: &[u8],
+	value_out: &mut [u8],
+	value_offset: u32,
+) -> Option<u32> {
+	let child_info = ChildInfo::new_default(storage_key);
+	match with_externalities(|ext| ext.child_storage(&child_info, key)) {
+		Some(value) => {
+			let value_offset = value_offset as usize;
+			let data = &value[value_offset.min(value.len())..];
+			let written = core::cmp::min(data.len(), value_out.len());
+			value_out[..written].copy_from_slice(&data[..written]);
+			Some(value.len() as u32)
+		},
+		None => None,
+	}
+}
+
+fn host_default_child_storage_set(storage_key: &[u8], key: &[u8], value: &[u8]) {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| {
+		ext.place_child_storage(&child_info, key.to_vec(), Some(value.to_vec()))
+	})
+}
+
+fn host_default_child_storage_clear(storage_key: &[u8], key: &[u8]) {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.place_child_storage(&child_info, key.to_vec(), None))
+}
+
+fn host_default_child_storage_storage_kill(
+	storage_key: &[u8],
+	limit: Option<u32>,
+) -> KillStorageResult {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.kill_child_storage(&child_info, limit, None).into())
+}
+
+fn host_default_child_storage_exists(storage_key: &[u8], key: &[u8]) -> bool {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.exists_child_storage(&child_info, key))
+}
+
+fn host_default_child_storage_clear_prefix(
+	storage_key: &[u8],
+	prefix: &[u8],
+	limit: Option<u32>,
+) -> KillStorageResult {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.clear_child_prefix(&child_info, prefix, limit, None).into())
+}
+
+fn host_default_child_storage_root(storage_key: &[u8], version: StateVersion) -> Vec<u8> {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.child_storage_root(&child_info, version))
+}
+
+fn host_default_child_storage_next_key(storage_key: &[u8], key: &[u8]) -> Option<Vec<u8>> {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.next_child_storage_key(&child_info, key))
+}
+
+fn host_offchain_index_set(_key: &[u8], _value: &[u8]) {}
+
+fn host_offchain_index_clear(_key: &[u8]) {}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/mod.rs

@@ -0,0 +1,70 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! A module that enables a runtime to work as teyrchain.
+
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub mod implementation;
+#[cfg(test)]
+mod tests;
+
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub mod trie_cache;
+
+#[cfg(any(test, not(feature = "std")))]
+#[doc(hidden)]
+pub mod trie_recorder;
+
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use alloc::{boxed::Box, slice};
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use bytes;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use codec::decode_from_bytes;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use pezkuwi_teyrchain_primitives;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use pezsp_runtime::traits::GetRuntimeBlockType;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use pezsp_std;
+
+/// Basically the same as
+/// [`ValidationParams`](pezkuwi_teyrchain_primitives::primitives::ValidationParams), but a little
+/// bit optimized for our use case here.
+///
+/// `block_data` and `head_data` are represented as [`bytes::Bytes`] to make them reuse
+/// the memory of the input parameter of the exported `validate_blocks` function.
+///
+/// The layout of this type must match exactly the layout of
+/// [`ValidationParams`](pezkuwi_teyrchain_primitives::primitives::ValidationParams) to have the
+/// same SCALE encoding.
+#[derive(codec::Decode)]
+#[cfg_attr(feature = "std", derive(codec::Encode))]
+#[doc(hidden)]
+pub struct MemoryOptimizedValidationParams {
+	pub parent_head: bytes::Bytes,
+	pub block_data: bytes::Bytes,
+	pub relay_parent_number: pezcumulus_primitives_core::relay_chain::BlockNumber,
+	pub relay_parent_storage_root: pezcumulus_primitives_core::relay_chain::Hash,
+}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/tests.rs

@@ -0,0 +1,777 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::{validate_block::MemoryOptimizedValidationParams, *};
+use codec::{Decode, DecodeAll, Encode};
+use pezcumulus_primitives_core::{relay_chain, PersistedValidationData, TeyrchainBlockData};
+use pezcumulus_test_client::{
+	generate_extrinsic, generate_extrinsic_with_pair,
+	runtime::{
+		self as test_runtime, Block, Hash, Header, SudoCall, SystemCall, TestPalletCall,
+		UncheckedExtrinsic, WASM_BINARY,
+	},
+	seal_block, transfer, BlockData, BlockOrigin, BuildTeyrchainBlockData, Client,
+	DefaultTestClientBuilderExt, HeadData, InitBlockBuilder,
+	Sr25519Keyring::{Alice, Bob, Charlie},
+	TestClientBuilder, TestClientBuilderExt, ValidationParams,
+};
+use pezcumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
+use pezkuwi_teyrchain_primitives::primitives::ValidationResult;
+use pezsc_consensus::{BlockImport, BlockImportParams, ForkChoiceStrategy};
+use pezsp_api::{ApiExt, Core, ProofRecorder, ProvideRuntimeApi};
+use pezsp_consensus_slots::SlotDuration;
+use pezsp_core::{Hasher, H256};
+use pezsp_runtime::{
+	traits::{BlakeTwo256, Block as BlockT, Header as HeaderT},
+	DigestItem,
+};
+use pezsp_trie::{proof_size_extension::ProofSizeExt, recorder::IgnoredNodes, StorageProof};
+use std::{env, process::Command};
+
+fn call_validate_block_validation_result(
+	validation_code: &[u8],
+	parent_head: Header,
+	block_data: TeyrchainBlockData<Block>,
+	relay_parent_storage_root: Hash,
+) -> pezcumulus_test_client::ExecutorResult<ValidationResult> {
+	pezcumulus_test_client::validate_block(
+		ValidationParams {
+			block_data: BlockData(block_data.encode()),
+			parent_head: HeadData(parent_head.encode()),
+			relay_parent_number: 1,
+			relay_parent_storage_root,
+		},
+		validation_code,
+	)
+}
+
+fn call_validate_block(
+	parent_head: Header,
+	block_data: TeyrchainBlockData<Block>,
+	relay_parent_storage_root: Hash,
+) -> pezcumulus_test_client::ExecutorResult<Header> {
+	call_validate_block_validation_result(
+		WASM_BINARY.expect("You need to build the WASM binaries to run the tests!"),
+		parent_head,
+		block_data,
+		relay_parent_storage_root,
+	)
+	.map(|v| Header::decode(&mut &v.head_data.0[..]).expect("Decodes `Header`."))
+}
+
+/// Call `validate_block` in the runtime with `elastic-scaling` activated.
+fn call_validate_block_elastic_scaling(
+	parent_head: Header,
+	block_data: TeyrchainBlockData<Block>,
+	relay_parent_storage_root: Hash,
+) -> pezcumulus_test_client::ExecutorResult<Header> {
+	call_validate_block_validation_result(
+		test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!"),
+		parent_head,
+		block_data,
+		relay_parent_storage_root,
+	)
+	.map(|v| Header::decode(&mut &v.head_data.0[..]).expect("Decodes `Header`."))
+}
+
+fn create_test_client() -> (Client, Header) {
+	let client = TestClientBuilder::new().enable_import_proof_recording().build();
+
+	let genesis_header = client
+		.header(client.chain_info().genesis_hash)
+		.ok()
+		.flatten()
+		.expect("Genesis header exists; qed");
+
+	(client, genesis_header)
+}
+
+/// Create test client using the runtime with `elastic-scaling` feature enabled.
+fn create_elastic_scaling_test_client() -> (Client, Header) {
+	let mut builder = TestClientBuilder::new();
+	builder.genesis_init_mut().wasm = Some(
+		test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!")
+			.to_vec(),
+	);
+	let client = builder.enable_import_proof_recording().build();
+
+	let genesis_header = client
+		.header(client.chain_info().genesis_hash)
+		.ok()
+		.flatten()
+		.expect("Genesis header exists; qed");
+
+	(client, genesis_header)
+}
+
+fn pop_seal(mut block: Block) -> Block {
+	assert!(block.header.digest.pop().unwrap().as_seal().is_some());
+	block
+}
+
+struct TestBlockData {
+	block: TeyrchainBlockData<Block>,
+	validation_data: PersistedValidationData,
+}
+
+fn build_block_with_witness(
+	client: &Client,
+	extra_extrinsics: Vec<UncheckedExtrinsic>,
+	parent_head: Header,
+	mut sproof_builder: RelayStateSproofBuilder,
+	pre_digests: Vec<DigestItem>,
+) -> TestBlockData {
+	sproof_builder.para_id = test_runtime::TEYRCHAIN_ID.into();
+	sproof_builder.included_para_head = Some(HeadData(parent_head.encode()));
+
+	let validation_data = PersistedValidationData {
+		relay_parent_number: 1,
+		parent_head: parent_head.encode().into(),
+		..Default::default()
+	};
+
+	let pezcumulus_test_client::BlockBuilderAndSupportData {
+		mut block_builder,
+		persisted_validation_data,
+	} = client.init_block_builder_with_pre_digests(Some(validation_data), sproof_builder, pre_digests);
+
+	extra_extrinsics.into_iter().for_each(|e| block_builder.push(e).unwrap());
+
+	let mut block = block_builder.build_teyrchain_block(*parent_head.state_root());
+
+	block.blocks_mut()[0] = seal_block(block.blocks()[0].clone(), client);
+
+	TestBlockData { block, validation_data: persisted_validation_data }
+}
+
+fn build_multiple_blocks_with_witness(
+	client: &Client,
+	mut parent_head: Header,
+	mut sproof_builder: RelayStateSproofBuilder,
+	num_blocks: u32,
+	extra_extrinsics: impl Fn(u32) -> Vec<UncheckedExtrinsic>,
+) -> TestBlockData {
+	let parent_head_root = *parent_head.state_root();
+	sproof_builder.para_id = test_runtime::TEYRCHAIN_ID.into();
+	sproof_builder.included_para_head = Some(HeadData(parent_head.encode()));
+
+	let timestamp = if sproof_builder.current_slot == 0u64 {
+		let timestamp = std::time::SystemTime::now()
+			.duration_since(std::time::SystemTime::UNIX_EPOCH)
+			.expect("Time is always after UNIX_EPOCH; qed")
+			.as_millis() as u64;
+		sproof_builder.current_slot = (timestamp / 6000).into();
+
+		timestamp
+	} else {
+		sproof_builder
+			.current_slot
+			.timestamp(SlotDuration::from_millis(6000))
+			.unwrap()
+			.as_millis()
+	};
+
+	let validation_data = PersistedValidationData {
+		relay_parent_number: 1,
+		parent_head: parent_head.encode().into(),
+		..Default::default()
+	};
+
+	let mut persisted_validation_data = None;
+	let mut blocks = Vec::new();
+	let mut proof = StorageProof::empty();
+	let mut ignored_nodes = IgnoredNodes::<H256>::default();
+
+	for i in 0..num_blocks {
+		let pezcumulus_test_client::BlockBuilderAndSupportData {
+			mut block_builder,
+			persisted_validation_data: p_v_data,
+		} = client.init_block_builder_with_ignored_nodes(
+			parent_head.hash(),
+			Some(validation_data.clone()),
+			sproof_builder.clone(),
+			timestamp,
+			ignored_nodes.clone(),
+		);
+
+		persisted_validation_data = Some(p_v_data);
+
+		for ext in (extra_extrinsics)(i) {
+			block_builder.push(ext).unwrap();
+		}
+
+		let mut built_block = block_builder.build().unwrap();
+		built_block.block = seal_block(built_block.block, &client);
+
+		futures::executor::block_on({
+			let parent_hash = *built_block.block.header.parent_hash();
+			let state = client.state_at(parent_hash).unwrap();
+
+			let mut api = client.runtime_api();
+			let proof_recorder = ProofRecorder::<Block>::with_ignored_nodes(ignored_nodes.clone());
+			api.record_proof_with_recorder(proof_recorder.clone());
+			api.register_extension(ProofSizeExt::new(proof_recorder));
+			api.execute_block(parent_hash, pop_seal(built_block.block.clone()).into())
+				.unwrap();
+
+			let (mut header, extrinsics) = built_block.block.clone().deconstruct();
+			let seal = header.digest.pop().unwrap();
+
+			let mut import = BlockImportParams::new(BlockOrigin::Own, header);
+			import.body = Some(extrinsics);
+			import.post_digests.push(seal);
+			import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+			import.state_action = api.into_storage_changes(&state, parent_hash).unwrap().into();
+
+			BlockImport::import_block(&client, import)
+		})
+		.unwrap();
+
+		ignored_nodes.extend(IgnoredNodes::from_storage_proof::<BlakeTwo256>(
+			&built_block.proof.clone().unwrap(),
+		));
+		ignored_nodes.extend(IgnoredNodes::from_memory_db(built_block.storage_changes.transaction));
+		proof = StorageProof::merge([proof, built_block.proof.unwrap()]);
+
+		parent_head = built_block.block.header.clone();
+
+		blocks.push(built_block.block);
+	}
+
+	let proof = proof.into_compact_proof::<BlakeTwo256>(parent_head_root).unwrap();
+
+	TestBlockData {
+		block: TeyrchainBlockData::new(blocks, proof),
+		validation_data: persisted_validation_data.unwrap(),
+	}
+}
+
+#[test]
+fn validate_block_works() {
+	pezsp_tracing::try_init_simple();
+
+	let (client, parent_head) = create_test_client();
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		Vec::new(),
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+
+	let header = block.blocks()[0].header().clone();
+	let res_header =
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}
+
+#[test]
+fn validate_multiple_blocks_work() {
+	pezsp_tracing::try_init_simple();
+
+	let blocks_per_pov = 4;
+	let (client, parent_head) = create_elastic_scaling_test_client();
+	let TestBlockData { block, validation_data } = build_multiple_blocks_with_witness(
+		&client,
+		parent_head.clone(),
+		Default::default(),
+		blocks_per_pov,
+		|i| {
+			vec![generate_extrinsic_with_pair(
+				&client,
+				Charlie.into(),
+				TestPalletCall::read_and_write_big_value {},
+				Some(i),
+			)]
+		},
+	);
+
+	assert!(block.proof().encoded_size() < 3 * 1024 * 1024);
+
+	let header = block.blocks().last().unwrap().header().clone();
+	let res_header = call_validate_block_elastic_scaling(
+		parent_head,
+		block,
+		validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}
+
+#[test]
+fn validate_block_with_extra_extrinsics() {
+	pezsp_tracing::try_init_simple();
+
+	let (client, parent_head) = create_test_client();
+	let extra_extrinsics = vec![
+		transfer(&client, Alice, Bob, 69),
+		transfer(&client, Bob, Charlie, 100),
+		transfer(&client, Charlie, Alice, 500),
+	];
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		extra_extrinsics,
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+	let header = block.blocks()[0].header().clone();
+
+	let res_header =
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}
+
+#[test]
+fn validate_block_returns_custom_head_data() {
+	pezsp_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(),
+		Default::default(),
+	);
+	let header = block.blocks()[0].header().clone();
+	assert_ne!(expected_header, header.encode());
+
+	let res_header = call_validate_block_validation_result(
+		WASM_BINARY.expect("You need to build the WASM binaries to run the tests!"),
+		parent_head,
+		block,
+		validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block`")
+	.head_data
+	.0;
+	assert_eq!(expected_header, res_header);
+}
+
+#[test]
+fn validate_block_rejects_invalid_seal() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, parent_head) = create_test_client();
+		let TestBlockData { mut block, validation_data, .. } = build_block_with_witness(
+			&client,
+			Vec::new(),
+			parent_head.clone(),
+			Default::default(),
+			Default::default(),
+		);
+		let (id, data) =
+			block.blocks_mut()[0].header.digest.logs.last().unwrap().as_seal().unwrap();
+		let mut data = data.to_vec();
+		let random = BlakeTwo256::hash(&data);
+		data[..random.as_ref().len()].copy_from_slice(random.as_ref());
+
+		*block.blocks_mut()[0].header.digest.logs.last_mut().unwrap() = DigestItem::Seal(id, data);
+
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args(["validate_block_rejects_invalid_seal", "--", "--nocapture"])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap()).contains("Invalid AuRa seal"));
+	}
+}
+#[test]
+fn validate_block_invalid_parent_hash() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, parent_head) = create_test_client();
+		let TestBlockData { mut block, validation_data, .. } = build_block_with_witness(
+			&client,
+			Vec::new(),
+			parent_head.clone(),
+			Default::default(),
+			Default::default(),
+		);
+		block.blocks_mut()[0].header.set_parent_hash(Hash::from_low_u64_be(1));
+
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args(["validate_block_invalid_parent_hash", "--", "--nocapture"])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("Teyrchain head needs to be the parent of the first block"));
+	}
+}
+
+#[test]
+fn validate_block_fails_on_invalid_validation_data() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, parent_head) = create_test_client();
+		let TestBlockData { block, .. } = build_block_with_witness(
+			&client,
+			Vec::new(),
+			parent_head.clone(),
+			Default::default(),
+			Default::default(),
+		);
+
+		call_validate_block(parent_head, block, Hash::random()).unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args(["validate_block_fails_on_invalid_validation_data", "--", "--nocapture"])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("Relay parent storage root doesn't match"));
+	}
+}
+
+/// Test that ensures that `ValidationParams` and `MemoryOptimizedValidationParams`
+/// are encoding/decoding.
+#[test]
+fn validation_params_and_memory_optimized_validation_params_encode_and_decode() {
+	const BLOCK_DATA: &[u8] = &[1, 2, 3, 4, 5];
+	const PARENT_HEAD: &[u8] = &[1, 3, 4, 5, 6, 7, 9];
+
+	let validation_params = ValidationParams {
+		block_data: BlockData(BLOCK_DATA.encode()),
+		parent_head: HeadData(PARENT_HEAD.encode()),
+		relay_parent_number: 1,
+		relay_parent_storage_root: Hash::random(),
+	};
+
+	let encoded = validation_params.encode();
+
+	let decoded = MemoryOptimizedValidationParams::decode_all(&mut &encoded[..]).unwrap();
+	assert_eq!(decoded.relay_parent_number, validation_params.relay_parent_number);
+	assert_eq!(decoded.relay_parent_storage_root, validation_params.relay_parent_storage_root);
+	assert_eq!(decoded.block_data, validation_params.block_data.0);
+	assert_eq!(decoded.parent_head, validation_params.parent_head.0);
+
+	let encoded = decoded.encode();
+
+	let decoded = ValidationParams::decode_all(&mut &encoded[..]).unwrap();
+	assert_eq!(decoded, validation_params);
+}
+
+/// Test for ensuring that we are differentiating in the `validation::trie_cache` between different
+/// child tries.
+///
+/// This is achieved by first building a block using `read_and_write_child_tries` that should set
+/// the values in the child tries. In the second step we are building a second block with the same
+/// extrinsic that reads the values from the child tries and it asserts that we read the correct
+/// data from the state.
+#[test]
+fn validate_block_works_with_child_tries() {
+	pezsp_tracing::try_init_simple();
+
+	let (client, parent_head) = create_test_client();
+	let TestBlockData { block, .. } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic(&client, Charlie, TestPalletCall::read_and_write_child_tries {})],
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+
+	let (mut header, extrinsics) = block.blocks()[0].clone().deconstruct();
+	let seal = header.digest.pop().unwrap();
+
+	let mut import = BlockImportParams::new(BlockOrigin::Own, header.clone());
+	import.body = Some(extrinsics);
+	import.post_digests.push(seal);
+	import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+
+	futures::executor::block_on(BlockImport::import_block(&client, import)).unwrap();
+
+	let parent_head = block.blocks()[0].header.clone();
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic(&client, Alice, TestPalletCall::read_and_write_child_tries {})],
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+
+	let header = block.blocks()[0].header().clone();
+	let res_header =
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}
+
+#[test]
+fn state_changes_in_multiple_blocks_are_applied_in_exact_order() {
+	pezsp_tracing::try_init_simple();
+
+	let blocks_per_pov = 12;
+	// disable the core selection logic
+	let (client, genesis_head) = create_elastic_scaling_test_client();
+
+	// 1. Build the initial block that stores values in the map.
+	let TestBlockData { block: initial_block_data, .. } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic_with_pair(
+			&client,
+			Alice.into(),
+			TestPalletCall::store_values_in_map { max_key: 4095 },
+			Some(0),
+		)],
+		genesis_head.clone(),
+		RelayStateSproofBuilder { current_slot: 1.into(), ..Default::default() },
+		Vec::new(),
+	);
+
+	let initial_block = initial_block_data.blocks()[0].clone();
+	let (mut header, extrinsics) = initial_block.clone().deconstruct();
+	let seal = header.digest.pop().unwrap();
+
+	let mut import = BlockImportParams::new(BlockOrigin::Own, header.clone());
+	import.body = Some(extrinsics);
+	import.post_digests.push(seal);
+	import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+
+	futures::executor::block_on(BlockImport::import_block(&client, import)).unwrap();
+	let initial_block_header = initial_block.header().clone();
+
+	// 2. Build the PoV block that removes values from the map.
+	let TestBlockData { block: pov_block_data, validation_data: pov_validation_data } =
+		build_multiple_blocks_with_witness(
+			&client,
+			initial_block_header.clone(), // Start building PoV from the initial block's header
+			RelayStateSproofBuilder { current_slot: 2.into(), ..Default::default() },
+			blocks_per_pov,
+			|i| {
+				// Each block `i` (0-11) removes key `116 + i`.
+				let key_to_remove = 116 + i;
+				vec![generate_extrinsic_with_pair(
+					&client,
+					Bob.into(), // Use Bob to avoid nonce conflicts with Alice
+					TestPalletCall::remove_value_from_map { key: key_to_remove },
+					Some(i),
+				)]
+			},
+		);
+
+	// 3. Validate the PoV.
+	let final_pov_header = pov_block_data.blocks().last().unwrap().header().clone();
+	let res_header = call_validate_block_elastic_scaling(
+		initial_block_header, // The parent is the head of the initial block before the PoV
+		pov_block_data,
+		pov_validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block` after building the PoV");
+	assert_eq!(final_pov_header, res_header);
+}
+
+#[test]
+fn validate_block_handles_ump_signal() {
+	use pezcumulus_primitives_core::{
+		relay_chain::{UMPSignal, UMP_SEPARATOR},
+		ClaimQueueOffset, CoreInfo, CoreSelector,
+	};
+
+	pezsp_tracing::try_init_simple();
+
+	let (client, parent_head) = create_elastic_scaling_test_client();
+	let extra_extrinsics =
+		vec![transfer(&client, Alice, Bob, 69), transfer(&client, Bob, Charlie, 100)];
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		extra_extrinsics,
+		parent_head.clone(),
+		Default::default(),
+		vec![CumulusDigestItem::CoreInfo(CoreInfo {
+			selector: CoreSelector(0),
+			claim_queue_offset: ClaimQueueOffset(0),
+			number_of_cores: 1.into(),
+		})
+		.to_digest_item()],
+	);
+
+	let upward_messages = call_validate_block_validation_result(
+		test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!"),
+		parent_head,
+		block,
+		validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block`")
+	.upward_messages;
+
+	assert_eq!(
+		upward_messages,
+		vec![UMP_SEPARATOR, UMPSignal::SelectCore(CoreSelector(0), ClaimQueueOffset(0)).encode()]
+	);
+}
+
+#[test]
+fn ensure_we_only_like_blockchains() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, parent_head) = create_elastic_scaling_test_client();
+		let TestBlockData { mut block, validation_data } = build_multiple_blocks_with_witness(
+			&client,
+			parent_head.clone(),
+			Default::default(),
+			4,
+			|_| Default::default(),
+		);
+
+		// Reference some non existing parent.
+		block.blocks_mut()[2].header.parent_hash = Hash::default();
+
+		call_validate_block_elastic_scaling(
+			parent_head,
+			block,
+			validation_data.relay_parent_storage_root,
+		)
+		.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args(["ensure_we_only_like_blockchains", "--", "--nocapture"])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("Not a valid chain of blocks :("));
+	}
+}
+
+#[test]
+fn rejects_multiple_blocks_per_pov_when_applying_runtime_upgrade() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, genesis_head) = create_elastic_scaling_test_client();
+
+		let code = test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!")
+			.to_vec();
+		let code_len = code.len() as u32;
+
+		let mut proof_builder =
+			RelayStateSproofBuilder { current_slot: 1.into(), ..Default::default() };
+		proof_builder.host_config.max_code_size = code_len * 2;
+
+		// Build the block that send the runtime upgrade.
+		let TestBlockData { block: initial_block_data, .. } = build_block_with_witness(
+			&client,
+			vec![generate_extrinsic_with_pair(
+				&client,
+				Alice.into(),
+				SudoCall::sudo {
+					call: Box::new(SystemCall::set_code_without_checks { code }.into()),
+				},
+				Some(0),
+			)],
+			genesis_head.clone(),
+			proof_builder,
+			Vec::new(),
+		);
+
+		let initial_block = initial_block_data.blocks()[0].clone();
+		let (mut header, extrinsics) = initial_block.clone().deconstruct();
+		let seal = header.digest.pop().unwrap();
+
+		let mut import = BlockImportParams::new(BlockOrigin::Own, header.clone());
+		import.body = Some(extrinsics);
+		import.post_digests.push(seal);
+		import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+
+		futures::executor::block_on(BlockImport::import_block(&client, import)).unwrap();
+		let initial_block_header = initial_block.header().clone();
+
+		let mut proof_builder = RelayStateSproofBuilder {
+			current_slot: 2.into(),
+			upgrade_go_ahead: Some(relay_chain::UpgradeGoAhead::GoAhead),
+			..Default::default()
+		};
+		proof_builder.host_config.max_code_size = code_len * 2;
+
+		// 2. Build a PoV that consists of multiple blocks.
+		let TestBlockData { block: pov_block_data, validation_data: pov_validation_data } =
+			build_multiple_blocks_with_witness(
+				&client,
+				initial_block_header.clone(), // Start building PoV from the initial block's header
+				proof_builder,
+				4,
+				|_| Vec::new(),
+			);
+
+		// 3. Validate the PoV.
+		call_validate_block_elastic_scaling(
+			initial_block_header, // The parent is the head of the initial block before the PoV
+			pov_block_data,
+			pov_validation_data.relay_parent_storage_root,
+		)
+		.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args([
+				"rejects_multiple_blocks_per_pov_when_applying_runtime_upgrade",
+				"--",
+				"--nocapture",
+			])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("only one block per PoV is allowed"));
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/trie_cache.rs

@@ -0,0 +1,158 @@
+// This file is part of Pezcumulus.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use alloc::boxed::Box;
+
+use core::cell::{RefCell, RefMut};
+use hashbrown::{hash_map::Entry, HashMap};
+use pezsp_state_machine::TrieCacheProvider;
+use pezsp_trie::{NodeCodec, RandomState};
+use trie_db::{node::NodeOwned, Hasher};
+
+/// Special purpose trie cache implementation that is able to cache an unlimited number
+/// of values. To be used in `validate_block` to serve values and nodes that
+/// have already been loaded and decoded from the storage proof.
+pub struct TrieCache<'a, H: Hasher> {
+	node_cache: RefMut<'a, HashMap<H::Out, NodeOwned<H::Out>, RandomState>>,
+	value_cache: Option<RefMut<'a, HashMap<Box<[u8]>, trie_db::CachedValue<H::Out>, RandomState>>>,
+}
+
+impl<'a, H: Hasher> trie_db::TrieCache<NodeCodec<H>> for TrieCache<'a, H> {
+	fn lookup_value_for_key(&mut self, key: &[u8]) -> Option<&trie_db::CachedValue<H::Out>> {
+		self.value_cache.as_ref().and_then(|cache| cache.get(key))
+	}
+
+	fn cache_value_for_key(&mut self, key: &[u8], value: trie_db::CachedValue<H::Out>) {
+		self.value_cache.as_mut().and_then(|cache| cache.insert(key.into(), value));
+	}
+
+	fn get_or_insert_node(
+		&mut self,
+		hash: <NodeCodec<H> as trie_db::NodeCodec>::HashOut,
+		fetch_node: &mut dyn FnMut() -> trie_db::Result<
+			NodeOwned<H::Out>,
+			H::Out,
+			<NodeCodec<H> as trie_db::NodeCodec>::Error,
+		>,
+	) -> trie_db::Result<&NodeOwned<H::Out>, H::Out, <NodeCodec<H> as trie_db::NodeCodec>::Error> {
+		match self.node_cache.entry(hash) {
+			Entry::Occupied(entry) => Ok(entry.into_mut()),
+			Entry::Vacant(entry) => Ok(entry.insert(fetch_node()?)),
+		}
+	}
+
+	fn get_node(
+		&mut self,
+		hash: &H::Out,
+	) -> Option<&NodeOwned<<NodeCodec<H> as trie_db::NodeCodec>::HashOut>> {
+		self.node_cache.get(hash)
+	}
+}
+
+/// Provider of [`TrieCache`] instances.
+pub struct CacheProvider<H: Hasher> {
+	node_cache: RefCell<HashMap<H::Out, NodeOwned<H::Out>, RandomState>>,
+	/// Cache: `storage_root` => `storage_key` => `value`.
+	///
+	/// One `block` can for example use multiple tries (child tries) and we need to distinguish the
+	/// cached (`storage_key`, `value`) between them. For this we are using the `storage_root` to
+	/// distinguish them (even if the storage root is the same for two child tries, it just means
+	/// that both are exactly the same trie and there would happen no collision).
+	value_cache: RefCell<
+		HashMap<H::Out, HashMap<Box<[u8]>, trie_db::CachedValue<H::Out>, RandomState>, RandomState>,
+	>,
+}
+
+impl<H: Hasher> CacheProvider<H> {
+	/// Constructs a new instance of [`CacheProvider`] with an uninitialized state
+	/// and empty node and value caches.
+	pub fn new() -> Self {
+		CacheProvider { node_cache: Default::default(), value_cache: Default::default() }
+	}
+}
+
+impl<H: Hasher> TrieCacheProvider<H> for &&CacheProvider<H> {
+	type Cache<'a>
+		= TrieCache<'a, H>
+	where
+		Self: 'a,
+		H: 'a;
+
+	fn as_trie_db_cache(&self, storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+		TrieCacheProvider::<H>::as_trie_db_cache(**self, storage_root)
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		TrieCacheProvider::<H>::as_trie_db_mut_cache(**self)
+	}
+
+	fn merge<'a>(&'a self, other: Self::Cache<'a>, new_root: <H as Hasher>::Out) {
+		TrieCacheProvider::merge(**self, other, new_root)
+	}
+}
+
+impl<H: Hasher> TrieCacheProvider<H> for &CacheProvider<H> {
+	type Cache<'a>
+		= TrieCache<'a, H>
+	where
+		Self: 'a,
+		H: 'a;
+
+	fn as_trie_db_cache(&self, storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+		TrieCacheProvider::<H>::as_trie_db_cache(*self, storage_root)
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		TrieCacheProvider::<H>::as_trie_db_mut_cache(*self)
+	}
+
+	fn merge<'a>(&'a self, other: Self::Cache<'a>, new_root: <H as Hasher>::Out) {
+		TrieCacheProvider::merge(*self, other, new_root)
+	}
+}
+
+impl<H: Hasher> TrieCacheProvider<H> for CacheProvider<H> {
+	type Cache<'a>
+		= TrieCache<'a, H>
+	where
+		H: 'a;
+
+	fn as_trie_db_cache(&self, storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+		TrieCache {
+			value_cache: Some(RefMut::map(self.value_cache.borrow_mut(), |c| {
+				c.entry(storage_root).or_default()
+			})),
+			node_cache: self.node_cache.borrow_mut(),
+		}
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		// This method is called when we calculate the storage root.
+		// We are not interested in caching new values (as we would throw them away directly after a
+		// block is validated) and thus, we don't pass any `value_cache`.
+		TrieCache { value_cache: None, node_cache: self.node_cache.borrow_mut() }
+	}
+
+	fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: <H as Hasher>::Out) {
+		// This is called to merge the `value_cache` from `as_trie_db_mut_cache`, which is not
+		// activated, so we don't need to do anything here.
+	}
+}
+
+// This is safe here since we are single-threaded in WASM
+unsafe impl<H: Hasher> Send for CacheProvider<H> {}
+unsafe impl<H: Hasher> Sync for CacheProvider<H> {}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/trie_recorder.rs

@@ -0,0 +1,296 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Provide a specialized trie-recorder and provider for use in validate-block.
+//!
+//! This file defines two main structs, [`SizeOnlyRecorder`] and
+//! [`SizeOnlyRecorderProvider`]. They are used to track the current
+//! proof-size without actually recording the accessed nodes themselves.
+
+use alloc::rc::Rc;
+use codec::Encode;
+use core::cell::{RefCell, RefMut};
+use hashbrown::{HashMap, HashSet};
+use pezsp_trie::{NodeCodec, ProofSizeProvider, RandomState, StorageProof};
+use trie_db::{Hasher, RecordedForKey, TrieAccess};
+
+pub(crate) type SeenNodes<H> = Rc<RefCell<HashSet<<H as Hasher>::Out, RandomState>>>;
+
+/// A trie recorder that only keeps track of the proof size.
+///
+/// The internal size counting logic should align
+/// with ['pezsp_trie::recorder::Recorder'].
+pub struct SizeOnlyRecorder<'a, H: Hasher> {
+	seen_nodes: RefMut<'a, HashSet<H::Out, RandomState>>,
+	encoded_size: RefMut<'a, usize>,
+	recorded_keys: RefMut<'a, HashMap<Rc<[u8]>, RecordedForKey, RandomState>>,
+}
+
+impl<'a, H: trie_db::Hasher> trie_db::TrieRecorder<H::Out> for SizeOnlyRecorder<'a, H> {
+	fn record(&mut self, access: TrieAccess<'_, H::Out>) {
+		let mut encoded_size_update = 0;
+		match access {
+			TrieAccess::NodeOwned { hash, node_owned } =>
+				if self.seen_nodes.insert(hash) {
+					let node = node_owned.to_encoded::<NodeCodec<H>>();
+					encoded_size_update += node.encoded_size();
+				},
+			TrieAccess::EncodedNode { hash, encoded_node } =>
+				if self.seen_nodes.insert(hash) {
+					encoded_size_update += encoded_node.encoded_size();
+				},
+			TrieAccess::Value { hash, value, full_key } => {
+				if self.seen_nodes.insert(hash) {
+					encoded_size_update += value.encoded_size();
+				}
+				self.recorded_keys
+					.entry(full_key.into())
+					.and_modify(|e| *e = RecordedForKey::Value)
+					.or_insert_with(|| RecordedForKey::Value);
+			},
+			TrieAccess::Hash { full_key } => {
+				self.recorded_keys
+					.entry(full_key.into())
+					.or_insert_with(|| RecordedForKey::Hash);
+			},
+			TrieAccess::NonExisting { full_key } => {
+				self.recorded_keys
+					.entry(full_key.into())
+					.and_modify(|e| *e = RecordedForKey::Value)
+					.or_insert_with(|| RecordedForKey::Value);
+			},
+			TrieAccess::InlineValue { full_key } => {
+				self.recorded_keys
+					.entry(full_key.into())
+					.and_modify(|e| *e = RecordedForKey::Value)
+					.or_insert_with(|| RecordedForKey::Value);
+			},
+		};
+
+		*self.encoded_size += encoded_size_update;
+	}
+
+	fn trie_nodes_recorded_for_key(&self, key: &[u8]) -> RecordedForKey {
+		self.recorded_keys.get(key).copied().unwrap_or(RecordedForKey::None)
+	}
+}
+
+#[derive(Clone)]
+pub struct SizeOnlyRecorderProvider<H: Hasher> {
+	seen_nodes: SeenNodes<H>,
+	encoded_size: Rc<RefCell<usize>>,
+	recorded_keys: Rc<RefCell<HashMap<Rc<[u8]>, RecordedForKey, RandomState>>>,
+}
+
+impl<H: Hasher> Default for SizeOnlyRecorderProvider<H> {
+	fn default() -> Self {
+		Self {
+			seen_nodes: Default::default(),
+			encoded_size: Default::default(),
+			recorded_keys: Default::default(),
+		}
+	}
+}
+
+impl<H: Hasher> SizeOnlyRecorderProvider<H> {
+	/// Use the given `seen_nodes` to populate the internal state.
+	#[cfg(not(feature = "std"))]
+	pub(crate) fn with_seen_nodes(seen_nodes: SeenNodes<H>) -> Self {
+		Self { seen_nodes, ..Default::default() }
+	}
+}
+
+impl<H: trie_db::Hasher> pezsp_trie::TrieRecorderProvider<H> for SizeOnlyRecorderProvider<H> {
+	type Recorder<'a>
+		= SizeOnlyRecorder<'a, H>
+	where
+		H: 'a;
+
+	fn drain_storage_proof(self) -> Option<StorageProof> {
+		None
+	}
+
+	fn as_trie_recorder(&self, _storage_root: H::Out) -> Self::Recorder<'_> {
+		SizeOnlyRecorder {
+			encoded_size: self.encoded_size.borrow_mut(),
+			seen_nodes: self.seen_nodes.borrow_mut(),
+			recorded_keys: self.recorded_keys.borrow_mut(),
+		}
+	}
+}
+
+impl<H: trie_db::Hasher> ProofSizeProvider for SizeOnlyRecorderProvider<H> {
+	fn estimate_encoded_size(&self) -> usize {
+		*self.encoded_size.borrow()
+	}
+}
+
+// This is safe here since we are single-threaded in WASM
+unsafe impl<H: Hasher> Send for SizeOnlyRecorderProvider<H> {}
+unsafe impl<H: Hasher> Sync for SizeOnlyRecorderProvider<H> {}
+
+#[cfg(test)]
+mod tests {
+	use rand::Rng;
+	use pezsp_trie::{
+		cache::{CacheSize, SharedTrieCache},
+		MemoryDB, ProofSizeProvider, TrieRecorderProvider,
+	};
+	use trie_db::{Trie, TrieDBBuilder, TrieDBMutBuilder, TrieHash, TrieMut, TrieRecorder};
+	use trie_standardmap::{Alphabet, StandardMap, ValueMode};
+
+	use super::*;
+
+	type Recorder = pezsp_trie::recorder::Recorder<pezsp_core::Blake2Hasher>;
+
+	fn create_trie() -> (
+		pezsp_trie::MemoryDB<pezsp_core::Blake2Hasher>,
+		TrieHash<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>,
+		Vec<(Vec<u8>, Vec<u8>)>,
+	) {
+		let mut db = MemoryDB::default();
+		let mut root = Default::default();
+
+		let mut seed = Default::default();
+		let test_data: Vec<(Vec<u8>, Vec<u8>)> = StandardMap {
+			alphabet: Alphabet::Low,
+			min_key: 16,
+			journal_key: 0,
+			value_mode: ValueMode::Random,
+			count: 1000,
+		}
+		.make_with(&mut seed)
+		.into_iter()
+		.map(|(k, v)| {
+			// Double the length so we end up with some values of 2 bytes and some of 64
+			let v = [v.clone(), v].concat();
+			(k, v)
+		})
+		.collect();
+
+		// Fill database with values
+		{
+			let mut trie = TrieDBMutBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(
+				&mut db, &mut root,
+			)
+			.build();
+			for (k, v) in &test_data {
+				trie.insert(k, v).expect("Inserts data");
+			}
+		}
+
+		(db, root, test_data)
+	}
+
+	#[test]
+	fn recorder_equivalence_cache() {
+		let (db, root, test_data) = create_trie();
+
+		let mut rng = rand::thread_rng();
+		for _ in 1..10 {
+			let reference_recorder = Recorder::default();
+			let recorder_for_test: SizeOnlyRecorderProvider<pezsp_core::Blake2Hasher> =
+				SizeOnlyRecorderProvider::default();
+			let reference_cache: SharedTrieCache<pezsp_core::Blake2Hasher> =
+				SharedTrieCache::new(CacheSize::new(1024 * 5), None);
+			let cache_for_test: SharedTrieCache<pezsp_core::Blake2Hasher> =
+				SharedTrieCache::new(CacheSize::new(1024 * 5), None);
+			{
+				let local_cache = cache_for_test.local_cache_untrusted();
+				let mut trie_cache_for_reference = local_cache.as_trie_db_cache(root);
+				let mut reference_trie_recorder = reference_recorder.as_trie_recorder(root);
+				let reference_trie =
+					TrieDBBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(&db, &root)
+						.with_recorder(&mut reference_trie_recorder)
+						.with_cache(&mut trie_cache_for_reference)
+						.build();
+
+				let local_cache_for_test = reference_cache.local_cache_untrusted();
+				let mut trie_cache_for_test = local_cache_for_test.as_trie_db_cache(root);
+				let mut trie_recorder_under_test = recorder_for_test.as_trie_recorder(root);
+				let test_trie =
+					TrieDBBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(&db, &root)
+						.with_recorder(&mut trie_recorder_under_test)
+						.with_cache(&mut trie_cache_for_test)
+						.build();
+
+				// Access random values from the test data
+				for _ in 0..100 {
+					let index: usize = rng.gen_range(0..test_data.len());
+					test_trie.get(&test_data[index].0).unwrap().unwrap();
+					reference_trie.get(&test_data[index].0).unwrap().unwrap();
+				}
+
+				// Check that we have the same nodes recorded for both recorders
+				for (key, _) in test_data.iter() {
+					let reference = reference_trie_recorder.trie_nodes_recorded_for_key(key);
+					let test_value = trie_recorder_under_test.trie_nodes_recorded_for_key(key);
+					assert_eq!(format!("{:?}", reference), format!("{:?}", test_value));
+				}
+			}
+
+			// Check that we have the same size recorded for both recorders
+			assert_eq!(
+				reference_recorder.estimate_encoded_size(),
+				recorder_for_test.estimate_encoded_size()
+			);
+		}
+	}
+
+	#[test]
+	fn recorder_equivalence_no_cache() {
+		let (db, root, test_data) = create_trie();
+
+		let mut rng = rand::thread_rng();
+		for _ in 1..10 {
+			let reference_recorder = Recorder::default();
+			let recorder_for_test: SizeOnlyRecorderProvider<pezsp_core::Blake2Hasher> =
+				SizeOnlyRecorderProvider::default();
+			{
+				let mut reference_trie_recorder = reference_recorder.as_trie_recorder(root);
+				let reference_trie =
+					TrieDBBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(&db, &root)
+						.with_recorder(&mut reference_trie_recorder)
+						.build();
+
+				let mut trie_recorder_under_test = recorder_for_test.as_trie_recorder(root);
+				let test_trie =
+					TrieDBBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(&db, &root)
+						.with_recorder(&mut trie_recorder_under_test)
+						.build();
+
+				for _ in 0..200 {
+					let index: usize = rng.gen_range(0..test_data.len());
+					test_trie.get(&test_data[index].0).unwrap().unwrap();
+					reference_trie.get(&test_data[index].0).unwrap().unwrap();
+				}
+
+				// Check that we have the same nodes recorded for both recorders
+				for (key, _) in test_data.iter() {
+					let reference = reference_trie_recorder.trie_nodes_recorded_for_key(key);
+					let test_value = trie_recorder_under_test.trie_nodes_recorded_for_key(key);
+					assert_eq!(format!("{:?}", reference), format!("{:?}", test_value));
+				}
+			}
+
+			// Check that we have the same size recorded for both recorders
+			assert_eq!(
+				reference_recorder.estimate_encoded_size(),
+				recorder_for_test.estimate_encoded_size()
+			);
+		}
+	}
+}