Feat: Add API and mechanism to retrieve additional top-level and child proofs via the relay state proof

cumulus/client/consensus/aura/src/collator.rs

@@ -33,6 +33,7 @@ use cumulus_client_consensus_common::{
 use cumulus_client_parachain_inherent::{ParachainInherentData, ParachainInherentDataProvider};
 use cumulus_primitives_core::{
 	relay_chain::Hash as PHash, DigestItem, ParachainBlockData, PersistedValidationData,
+	RelayProofRequest,
 };
 use cumulus_relay_chain_interface::RelayChainInterface;
 use sc_client_api::BackendTransaction;
@@ -177,6 +178,7 @@ where
 		parent_hash: Block::Hash,
 		timestamp: impl Into<Option<Timestamp>>,
 		relay_parent_descendants: Option<RelayParentData>,
+		relay_proof_request: RelayProofRequest,
 		collator_peer_id: PeerId,
 	) -> Result<(ParachainInherentData, InherentData), Box<dyn Error + Send + Sync + 'static>> {
 		let paras_inherent_data = ParachainInherentDataProvider::create_at(
@@ -188,6 +190,7 @@ where
 				.map(RelayParentData::into_inherent_descendant_list)
 				.unwrap_or_default(),
 			Vec::new(),
+			relay_proof_request,
 			collator_peer_id,
 		)
 		.await;
@@ -224,6 +227,7 @@ where
 		validation_data: &PersistedValidationData,
 		parent_hash: Block::Hash,
 		timestamp: impl Into<Option<Timestamp>>,
+		relay_proof_request: RelayProofRequest,
 		collator_peer_id: PeerId,
 	) -> Result<(ParachainInherentData, InherentData), Box<dyn Error + Send + Sync + 'static>> {
 		self.create_inherent_data_with_rp_offset(
@@ -232,6 +236,7 @@ where
 			parent_hash,
 			timestamp,
 			None,
+			relay_proof_request,
 			collator_peer_id,
 		)
 		.await

cumulus/client/consensus/aura/src/collators/basic.rs

@@ -238,6 +238,7 @@ where
 						&validation_data,
 						parent_hash,
 						claim.timestamp(),
+						Default::default(),
 						params.collator_peer_id,
 					)
 					.await

cumulus/client/consensus/aura/src/collators/lookahead.rs

@@ -36,7 +36,7 @@ use codec::{Codec, Encode};
 use cumulus_client_collator::service::ServiceInterface as CollatorServiceInterface;
 use cumulus_client_consensus_common::{self as consensus_common, ParachainBlockImportMarker};
 use cumulus_primitives_aura::AuraUnincludedSegmentApi;
-use cumulus_primitives_core::{CollectCollationInfo, PersistedValidationData};
+use cumulus_primitives_core::{CollectCollationInfo, KeyToIncludeInRelayProof, PersistedValidationData};
 use cumulus_relay_chain_interface::RelayChainInterface;
 use sp_consensus::Environment;
 
@@ -164,8 +164,10 @@ where
 		+ Send
 		+ Sync
 		+ 'static,
-	Client::Api:
-		AuraApi<Block, P::Public> + CollectCollationInfo<Block> + AuraUnincludedSegmentApi<Block>,
+	Client::Api: AuraApi<Block, P::Public>
+		+ CollectCollationInfo<Block>
+		+ AuraUnincludedSegmentApi<Block>
+		+ KeyToIncludeInRelayProof<Block>,
 	Backend: sc_client_api::Backend<Block> + 'static,
 	RClient: RelayChainInterface + Clone + 'static,
 	CIDP: CreateInherentDataProviders<Block, ()> + 'static,
@@ -216,8 +218,10 @@ where
 		+ Send
 		+ Sync
 		+ 'static,
