cumulus_test_relay_sproof_builder/
lib.rs

1// Copyright (C) Parity Technologies (UK) Ltd.
2// This file is part of Cumulus.
3// SPDX-License-Identifier: Apache-2.0
4
5// Licensed under the Apache License, Version 2.0 (the "License");
6// you may not use this file except in compliance with the License.
7// You may obtain a copy of the License at
8//
9// 	http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS,
13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14// See the License for the specific language governing permissions and
15// limitations under the License.
16
17extern crate alloc;
18
19use alloc::collections::btree_map::BTreeMap;
20use codec::{Decode, Encode};
21use cumulus_primitives_core::{
22	relay_chain, AbridgedHostConfiguration, AbridgedHrmpChannel, ParaId,
23};
24use polkadot_primitives::{Header, UpgradeGoAhead};
25use sp_consensus_babe::{
26	digests::{CompatibleDigestItem, PreDigest, PrimaryPreDigest},
27	AuthorityId, AuthorityPair, BabeAuthorityWeight,
28};
29use sp_core::{
30	sr25519::vrf::{VrfPreOutput, VrfProof, VrfSignature},
31	Pair, H256,
32};
33use sp_runtime::{
34	traits::{HashingFor, Header as HeaderT},
35	Digest, DigestItem,
36};
37use sp_trie::PrefixedMemoryDB;
38
39/// Builds a sproof (portmanteau of 'spoof' and 'proof') of the relay chain state.
40#[derive(Clone)]
41pub struct RelayStateSproofBuilder {
42	/// The para id of the current parachain.
43	///
44	/// This doesn't get into the storage proof produced by the builder, however, it is used for
45	/// generation of the storage image and by auxiliary methods.
46	///
47	/// It's recommended to change this value once in the very beginning of usage.
48	///
49	/// The default value is 200.
50	pub para_id: ParaId,
51
52	pub host_config: AbridgedHostConfiguration,
53	pub dmq_mqc_head: Option<relay_chain::Hash>,
54	pub upgrade_go_ahead: Option<UpgradeGoAhead>,
55	pub relay_dispatch_queue_remaining_capacity: Option<(u32, u32)>,
56	pub hrmp_ingress_channel_index: Option<Vec<ParaId>>,
57	pub hrmp_egress_channel_index: Option<Vec<ParaId>>,
58	pub hrmp_channels: BTreeMap<relay_chain::HrmpChannelId, AbridgedHrmpChannel>,
59	pub current_slot: relay_chain::Slot,
60	pub current_epoch: u64,
61	pub randomness: relay_chain::Hash,
62	pub additional_key_values: Vec<(Vec<u8>, Vec<u8>)>,
63	pub included_para_head: Option<relay_chain::HeadData>,
64	pub num_authorities: u64,
65}
66
67impl Default for RelayStateSproofBuilder {
68	fn default() -> Self {
69		RelayStateSproofBuilder {
70			para_id: ParaId::from(200),
71			host_config: cumulus_primitives_core::AbridgedHostConfiguration {
72				max_code_size: 2 * 1024 * 1024,
73				max_head_data_size: 1024 * 1024,
74				max_upward_queue_count: 8,
75				max_upward_queue_size: 1024,
76				max_upward_message_size: 256,
77				max_upward_message_num_per_candidate: 5,
78				hrmp_max_message_num_per_candidate: 5,
79				validation_upgrade_cooldown: 6,
80				validation_upgrade_delay: 6,
81				async_backing_params: relay_chain::AsyncBackingParams {
82					allowed_ancestry_len: 0,
83					max_candidate_depth: 0,
84				},
85			},
86			dmq_mqc_head: None,
87			upgrade_go_ahead: None,
88			relay_dispatch_queue_remaining_capacity: None,
89			hrmp_ingress_channel_index: None,
90			hrmp_egress_channel_index: None,
91			hrmp_channels: BTreeMap::new(),
92			current_slot: 0.into(),
93			current_epoch: 0u64,
94			randomness: relay_chain::Hash::default(),
95			additional_key_values: vec![],
96			included_para_head: None,
97			num_authorities: 1,
98		}
99	}
100}
101
102impl RelayStateSproofBuilder {
103	/// Returns a mutable reference to HRMP channel metadata for a channel (`sender`,
104	/// `self.para_id`).
105	///
106	/// If there is no channel, a new default one is created.
107	///
108	/// It also updates the `hrmp_ingress_channel_index`, creating it if needed.
109	pub fn upsert_inbound_channel(&mut self, sender: ParaId) -> &mut AbridgedHrmpChannel {
110		let in_index = self.hrmp_ingress_channel_index.get_or_insert_with(Vec::new);
111		if let Err(idx) = in_index.binary_search(&sender) {
112			in_index.insert(idx, sender);
113		}
114
115		self.upsert_channel(relay_chain::HrmpChannelId { sender, recipient: self.para_id })
116	}
117
118	/// Returns a mutable reference to HRMP channel metadata for a channel (`self.para_id`,
119	/// `recipient`).
120	///
121	/// If there is no channel, a new default one is created.
122	///
123	/// It also updates the `hrmp_egress_channel_index`, creating it if needed.
124	pub fn upsert_outbound_channel(&mut self, recipient: ParaId) -> &mut AbridgedHrmpChannel {
125		let in_index = self.hrmp_egress_channel_index.get_or_insert_with(Vec::new);
126		if let Err(idx) = in_index.binary_search(&recipient) {
127			in_index.insert(idx, recipient);
128		}
129
130		self.upsert_channel(relay_chain::HrmpChannelId { sender: self.para_id, recipient })
131	}
132
133	/// Creates a new default entry in the hrmp channels mapping if not exists, and returns mutable
134	/// reference to it.
135	fn upsert_channel(&mut self, id: relay_chain::HrmpChannelId) -> &mut AbridgedHrmpChannel {
136		self.hrmp_channels.entry(id).or_insert_with(|| AbridgedHrmpChannel {
137			max_capacity: 0,
138			max_total_size: 0,
139			max_message_size: 0,
140			msg_count: 0,
141			total_size: 0,
142			mqc_head: None,
143		})
144	}
145
146	/// Build sproof and generate relay parent descendants with the configured authorities.
147	///
148	/// Returns a tuple of (state_root, storage_proof, relay_parent_descendants).
149	pub fn into_state_root_proof_and_descendants(
150		self,
151		relay_parent_offset: u64,
152	) -> (polkadot_primitives::Hash, sp_state_machine::StorageProof, Vec<Header>) {
153		let authorities = generate_authority_pairs(self.num_authorities);
154		let (state_root, proof) = self.into_state_root_and_proof();
155		let descendants =
156			build_relay_parent_descendants(relay_parent_offset + 1, state_root.into(), authorities);
157		(state_root, proof, descendants)
158	}
159
160	pub fn into_state_root_and_proof(
161		mut self,
162	) -> (polkadot_primitives::Hash, sp_state_machine::StorageProof) {
163		// Generate and add authorities if num_authorities is set
164		if self.num_authorities > 0 {
165			let authorities = generate_authority_pairs(self.num_authorities);
166			let auth_pair = convert_to_authority_weight_pair(&authorities);
167
168			// Add authorities to the sproof builder
169			self.additional_key_values.push((
170				relay_chain::well_known_keys::AUTHORITIES.to_vec(),
171				auth_pair.clone().encode(),
172			));
173			self.additional_key_values.push((
174				relay_chain::well_known_keys::NEXT_AUTHORITIES.to_vec(),
175				auth_pair.encode(),
176			));
177		}
178
179		let (db, root) =
180			PrefixedMemoryDB::<HashingFor<polkadot_primitives::Block>>::default_with_root();
181		let state_version = Default::default(); // for test using default.
182		let mut backend = sp_state_machine::TrieBackendBuilder::new(db, root).build();
183
184		let mut relevant_keys = Vec::new();
185		{
186			use codec::Encode as _;
187
188			let mut insert = |key: Vec<u8>, value: Vec<u8>| {
189				relevant_keys.push(key.clone());
190				backend.insert(vec![(None, vec![(key, Some(value))])], state_version);
191			};
192
193			insert(relay_chain::well_known_keys::ACTIVE_CONFIG.to_vec(), self.host_config.encode());
194			if let Some(dmq_mqc_head) = self.dmq_mqc_head {
195				insert(
196					relay_chain::well_known_keys::dmq_mqc_head(self.para_id),
197					dmq_mqc_head.encode(),
198				);
199			}
200			if let Some(para_head) = self.included_para_head {
201				insert(relay_chain::well_known_keys::para_head(self.para_id), para_head.encode());
202			}
203			if let Some(relay_dispatch_queue_remaining_capacity) =
204				self.relay_dispatch_queue_remaining_capacity
205			{
206				insert(
207					relay_chain::well_known_keys::relay_dispatch_queue_remaining_capacity(
208						self.para_id,
209					)
210					.key,
211					relay_dispatch_queue_remaining_capacity.encode(),
212				);
213			}
214			if let Some(upgrade_go_ahead) = self.upgrade_go_ahead {
215				insert(
216					relay_chain::well_known_keys::upgrade_go_ahead_signal(self.para_id),
217					upgrade_go_ahead.encode(),
218				);
219			}
220			if let Some(hrmp_ingress_channel_index) = self.hrmp_ingress_channel_index {
221				let mut sorted = hrmp_ingress_channel_index.clone();
222				sorted.sort();
223				assert_eq!(sorted, hrmp_ingress_channel_index);
224
225				insert(
226					relay_chain::well_known_keys::hrmp_ingress_channel_index(self.para_id),
227					hrmp_ingress_channel_index.encode(),
228				);
229			}
230			if let Some(hrmp_egress_channel_index) = self.hrmp_egress_channel_index {
231				let mut sorted = hrmp_egress_channel_index.clone();
232				sorted.sort();
233				assert_eq!(sorted, hrmp_egress_channel_index);
234
235				insert(
236					relay_chain::well_known_keys::hrmp_egress_channel_index(self.para_id),
237					hrmp_egress_channel_index.encode(),
238				);
239			}
240			for (channel, metadata) in self.hrmp_channels {
241				insert(relay_chain::well_known_keys::hrmp_channels(channel), metadata.encode());
242			}
243			insert(relay_chain::well_known_keys::EPOCH_INDEX.to_vec(), self.current_epoch.encode());
244			insert(
245				relay_chain::well_known_keys::ONE_EPOCH_AGO_RANDOMNESS.to_vec(),
246				self.randomness.encode(),
247			);
248			insert(relay_chain::well_known_keys::CURRENT_SLOT.to_vec(), self.current_slot.encode());
249
250			for (key, value) in self.additional_key_values {
251				insert(key, value);
252			}
253		}
254
255		let root = *backend.root();
256		let proof = sp_state_machine::prove_read(backend, relevant_keys).expect("prove read");
257		(root, proof)
258	}
259}
260
261/// Generate a vector of AuthorityPairs
262pub fn generate_authority_pairs(num_authorities: u64) -> Vec<AuthorityPair> {
263	(0..num_authorities).map(|i| AuthorityPair::from_seed(&[i as u8; 32])).collect()
264}
265
266/// Convert AuthorityPair to (AuthorityId, BabeAuthorityWeight)
267fn convert_to_authority_weight_pair(
268	authorities: &[AuthorityPair],
269) -> Vec<(AuthorityId, BabeAuthorityWeight)> {
270	authorities
271		.iter()
272		.map(|auth| (auth.public().into(), Default::default()))
273		.collect()
274}
275
276/// Add a BABE pre-digest to a generic header
277fn add_babe_pre_digest(header: &mut Header, authority_index: u32, block_number: u64) {
278	/// This method generates some vrf data, but only to make the compiler happy
279	fn generate_testing_vrf() -> VrfSignature {
280		let vrf_proof_bytes = [0u8; 64];
281		let proof: VrfProof = VrfProof::decode(&mut vrf_proof_bytes.as_slice()).unwrap();
282		let vrf_pre_out_bytes = [0u8; 32];
283		let pre_output: VrfPreOutput =
284			VrfPreOutput::decode(&mut vrf_pre_out_bytes.as_slice()).unwrap();
285		VrfSignature { pre_output, proof }
286	}
287
288	let pre_digest = PrimaryPreDigest {
289		authority_index,
290		slot: block_number.into(),
291		vrf_signature: generate_testing_vrf(),
292	};
293
294	header
295		.digest_mut()
296		.push(DigestItem::babe_pre_digest(PreDigest::Primary(pre_digest)));
297}
298
299/// Create a mock chain of relay headers as descendants of the relay parent
300pub fn build_relay_parent_descendants(
301	num_headers: u64,
302	state_root: H256,
303	authorities: Vec<AuthorityPair>,
304) -> Vec<Header> {
305	let mut headers = Vec::with_capacity(num_headers as usize);
306
307	let mut previous_hash = None;
308
309	for block_number in 0..=num_headers as u32 {
310		let mut header = Header {
311			number: block_number,
312			parent_hash: previous_hash.unwrap_or_default(),
313			state_root,
314			extrinsics_root: H256::default(),
315			digest: Digest::default(),
316		};
317		let authority_index = block_number % (authorities.len() as u32);
318
319		// Add pre-digest
320		add_babe_pre_digest(&mut header, authority_index, block_number as u64);
321
322		// Sign and seal the header
323		let signature = authorities[authority_index as usize].sign(header.hash().as_bytes());
324		header.digest_mut().push(DigestItem::babe_seal(signature.into()));
325
326		previous_hash = Some(header.hash());
327		headers.push(header);
328	}
329
330	headers
331}
cumulus_test_relay_sproof_builder/lib.rs

cumulus_test_relay_sproof_builder/
lib.rs
