feat: imp KeyToIncludeInRelayProof for test pallet

Author: metricaez

Cargo.lock

@@ -41,6 +41,9 @@ sp-session = { workspace = true }
 sp-transaction-pool = { workspace = true }
 sp-version = { workspace = true }
 
+# Polkadot
+polkadot-primitives = { workspace = true }
+
 # Cumulus
 cumulus-pallet-aura-ext = { workspace = true }
 cumulus-pallet-parachain-system = { workspace = true }
@@ -75,6 +78,7 @@ std = [
 	"pallet-timestamp/std",
 	"pallet-transaction-payment/std",
 	"parachain-info/std",
+	"polkadot-primitives/std",
 	"scale-info/std",
 	"serde_json/std",
 	"sp-api/std",

cumulus/test/runtime/Cargo.toml

@@ -374,7 +374,7 @@ impl cumulus_pallet_parachain_system::Config for Runtime {
 	type WeightInfo = ();
 	type SelfParaId = parachain_info::Pallet<Runtime>;
 	type RuntimeEvent = RuntimeEvent;
-	type OnSystemEvent = ();
+	type OnSystemEvent = TestPallet;
 	type OutboundXcmpMessageSource = ();
 	// Ignore all DMP messages by enqueueing them into `()`:
 	type DmpQueue = frame_support::traits::EnqueueWithOrigin<(), sp_core::ConstU8<0>>;
@@ -643,8 +643,15 @@ impl_runtime_apis! {
 	}
 
 	impl cumulus_primitives_core::KeyToIncludeInRelayProof<Block> for Runtime {
-		fn keys_to_prove() -> RelayProofRequest {
-			Default::default()
+		fn keys_to_prove() -> cumulus_primitives_core::RelayProofRequest {
+			use cumulus_primitives_core::RelayStorageKey;
+
+			RelayProofRequest {
+				keys: vec![
+					// Request a well-known key to verify its inclusion in the relay proof.
+					RelayStorageKey::Top(test_pallet::RELAY_EPOCH_INDEX_KEY.to_vec()),
+				],
+			}
 		}
 	}
 }

cumulus/test/runtime/src/lib.rs

@@ -17,10 +17,15 @@
 /// A special pallet that exposes dispatchables that are only useful for testing.
 pub use pallet::*;
 
+use polkadot_primitives::well_known_keys;
+
 /// Some key that we set in genesis and only read in [`TestOnRuntimeUpgrade`] to ensure that
 /// [`OnRuntimeUpgrade`] works as expected.
 pub const TEST_RUNTIME_UPGRADE_KEY: &[u8] = b"+test_runtime_upgrade_key+";
 
+/// A well-known key to request for inclusion in the proof.
+pub use well_known_keys::EPOCH_INDEX as RELAY_EPOCH_INDEX_KEY;
+
 #[frame_support::pallet(dev_mode)]
 pub mod pallet {
 	use crate::test_pallet::TEST_RUNTIME_UPGRADE_KEY;
@@ -121,3 +126,27 @@ pub mod pallet {
 		}
 	}
 }
+
+impl<T: Config> cumulus_pallet_parachain_system::OnSystemEvent for Pallet<T> {
+	fn on_validation_data(_data: &cumulus_primitives_core::PersistedValidationData) {
+		// Nothing to do here for tests
+	}
+
+	fn on_validation_code_applied() {
+		// Nothing to do here for tests
+	}
+
+	fn on_relay_state_proof(
+		relay_state_proof: &cumulus_pallet_parachain_system::relay_state_snapshot::RelayChainStateProof,
+	) -> frame_support::weights::Weight {
+		use crate::test_pallet::RELAY_EPOCH_INDEX_KEY;
+
+		// Expect the requested key to be part of the proof.
+		relay_state_proof
+			.read_optional_entry::<u64>(RELAY_EPOCH_INDEX_KEY)
+			.expect("Invalid relay chain state proof")
+			.expect("EPOCH_INDEX must be present");
+
+		frame_support::weights::Weight::zero()
+	}
+}