-	Client::Api:
-		AuraApi<Block, P::Public> + CollectCollationInfo<Block> + AuraUnincludedSegmentApi<Block>,
+	Client::Api: AuraApi<Block, P::Public>
+		+ CollectCollationInfo<Block>
+		+ AuraUnincludedSegmentApi<Block>
+		+ KeyToIncludeInRelayProof<Block>,
 	Backend: sc_client_api::Backend<Block> + 'static,
 	RClient: RelayChainInterface + Clone + 'static,
 	CIDP: CreateInherentDataProviders<Block, ()> + 'static,
@@ -392,12 +396,15 @@ where
 
 				// Build and announce collations recursively until
 				// `can_build_upon` fails or building a collation fails.
+				let relay_proof_request = super::get_relay_proof_request(&*params.para_client, parent_hash);
+
 				let (parachain_inherent_data, other_inherent_data) = match collator
 					.create_inherent_data(
 						relay_parent,
 						&validation_data,
 						parent_hash,
 						slot_claim.timestamp(),
+						relay_proof_request,
 						params.collator_peer_id,
 					)
 					.await

cumulus/client/consensus/aura/src/collators/mod.rs

@@ -25,7 +25,9 @@ use crate::collator::SlotClaim;
 use codec::Codec;
 use cumulus_client_consensus_common::{self as consensus_common, ParentSearchParams};
 use cumulus_primitives_aura::{AuraUnincludedSegmentApi, Slot};
-use cumulus_primitives_core::{relay_chain::Header as RelayHeader, BlockT};
+use cumulus_primitives_core::{
+	relay_chain::Header as RelayHeader, BlockT, KeyToIncludeInRelayProof, RelayProofRequest,
+};
 use cumulus_relay_chain_interface::{OverseerHandle, RelayChainInterface};
 use polkadot_node_subsystem::messages::{CollatorProtocolMessage, RuntimeApiRequest};
 use polkadot_node_subsystem_util::runtime::ClaimQueueSnapshot;
@@ -662,6 +664,34 @@ mod tests {
 	}
 }
 
+/// Fetches relay chain storage proof requests from the parachain runtime.
+///
+/// Queries the runtime API to determine which relay chain storage keys
+/// (both top-level and child trie keys) should be included in the relay chain state proof.
+///
+/// Falls back to an empty request if the runtime API call fails or is not implemented.
+fn get_relay_proof_request<Block, Client>(
+	client: &Client,
+	parent_hash: Block::Hash,
+) -> RelayProofRequest
+where
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block>,
+	Client::Api: KeyToIncludeInRelayProof<Block>,
+{
+	client
+		.runtime_api()
+		.keys_to_prove(parent_hash)
+		.unwrap_or_else(|e| {
+			tracing::debug!(
+				target: crate::LOG_TARGET,
+				error = ?e,
+				"Failed to fetch relay proof requests from runtime, using empty request"
+			);
+			Default::default()
+		})
+}
+
 /// Holds a relay parent and its descendants.
 pub struct RelayParentData {
 	/// The relay parent block header

cumulus/client/consensus/aura/src/collators/slot_based/block_builder_task.rs

@@ -359,14 +359,14 @@ where
 				relay_parent_storage_root: *relay_parent_header.state_root(),
 				max_pov_size: *max_pov_size,
 			};
-
 			let (parachain_inherent_data, other_inherent_data) = match collator
 				.create_inherent_data_with_rp_offset(
 					relay_parent,
 					&validation_data,
 					parent_hash,
 					slot_claim.timestamp(),
 					Some(rp_data),
+					Default::default(),
 					collator_peer_id,
 				)
 				.await

cumulus/client/consensus/aura/src/collators/slot_based/tests.rs

@@ -566,6 +566,15 @@ impl RelayChainInterface for TestRelayClient {
 		unimplemented!("Not needed for test")
 	}
 
+	async fn prove_child_read(
+		&self,
+		_: RelayHash,
+		_: &cumulus_relay_chain_interface::ChildInfo,
+		_: &[Vec<u8>],
+	) -> RelayChainResult<sc_client_api::StorageProof> {
+		unimplemented!("Not needed for test")
+	}
+
 	async fn wait_for_block(&self, _: RelayHash) -> RelayChainResult<()> {
 		unimplemented!("Not needed for test")
 	}

