|
17 | 17 | extern crate alloc; |
18 | 18 |
|
19 | 19 | use alloc::collections::btree_map::BTreeMap; |
| 20 | +use codec::{Decode, Encode}; |
20 | 21 | use cumulus_primitives_core::{ |
21 | 22 | relay_chain, AbridgedHostConfiguration, AbridgedHrmpChannel, ParaId, |
22 | 23 | }; |
23 | | -use polkadot_primitives::UpgradeGoAhead; |
24 | | -use sp_runtime::traits::HashingFor; |
| 24 | +use polkadot_primitives::{Header, UpgradeGoAhead}; |
| 25 | +use sp_consensus_babe::{ |
| 26 | + digests::{CompatibleDigestItem, PreDigest, PrimaryPreDigest}, |
| 27 | + AuthorityId, AuthorityPair, BabeAuthorityWeight, |
| 28 | +}; |
| 29 | +use sp_core::{ |
| 30 | + sr25519::vrf::{VrfPreOutput, VrfProof, VrfSignature}, |
| 31 | + Pair, H256, |
| 32 | +}; |
| 33 | +use sp_runtime::{ |
| 34 | + traits::{HashingFor, Header as HeaderT}, |
| 35 | + Digest, DigestItem, |
| 36 | +}; |
25 | 37 | use sp_trie::PrefixedMemoryDB; |
26 | 38 |
|
27 | 39 | /// Builds a sproof (portmanteau of 'spoof' and 'proof') of the relay chain state. |
@@ -49,6 +61,7 @@ pub struct RelayStateSproofBuilder { |
49 | 61 | pub randomness: relay_chain::Hash, |
50 | 62 | pub additional_key_values: Vec<(Vec<u8>, Vec<u8>)>, |
51 | 63 | pub included_para_head: Option<relay_chain::HeadData>, |
| 64 | + pub num_authorities: u64, |
52 | 65 | } |
53 | 66 |
|
54 | 67 | impl Default for RelayStateSproofBuilder { |
@@ -81,6 +94,7 @@ impl Default for RelayStateSproofBuilder { |
81 | 94 | randomness: relay_chain::Hash::default(), |
82 | 95 | additional_key_values: vec![], |
83 | 96 | included_para_head: None, |
| 97 | + num_authorities: 1, |
84 | 98 | } |
85 | 99 | } |
86 | 100 | } |
@@ -129,9 +143,39 @@ impl RelayStateSproofBuilder { |
129 | 143 | }) |
130 | 144 | } |
131 | 145 |
|
132 | | - pub fn into_state_root_and_proof( |
| 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( |
133 | 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, |
134 | 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 | + |
135 | 179 | let (db, root) = |
136 | 180 | PrefixedMemoryDB::<HashingFor<polkadot_primitives::Block>>::default_with_root(); |
137 | 181 | let state_version = Default::default(); // for test using default. |
@@ -213,3 +257,75 @@ impl RelayStateSproofBuilder { |
213 | 257 | (root, proof) |
214 | 258 | } |
215 | 259 | } |
| 260 | + |
| 261 | +/// Generate a vector of AuthorityPairs |
| 262 | +pub 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) |
| 267 | +fn 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 |
| 277 | +fn 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 |
| 300 | +pub 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 | +} |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?v=4&size=40){kind=avatar}
0 commit comments