cumulus/client/parachain-inherent/src/lib.rs

@@ -24,25 +24,26 @@ use cumulus_primitives_core::{
 		self, ApprovedPeerId, Block as RelayBlock, Hash as PHash, Header as RelayHeader,
 		HrmpChannelId,
 	},
-	ParaId, PersistedValidationData,
+	ParaId, PersistedValidationData, RelayProofRequest, RelayStorageKey,
 };
 pub use cumulus_primitives_parachain_inherent::{ParachainInherentData, INHERENT_IDENTIFIER};
 use cumulus_relay_chain_interface::RelayChainInterface;
 pub use mock::{MockValidationDataInherentDataProvider, MockXcmConfig};
 use sc_network_types::PeerId;
+use sp_state_machine::StorageProof;
+use sp_storage::ChildInfo;
 
 const LOG_TARGET: &str = "parachain-inherent";
 
-/// Collect the relevant relay chain state in form of a proof for putting it into the validation
-/// data inherent.
-async fn collect_relay_storage_proof(
+/// Builds the list of static relay chain storage keys that are always needed for parachain
+/// validation.
+async fn get_static_relay_storage_keys(
 	relay_chain_interface: &impl RelayChainInterface,
 	para_id: ParaId,
 	relay_parent: PHash,
 	include_authorities: bool,
 	include_next_authorities: bool,
-	additional_relay_state_keys: Vec<Vec<u8>>,
-) -> Option<sp_state_machine::StorageProof> {
+) -> Option<Vec<Vec<u8>>> {
 	use relay_chain::well_known_keys as relay_well_known_keys;
 
 	let ingress_channels = relay_chain_interface
@@ -136,25 +137,95 @@ async fn collect_relay_storage_proof(
 		relevant_keys.push(relay_well_known_keys::NEXT_AUTHORITIES.to_vec());
 	}
 
-	// Add additional relay state keys
+	Some(relevant_keys)
+}
+
+/// Collect the relevant relay chain state in form of a proof for putting it into the validation
+/// data inherent.
+async fn collect_relay_storage_proof(
+	relay_chain_interface: &impl RelayChainInterface,
+	para_id: ParaId,
+	relay_parent: PHash,
+	include_authorities: bool,
+	include_next_authorities: bool,
+	additional_relay_state_keys: Vec<Vec<u8>>,
+	relay_proof_request: RelayProofRequest,
+) -> Option<StorageProof> {
+	// Get static keys that are always needed
+	let mut all_top_keys = get_static_relay_storage_keys(
+		relay_chain_interface,
+		para_id,
+		relay_parent,
+		include_authorities,
+		include_next_authorities,
+	)
+	.await?;
+
+	// Add additional_relay_state_keys
 	let unique_keys: Vec<Vec<u8>> = additional_relay_state_keys
 		.into_iter()
-		.filter(|key| !relevant_keys.contains(key))
+		.filter(|key| !all_top_keys.contains(key))
 		.collect();
-	relevant_keys.extend(unique_keys);
+	all_top_keys.extend(unique_keys);
 
-	relay_chain_interface
-		.prove_read(relay_parent, &relevant_keys)
-		.await
-		.map_err(|e| {
+	// Group requested keys by storage type
+	let RelayProofRequest { keys } = relay_proof_request;
+	let mut child_keys: std::collections::BTreeMap<Vec<u8>, Vec<Vec<u8>>> =
+		std::collections::BTreeMap::new();
+
+	for key in keys {
+		match key {
+			RelayStorageKey::Top(k) => {
+				if !all_top_keys.contains(&k) {
+					all_top_keys.push(k);
+				}
+			},
+			RelayStorageKey::Child { storage_key, key } => {
+				child_keys.entry(storage_key).or_default().push(key);
+			},
+		}
+	}
+
+	// Collect all storage proofs
+	let mut all_proofs = Vec::new();
+
+	// Collect top-level storage proof.
+	match relay_chain_interface.prove_read(relay_parent, &all_top_keys).await {
+		Ok(top_proof) => {
+			all_proofs.push(top_proof);
+		},
+		Err(e) => {
 			tracing::error!(
 				target: LOG_TARGET,
 				relay_parent = ?relay_parent,
 				error = ?e,
-				"Cannot obtain read proof from relay chain.",
+				"Cannot obtain relay chain storage proof.",
 			);
-		})
-		.ok()
+			return None;
+		},
+	}
+
+	// Collect child trie proofs
+	for (storage_key, data_keys) in child_keys {
+		let child_info = ChildInfo::new_default(&storage_key);
+		match relay_chain_interface.prove_child_read(relay_parent, &child_info, &data_keys).await {
+			Ok(child_proof) => {
+				all_proofs.push(child_proof);
+			},
+			Err(e) => {
+				tracing::error!(
+					target: LOG_TARGET,
+					relay_parent = ?relay_parent,
+					child_trie_id = ?child_info.storage_key(),
+					error = ?e,
+					"Cannot obtain child trie proof from relay chain.",
+				);
+			},
+		}
+	}
+
+	// Merge all proofs
+	Some(StorageProof::merge(all_proofs))
 }
 
 pub struct ParachainInherentDataProvider;
@@ -170,6 +241,7 @@ impl ParachainInherentDataProvider {
 		para_id: ParaId,
 		relay_parent_descendants: Vec<RelayHeader>,
 		additional_relay_state_keys: Vec<Vec<u8>>,
+		relay_proof_request: RelayProofRequest,
 		collator_peer_id: PeerId,
 	) -> Option<ParachainInherentData> {
 		let collator_peer_id = ApprovedPeerId::try_from(collator_peer_id.to_bytes())
@@ -195,6 +267,7 @@ impl ParachainInherentDataProvider {
 			!relay_parent_descendants.is_empty(),
 			include_next_authorities,
 			additional_relay_state_keys,
+			relay_proof_request,
 		)
 		.await?;
 

cumulus/client/relay-chain-inprocess-interface/src/lib.rs

@@ -33,7 +33,9 @@ use cumulus_primitives_core::{
 	},
 	InboundDownwardMessage, ParaId, PersistedValidationData,
 };
-use cumulus_relay_chain_interface::{RelayChainError, RelayChainInterface, RelayChainResult};
+use cumulus_relay_chain_interface::{
+	ChildInfo, RelayChainError, RelayChainInterface, RelayChainResult,
+};
 use futures::{FutureExt, Stream, StreamExt};
 use polkadot_primitives::CandidateEvent;
 use polkadot_service::{
@@ -240,6 +242,18 @@ impl RelayChainInterface for RelayChainInProcessInterface {
 			.map_err(RelayChainError::StateMachineError)
 	}
 
+	async fn prove_child_read(
+		&self,
+		relay_parent: PHash,
+		child_info: &ChildInfo,
+		child_keys: &[Vec<u8>],
+	) -> RelayChainResult<StorageProof> {
+		let state_backend = self.backend.state_at(relay_parent, TrieCacheContext::Untrusted)?;
+
+		sp_state_machine::prove_child_read(state_backend, child_info, child_keys)
+			.map_err(RelayChainError::StateMachineError)
+	}
+
 	/// Wait for a given relay chain block in an async way.
 	///
 	/// The caller needs to pass the hash of a block it waits for and the function will return when

cumulus/client/relay-chain-interface/Cargo.toml

@@ -21,6 +21,7 @@ sc-network = { workspace = true, default-features = true }
 sp-api = { workspace = true, default-features = true }
 sp-blockchain = { workspace = true, default-features = true }
 sp-state-machine = { workspace = true, default-features = true }
+sp-storage = { workspace = true, default-features = true }
 sp-version = { workspace = true }
 
 async-trait = { workspace = true }