zombienet_orchestrator/generators/
chain_spec.rs

1use std::{
2    collections::HashMap,
3    path::{Path, PathBuf},
4};
5
6use anyhow::anyhow;
7use configuration::{
8    types::{AssetLocation, JsonOverrides},
9    HrmpChannelConfig,
10};
11use provider::{
12    constants::NODE_CONFIG_DIR,
13    types::{GenerateFileCommand, GenerateFilesOptions, TransferedFile},
14    DynNamespace, ProviderError,
15};
16use serde::{Deserialize, Serialize};
17use serde_json::json;
18use support::{constants::THIS_IS_A_BUG, fs::FileSystem, replacer::apply_replacements};
19use tokio::process::Command;
20use tracing::{debug, info, trace, warn};
21
22use super::errors::GeneratorError;
23use crate::{
24    network_spec::{node::NodeSpec, parachain::ParachainSpec, relaychain::RelaychainSpec},
25    ScopedFilesystem,
26};
27
28// TODO: (javier) move to state
29#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
30pub enum Context {
31    Relay,
32    Para,
33}
34
35enum ChainSpecFormat {
36    Plain,
37    Raw,
38}
39
40enum KeyType {
41    Session,
42    Aura,
43    Grandpa,
44}
45
46#[derive(Debug, Clone, Copy)]
47enum SessionKeyType {
48    Default,
49    Stash,
50    Evm,
51}
52
53impl Default for SessionKeyType {
54    fn default() -> Self {
55        Self::Default
56    }
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
60pub enum CommandInContext {
61    Local(String),
62    Remote(String),
63}
64
65impl CommandInContext {
66    fn cmd(&self) -> &str {
67        match self {
68            CommandInContext::Local(cmd) | CommandInContext::Remote(cmd) => cmd.as_ref(),
69        }
70    }
71}
72
73#[derive(Debug)]
74pub struct ParaGenesisConfig<T: AsRef<Path>> {
75    pub(crate) state_path: T,
76    pub(crate) wasm_path: T,
77    pub(crate) id: u32,
78    pub(crate) as_parachain: bool,
79}
80
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct ChainSpec {
83    // Name of the spec file, most of the times could be the same as the chain_name. (e.g rococo-local)
84    chain_spec_name: String,
85    asset_location: Option<AssetLocation>,
86    maybe_plain_path: Option<PathBuf>,
87    chain_name: Option<String>,
88    raw_path: Option<PathBuf>,
89    // The binary to build the chain-spec
90    command: Option<CommandInContext>,
91    // Imgae to use for build the chain-spec
92    image: Option<String>,
93    // Contex of the network (e.g relay or para)
94    context: Context,
95}
96
97impl ChainSpec {
98    pub(crate) fn new(chain_spec_name: impl Into<String>, context: Context) -> Self {
99        Self {
100            chain_spec_name: chain_spec_name.into(),
101            chain_name: None,
102            maybe_plain_path: None,
103            asset_location: None,
104            raw_path: None,
105            command: None,
106            image: None,
107            context,
108        }
109    }
110
111    pub(crate) fn chain_spec_name(&self) -> &str {
112        self.chain_spec_name.as_ref()
113    }
114
115    pub(crate) fn chain_name(&self) -> Option<&str> {
116        self.chain_name.as_deref()
117    }
118
119    pub(crate) fn set_chain_name(mut self, chain_name: impl Into<String>) -> Self {
120        self.chain_name = Some(chain_name.into());
121        self
122    }
123
124    pub(crate) fn asset_location(mut self, location: AssetLocation) -> Self {
125        self.asset_location = Some(location);
126        self
127    }
128
129    pub(crate) fn command(mut self, command: impl Into<String>, is_local: bool) -> Self {
130        let cmd = if is_local {
131            CommandInContext::Local(command.into())
132        } else {
133            CommandInContext::Remote(command.into())
134        };
135        self.command = Some(cmd);
136        self
137    }
138
139    pub(crate) fn image(mut self, image: Option<String>) -> Self {
140        self.image = image;
141        self
142    }
143
144    /// Build the chain-spec
145    pub async fn build<'a, T>(
146        &mut self,
147        ns: &DynNamespace,
148        scoped_fs: &ScopedFilesystem<'a, T>,
149    ) -> Result<(), GeneratorError>
150    where
151        T: FileSystem,
152    {
153        // TODO: Move this to state builder.
154        if self.asset_location.is_none() && self.command.is_none() {
155            return Err(GeneratorError::ChainSpecGeneration(
156                "Can not build the chain spec without set the command or asset_location"
157                    .to_string(),
158            ));
159        }
160
161        let maybe_plain_spec_path = PathBuf::from(format!("{}-plain.json", self.chain_spec_name));
162        // if we have a path, copy to the base_dir of the ns with the name `<name>-plain.json`
163        if let Some(location) = self.asset_location.as_ref() {
164            match location {
165                AssetLocation::FilePath(path) => {
166                    let file_to_transfer =
167                        TransferedFile::new(path.clone(), maybe_plain_spec_path.clone());
168
169                    scoped_fs
170                        .copy_files(vec![&file_to_transfer])
171                        .await
172                        .map_err(|_| {
173                            GeneratorError::ChainSpecGeneration(format!(
174                                "Error copying file: {file_to_transfer}"
175                            ))
176                        })?;
177                },
178                AssetLocation::Url(url) => {
179                    let res = reqwest::get(url.as_str())
180                        .await
181                        .map_err(|err| ProviderError::DownloadFile(url.to_string(), err.into()))?;
182
183                    let contents: &[u8] = &res.bytes().await.unwrap();
184                    trace!(
185                        "writing content from {} to: {maybe_plain_spec_path:?}",
186                        url.as_str()
187                    );
188                    scoped_fs.write(&maybe_plain_spec_path, contents).await?;
189                },
190            }
191        } else {
192            // we should create the chain-spec using command.
193            let mut replacement_value = String::default();
194            if let Some(chain_name) = self.chain_name.as_ref() {
195                if !chain_name.is_empty() {
196                    replacement_value.clone_from(chain_name);
197                }
198            };
199
200            // SAFETY: we ensure that command is some with the first check of the fn
201            // default as empty
202            let sanitized_cmd = if replacement_value.is_empty() {
203                // we need to remove the `--chain` flag
204                self.command.as_ref().unwrap().cmd().replace("--chain", "")
205            } else {
206                self.command.as_ref().unwrap().cmd().to_owned()
207            };
208
209            let full_cmd = apply_replacements(
210                &sanitized_cmd,
211                &HashMap::from([("chainName", replacement_value.as_str())]),
212            );
213            trace!("full_cmd: {:?}", full_cmd);
214
215            let parts: Vec<&str> = full_cmd.split_whitespace().collect();
216            let Some((cmd, args)) = parts.split_first() else {
217                return Err(GeneratorError::ChainSpecGeneration(format!(
218                    "Invalid generator command: {full_cmd}"
219                )));
220            };
221            trace!("cmd: {:?} - args: {:?}", cmd, args);
222
223            let generate_command =
224                GenerateFileCommand::new(cmd, maybe_plain_spec_path.clone()).args(args);
225            if let Some(CommandInContext::Local(_)) = self.command {
226                // local
227                build_locally(generate_command, scoped_fs).await?;
228            } else {
229                // remote
230                let options = GenerateFilesOptions::new(vec![generate_command], self.image.clone());
231                ns.generate_files(options).await?;
232            }
233        }
234
235        if is_raw(maybe_plain_spec_path.clone(), scoped_fs).await? {
236            let spec_path = PathBuf::from(format!("{}.json", self.chain_spec_name));
237            let tf_file = TransferedFile::new(
238                &PathBuf::from_iter([ns.base_dir(), &maybe_plain_spec_path]),
239                &spec_path,
240            );
241            scoped_fs.copy_files(vec![&tf_file]).await.map_err(|e| {
242                GeneratorError::ChainSpecGeneration(format!(
243                    "Error copying file: {tf_file}, err: {e}"
244                ))
245            })?;
246
247            self.raw_path = Some(spec_path);
248        } else {
249            self.maybe_plain_path = Some(maybe_plain_spec_path);
250        }
251        Ok(())
252    }
253
254    pub async fn build_raw<'a, T>(
255        &mut self,
256        ns: &DynNamespace,
257        scoped_fs: &ScopedFilesystem<'a, T>,
258    ) -> Result<(), GeneratorError>
259    where
260        T: FileSystem,
261    {
262        let None = self.raw_path else {
263            return Ok(());
264        };
265        // build raw
266        let temp_name = format!(
267            "temp-build-raw-{}-{}",
268            self.chain_spec_name,
269            rand::random::<u8>()
270        );
271        let raw_spec_path = PathBuf::from(format!("{}.json", self.chain_spec_name));
272        let cmd = self
273            .command
274            .as_ref()
275            .ok_or(GeneratorError::ChainSpecGeneration(
276                "Invalid command".into(),
277            ))?;
278        let maybe_plain_path =
279            self.maybe_plain_path
280                .as_ref()
281                .ok_or(GeneratorError::ChainSpecGeneration(
282                    "Invalid plain path".into(),
283                ))?;
284
285        // TODO: we should get the full path from the scoped filesystem
286        let chain_spec_path_local = format!(
287            "{}/{}",
288            ns.base_dir().to_string_lossy(),
289            maybe_plain_path.display()
290        );
291        // Remote path to be injected
292        let chain_spec_path_in_pod = format!("{}/{}", NODE_CONFIG_DIR, maybe_plain_path.display());
293        // Path in the context of the node, this can be different in the context of the providers (e.g native)
294        let chain_spec_path_in_args = if matches!(self.command, Some(CommandInContext::Local(_))) {
295            chain_spec_path_local.clone()
296        } else if ns.capabilities().prefix_with_full_path {
297            // In native
298            format!(
299                "{}/{}{}",
300                ns.base_dir().to_string_lossy(),
301                &temp_name,
302                &chain_spec_path_in_pod
303            )
304        } else {
305            chain_spec_path_in_pod.clone()
306        };
307
308        let mut full_cmd = apply_replacements(
309            cmd.cmd(),
310            &HashMap::from([("chainName", chain_spec_path_in_args.as_str())]),
311        );
312
313        if !full_cmd.contains("--raw") {
314            full_cmd = format!("{full_cmd} --raw");
315        }
316        trace!("full_cmd: {:?}", full_cmd);
317
318        let parts: Vec<&str> = full_cmd.split_whitespace().collect();
319        let Some((cmd, args)) = parts.split_first() else {
320            return Err(GeneratorError::ChainSpecGeneration(format!(
321                "Invalid generator command: {full_cmd}"
322            )));
323        };
324        trace!("cmd: {:?} - args: {:?}", cmd, args);
325
326        let generate_command = GenerateFileCommand::new(cmd, raw_spec_path.clone()).args(args);
327
328        if let Some(CommandInContext::Local(_)) = self.command {
329            // local
330            build_locally(generate_command, scoped_fs).await?;
331        } else {
332            // remote
333            let options = GenerateFilesOptions::with_files(
334                vec![generate_command],
335                self.image.clone(),
336                &[TransferedFile::new(
337                    chain_spec_path_local,
338                    chain_spec_path_in_pod,
339                )],
340            )
341            .temp_name(temp_name);
342            trace!("calling generate_files with options: {:#?}", options);
343            ns.generate_files(options).await?;
344        }
345
346        self.raw_path = Some(raw_spec_path);
347
348        Ok(())
349    }
350
351    /// Override the :code in chain-spec raw version
352    pub async fn override_code<'a, T>(
353        &mut self,
354        scoped_fs: &ScopedFilesystem<'a, T>,
355        wasm_override: &AssetLocation,
356    ) -> Result<(), GeneratorError>
357    where
358        T: FileSystem,
359    {
360        // first ensure we have the raw version of the chain-spec
361        let Some(_) = self.raw_path else {
362            return Err(GeneratorError::OverridingWasm(String::from(
363                "Raw path should be set at this point.",
364            )));
365        };
366        let (content, _) = self.read_spec(scoped_fs).await?;
367        // read override wasm
368        let override_content = wasm_override.get_asset().await.map_err(|_| {
369            GeneratorError::OverridingWasm(format!(
370                "Can not get asset to override wasm, asset: {wasm_override}"
371            ))
372        })?;
373
374        // read spec  to json value
375        let mut chain_spec_json: serde_json::Value =
376            serde_json::from_str(&content).map_err(|_| {
377                GeneratorError::ChainSpecGeneration("Can not parse chain-spec as json".into())
378            })?;
379
380        // override :code
381        let Some(code) = chain_spec_json.pointer_mut("/genesis/raw/top/0x3a636f6465") else {
382            return Err(GeneratorError::OverridingWasm(String::from(
383                "Pointer '/genesis/raw/top/0x3a636f6465' should be valid in the raw spec.",
384            )));
385        };
386
387        info!(
388            "🖋  Overriding ':code' (0x3a636f6465) in raw chain-spec with content of {}",
389            wasm_override
390        );
391        *code = json!(format!("0x{}", hex::encode(override_content)));
392
393        let overrided_content = serde_json::to_string_pretty(&chain_spec_json).map_err(|_| {
394            GeneratorError::ChainSpecGeneration("can not parse chain-spec value as json".into())
395        })?;
396        // save it
397        self.write_spec(scoped_fs, overrided_content).await?;
398
399        Ok(())
400    }
401
402    pub async fn override_raw_spec<'a, T>(
403        &mut self,
404        scoped_fs: &ScopedFilesystem<'a, T>,
405        raw_spec_overrides: &JsonOverrides,
406    ) -> Result<(), GeneratorError>
407    where
408        T: FileSystem,
409    {
410        // first ensure we have the raw version of the chain-spec
411        let Some(_) = self.raw_path else {
412            return Err(GeneratorError::OverridingRawSpec(String::from(
413                "Raw path should be set at this point.",
414            )));
415        };
416
417        let (content, _) = self.read_spec(scoped_fs).await?;
418
419        // read overrides to json value
420        let override_content: serde_json::Value = raw_spec_overrides.get().await.map_err(|_| {
421            GeneratorError::OverridingRawSpec(format!(
422                "Can not parse raw_spec_override contents as json: {raw_spec_overrides}"
423            ))
424        })?;
425
426        // read spec to json value
427        let mut chain_spec_json: serde_json::Value =
428            serde_json::from_str(&content).map_err(|_| {
429                GeneratorError::ChainSpecGeneration("Can not parse chain-spec as json".into())
430            })?;
431
432        // merge overrides with existing spec
433        merge(&mut chain_spec_json, &override_content);
434
435        // save changes
436        let overrided_content = serde_json::to_string_pretty(&chain_spec_json).map_err(|_| {
437            GeneratorError::ChainSpecGeneration("can not parse chain-spec value as json".into())
438        })?;
439        self.write_spec(scoped_fs, overrided_content).await?;
440
441        Ok(())
442    }
443
444    pub fn raw_path(&self) -> Option<&Path> {
445        self.raw_path.as_deref()
446    }
447
448    pub fn set_asset_location(&mut self, location: AssetLocation) {
449        self.asset_location = Some(location)
450    }
451
452    pub async fn read_chain_id<'a, T>(
453        &self,
454        scoped_fs: &ScopedFilesystem<'a, T>,
455    ) -> Result<String, GeneratorError>
456    where
457        T: FileSystem,
458    {
459        let (content, _) = self.read_spec(scoped_fs).await?;
460        ChainSpec::chain_id_from_spec(&content)
461    }
462
463    async fn read_spec<'a, T>(
464        &self,
465        scoped_fs: &ScopedFilesystem<'a, T>,
466    ) -> Result<(String, ChainSpecFormat), GeneratorError>
467    where
468        T: FileSystem,
469    {
470        let (path, format) = match (self.maybe_plain_path.as_ref(), self.raw_path.as_ref()) {
471            (Some(path), None) => (path, ChainSpecFormat::Plain),
472            (None, Some(path)) => (path, ChainSpecFormat::Raw),
473            (Some(_), Some(path)) => {
474                // if we have both paths return the raw
475                (path, ChainSpecFormat::Raw)
476            },
477            (None, None) => unreachable!(),
478        };
479
480        let content = scoped_fs.read_to_string(path.clone()).await.map_err(|_| {
481            GeneratorError::ChainSpecGeneration(format!(
482                "Can not read chain-spec from {}",
483                path.to_string_lossy()
484            ))
485        })?;
486
487        Ok((content, format))
488    }
489
490    async fn write_spec<'a, T>(
491        &self,
492        scoped_fs: &ScopedFilesystem<'a, T>,
493        content: impl Into<String>,
494    ) -> Result<(), GeneratorError>
495    where
496        T: FileSystem,
497    {
498        let (path, _format) = match (self.maybe_plain_path.as_ref(), self.raw_path.as_ref()) {
499            (Some(path), None) => (path, ChainSpecFormat::Plain),
500            (None, Some(path)) => (path, ChainSpecFormat::Raw),
501            (Some(_), Some(path)) => {
502                // if we have both paths return the raw
503                (path, ChainSpecFormat::Raw)
504            },
505            (None, None) => unreachable!(),
506        };
507
508        scoped_fs.write(path, content.into()).await.map_err(|_| {
509            GeneratorError::ChainSpecGeneration(format!(
510                "Can not write chain-spec from {}",
511                path.to_string_lossy()
512            ))
513        })?;
514
515        Ok(())
516    }
517
518    // TODO: (javier) move this fns to state aware
519    pub async fn customize_para<'a, T>(
520        &self,
521        para: &ParachainSpec,
522        relay_chain_id: &str,
523        scoped_fs: &ScopedFilesystem<'a, T>,
524    ) -> Result<(), GeneratorError>
525    where
526        T: FileSystem,
527    {
528        let (content, format) = self.read_spec(scoped_fs).await?;
529        let mut chain_spec_json: serde_json::Value =
530            serde_json::from_str(&content).map_err(|_| {
531                GeneratorError::ChainSpecGeneration("Can not parse chain-spec as json".into())
532            })?;
533
534        if let Some(para_id) = chain_spec_json.get_mut("para_id") {
535            *para_id = json!(para.id);
536        };
537        if let Some(para_id) = chain_spec_json.get_mut("paraId") {
538            *para_id = json!(para.id);
539        };
540
541        if let Some(relay_chain_id_field) = chain_spec_json.get_mut("relay_chain") {
542            *relay_chain_id_field = json!(relay_chain_id);
543        };
544
545        if let ChainSpecFormat::Plain = format {
546            let pointer = get_runtime_config_pointer(&chain_spec_json)
547                .map_err(GeneratorError::ChainSpecGeneration)?;
548
549            // make genesis overrides first.
550            if let Some(overrides) = &para.genesis_overrides {
551                let percolated_overrides = percolate_overrides(&pointer, overrides)
552                    .map_err(|e| GeneratorError::ChainSpecGeneration(e.to_string()))?;
553                if let Some(genesis) = chain_spec_json.pointer_mut(&pointer) {
554                    merge(genesis, percolated_overrides);
555                }
556            }
557
558            clear_authorities(&pointer, &mut chain_spec_json, &self.context);
559
560            let key_type_to_use = if para.is_evm_based {
561                SessionKeyType::Evm
562            } else {
563                SessionKeyType::Default
564            };
565
566            // Get validators to add as authorities
567            let validators: Vec<&NodeSpec> = para
568                .collators
569                .iter()
570                .filter(|node| node.is_validator)
571                .collect();
572
573            // check chain key types
574            if chain_spec_json
575                .pointer(&format!("{pointer}/session"))
576                .is_some()
577            {
578                add_authorities(&pointer, &mut chain_spec_json, &validators, key_type_to_use);
579            } else if chain_spec_json
580                .pointer(&format!("{pointer}/aura"))
581                .is_some()
582            {
583                add_aura_authorities(&pointer, &mut chain_spec_json, &validators, KeyType::Aura);
584            } else {
585                warn!("Can't customize keys, not `session` or `aura` find in the chain-spec file");
586            };
587
588            // Add nodes to collator
589            let invulnerables: Vec<&NodeSpec> = para
590                .collators
591                .iter()
592                .filter(|node| node.is_invulnerable)
593                .collect();
594
595            add_collator_selection(
596                &pointer,
597                &mut chain_spec_json,
598                &invulnerables,
599                key_type_to_use,
600            );
601
602            // override `parachainInfo/parachainId`
603            override_parachain_info(&pointer, &mut chain_spec_json, para.id);
604
605            // write spec
606            let content = serde_json::to_string_pretty(&chain_spec_json).map_err(|_| {
607                GeneratorError::ChainSpecGeneration("can not parse chain-spec value as json".into())
608            })?;
609            self.write_spec(scoped_fs, content).await?;
610        } else {
611            warn!("⚠️ Chain spec for para_id: {} is in raw mode", para.id);
612        }
613        Ok(())
614    }
615
616    pub async fn customize_relay<'a, T, U>(
617        &self,
618        relaychain: &RelaychainSpec,
619        hrmp_channels: &[HrmpChannelConfig],
620        para_artifacts: Vec<ParaGenesisConfig<U>>,
621        scoped_fs: &ScopedFilesystem<'a, T>,
622    ) -> Result<(), GeneratorError>
623    where
624        T: FileSystem,
625        U: AsRef<Path>,
626    {
627        let (content, format) = self.read_spec(scoped_fs).await?;
628        let mut chain_spec_json: serde_json::Value =
629            serde_json::from_str(&content).map_err(|_| {
630                GeneratorError::ChainSpecGeneration("Can not parse chain-spec as json".into())
631            })?;
632
633        if let ChainSpecFormat::Plain = format {
634            // get the tokenDecimals property or set the default (12)
635            let token_decimals =
636                if let Some(val) = chain_spec_json.pointer("/properties/tokenDecimals") {
637                    let val = val.as_u64().unwrap_or(12);
638                    if val > u8::MAX as u64 {
639                        12
640                    } else {
641                        val as u8
642                    }
643                } else {
644                    12
645                };
646            // get the config pointer
647            let pointer = get_runtime_config_pointer(&chain_spec_json)
648                .map_err(GeneratorError::ChainSpecGeneration)?;
649
650            // make genesis overrides first.
651            if let Some(overrides) = &relaychain.runtime_genesis_patch {
652                let percolated_overrides = percolate_overrides(&pointer, overrides)
653                    .map_err(|e| GeneratorError::ChainSpecGeneration(e.to_string()))?;
654                if let Some(patch_section) = chain_spec_json.pointer_mut(&pointer) {
655                    merge(patch_section, percolated_overrides);
656                }
657            }
658
659            // get min stake (to store if neede later)
660            let staking_min = get_staking_min(&pointer, &mut chain_spec_json);
661
662            // Clear authorities
663            clear_authorities(&pointer, &mut chain_spec_json, &self.context);
664
665            // add balances
666            add_balances(
667                &pointer,
668                &mut chain_spec_json,
669                &relaychain.nodes,
670                token_decimals,
671                staking_min,
672            );
673
674            // add staking
675            add_staking(
676                &pointer,
677                &mut chain_spec_json,
678                &relaychain.nodes,
679                staking_min,
680            );
681
682            // Get validators to add as authorities
683            let validators: Vec<&NodeSpec> = relaychain
684                .nodes
685                .iter()
686                .filter(|node| node.is_validator)
687                .collect();
688
689            // check chain key types
690            if chain_spec_json
691                .pointer(&format!("{pointer}/session"))
692                .is_some()
693            {
694                add_authorities(
695                    &pointer,
696                    &mut chain_spec_json,
697                    &validators,
698                    SessionKeyType::Stash,
699                );
700            } else {
701                add_aura_authorities(&pointer, &mut chain_spec_json, &validators, KeyType::Aura);
702                add_grandpa_authorities(&pointer, &mut chain_spec_json, &validators, KeyType::Aura);
703            }
704
705            // staking && nominators
706
707            if !hrmp_channels.is_empty() {
708                add_hrmp_channels(&pointer, &mut chain_spec_json, hrmp_channels);
709            }
710
711            // paras
712            for para_genesis_config in para_artifacts.iter() {
713                add_parachain_to_genesis(
714                    &pointer,
715                    &mut chain_spec_json,
716                    para_genesis_config,
717                    scoped_fs,
718                )
719                .await
720                .map_err(|e| GeneratorError::ChainSpecGeneration(e.to_string()))?;
721            }
722
723            // TODO:
724            // - staking
725            // - nominators
726
727            // write spec
728            let content = serde_json::to_string_pretty(&chain_spec_json).map_err(|_| {
729                GeneratorError::ChainSpecGeneration("can not parse chain-spec value as json".into())
730            })?;
731            self.write_spec(scoped_fs, content).await?;
732        } else {
733            warn!(
734                "⚠️ Chain Spec for chain {} is in raw mode, can't customize.",
735                self.chain_spec_name
736            );
737        }
738        Ok(())
739    }
740
741    pub async fn add_bootnodes<'a, T>(
742        &self,
743        scoped_fs: &ScopedFilesystem<'a, T>,
744        bootnodes: &[String],
745    ) -> Result<(), GeneratorError>
746    where
747        T: FileSystem,
748    {
749        let (content, _) = self.read_spec(scoped_fs).await?;
750        let mut chain_spec_json: serde_json::Value =
751            serde_json::from_str(&content).map_err(|_| {
752                GeneratorError::ChainSpecGeneration("Can not parse chain-spec as json".into())
753            })?;
754
755        if let Some(bootnodes_on_file) = chain_spec_json.get_mut("bootNodes") {
756            if let Some(bootnodes_on_file) = bootnodes_on_file.as_array_mut() {
757                let mut bootnodes_to_add =
758                    bootnodes.iter().map(|bootnode| json!(bootnode)).collect();
759                bootnodes_on_file.append(&mut bootnodes_to_add);
760            } else {
761                return Err(GeneratorError::ChainSpecGeneration(
762                    "id should be an string in the chain-spec, this is a bug".into(),
763                ));
764            };
765        } else {
766            return Err(GeneratorError::ChainSpecGeneration(
767                "'bootNodes' should be a fields in the chain-spec of the relaychain".into(),
768            ));
769        };
770
771        // write spec
772        let content = serde_json::to_string_pretty(&chain_spec_json).map_err(|_| {
773            GeneratorError::ChainSpecGeneration("can not parse chain-spec value as json".into())
774        })?;
775        self.write_spec(scoped_fs, content).await?;
776
777        Ok(())
778    }
779
780    /// Get the chain_is from the json content of a chain-spec file.
781    pub fn chain_id_from_spec(spec_content: &str) -> Result<String, GeneratorError> {
782        let chain_spec_json: serde_json::Value =
783            serde_json::from_str(spec_content).map_err(|_| {
784                GeneratorError::ChainSpecGeneration("Can not parse chain-spec as json".into())
785            })?;
786        if let Some(chain_id) = chain_spec_json.get("id") {
787            if let Some(chain_id) = chain_id.as_str() {
788                Ok(chain_id.to_string())
789            } else {
790                Err(GeneratorError::ChainSpecGeneration(
791                    "id should be an string in the chain-spec, this is a bug".into(),
792                ))
793            }
794        } else {
795            Err(GeneratorError::ChainSpecGeneration(
796                "'id' should be a fields in the chain-spec of the relaychain".into(),
797            ))
798        }
799    }
800}
801
802type GenesisNodeKey = (String, String, HashMap<String, String>);
803
804async fn build_locally<'a, T>(
805    generate_command: GenerateFileCommand,
806    scoped_fs: &ScopedFilesystem<'a, T>,
807) -> Result<(), GeneratorError>
808where
809    T: FileSystem,
810{
811    // generate_command.
812
813    let result = Command::new(generate_command.program.clone())
814        .args(generate_command.args.clone())
815        .current_dir(scoped_fs.base_dir)
816        .output()
817        .await
818        .map_err(|err| {
819            GeneratorError::ChainSpecGeneration(format!(
820                "Error running cmd: {} args: {}, err: {}",
821                &generate_command.program,
822                &generate_command.args.join(" "),
823                err
824            ))
825        })?;
826
827    if result.status.success() {
828        scoped_fs
829            .write(
830                generate_command.local_output_path,
831                String::from_utf8_lossy(&result.stdout).to_string(),
832            )
833            .await?;
834        Ok(())
835    } else {
836        Err(GeneratorError::ChainSpecGeneration(format!(
837            "Error running cmd: {} args: {}, err: {}",
838            &generate_command.program,
839            &generate_command.args.join(" "),
840            String::from_utf8_lossy(&result.stderr)
841        )))
842    }
843}
844
845async fn is_raw<'a, T>(
846    file: PathBuf,
847    scoped_fs: &ScopedFilesystem<'a, T>,
848) -> Result<bool, ProviderError>
849where
850    T: FileSystem,
851{
852    let content = scoped_fs.read_to_string(file).await?;
853    let chain_spec_json: serde_json::Value = serde_json::from_str(&content).unwrap();
854
855    Ok(chain_spec_json.pointer("/genesis/raw/top").is_some())
856}
857
858// Internal Chain-spec customizations
859
860async fn add_parachain_to_genesis<'a, T, U>(
861    runtime_config_ptr: &str,
862    chain_spec_json: &mut serde_json::Value,
863    para_genesis_config: &ParaGenesisConfig<U>,
864    scoped_fs: &ScopedFilesystem<'a, T>,
865) -> Result<(), anyhow::Error>
866where
867    T: FileSystem,
868    U: AsRef<Path>,
869{
870    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
871        let paras_pointer = if val.get("paras").is_some() {
872            "/paras/paras"
873        } else if val.get("parachainsParas").is_some() {
874            // For retro-compatibility with substrate pre Polkadot 0.9.5
875            "/parachainsParas/paras"
876        } else {
877            // The config may not contain paras. Since chainspec allows to contain the RuntimeGenesisConfig patch we can inject it.
878            val["paras"] = json!({ "paras": [] });
879            "/paras/paras"
880        };
881
882        let paras = val.pointer_mut(paras_pointer).ok_or(anyhow!(
883            "paras pointer should be valid {:?} ",
884            paras_pointer
885        ))?;
886        let paras_vec = paras
887            .as_array_mut()
888            .ok_or(anyhow!("paras should be an array"))?;
889
890        let head = scoped_fs
891            .read_to_string(para_genesis_config.state_path.as_ref())
892            .await?;
893        let wasm = scoped_fs
894            .read_to_string(para_genesis_config.wasm_path.as_ref())
895            .await?;
896
897        paras_vec.push(json!([
898            para_genesis_config.id,
899            [head.trim(), wasm.trim(), para_genesis_config.as_parachain]
900        ]));
901
902        Ok(())
903    } else {
904        unreachable!("pointer to runtime config should be valid!")
905    }
906}
907
908fn get_runtime_config_pointer(chain_spec_json: &serde_json::Value) -> Result<String, String> {
909    // runtime_genesis_config is no longer in ChainSpec after rococo runtime rework (refer to: https://github.com/paritytech/polkadot-sdk/pull/1256)
910    // ChainSpec may contain a RuntimeGenesisConfigPatch
911    let pointers = [
912        "/genesis/runtimeGenesis/config",
913        "/genesis/runtimeGenesis/patch",
914        "/genesis/runtimeGenesisConfigPatch",
915        "/genesis/runtime/runtime_genesis_config",
916        "/genesis/runtime",
917    ];
918
919    for pointer in pointers {
920        if chain_spec_json.pointer(pointer).is_some() {
921            return Ok(pointer.to_string());
922        }
923    }
924
925    Err("Can not find the runtime pointer".into())
926}
927
928fn percolate_overrides<'a>(
929    pointer: &str,
930    overrides: &'a serde_json::Value,
931) -> Result<&'a serde_json::Value, anyhow::Error> {
932    let pointer_parts = pointer.split('/').collect::<Vec<&str>>();
933    trace!("pointer_parts: {pointer_parts:?}");
934
935    let top_level = overrides
936        .as_object()
937        .ok_or_else(|| anyhow!("Overrides must be an object"))?;
938    let top_level_key = top_level
939        .keys()
940        .next()
941        .ok_or_else(|| anyhow!("Invalid override value: {:?}", overrides))?;
942    trace!("top_level_key: {top_level_key}");
943    let index = pointer_parts.iter().position(|x| *x == top_level_key);
944    let Some(i) = index else {
945        warn!("Top level key '{top_level_key}' isn't part of the pointer ({pointer}), returning without percolating");
946        return Ok(overrides);
947    };
948
949    let p = if i == pointer_parts.len() - 1 {
950        // top level key is at end of the pointer
951        let p = format!("/{}", pointer_parts[i]);
952        trace!("overrides pointer {p}");
953        p
954    } else {
955        // example: pointer is `/genesis/runtimeGenesis/patch` and the overrides start at  `runtimeGenesis`
956        let p = format!("/{}", pointer_parts[i..].join("/"));
957        trace!("overrides pointer {p}");
958        p
959    };
960    let overrides_to_use = overrides
961        .pointer(&p)
962        .ok_or_else(|| anyhow!("Invalid override value: {:?}", overrides))?;
963    Ok(overrides_to_use)
964}
965
966#[allow(dead_code)]
967fn construct_runtime_pointer_from_overrides(
968    overrides: &serde_json::Value,
969) -> Result<String, anyhow::Error> {
970    if overrides.get("genesis").is_some() {
971        // overrides already start with /genesis
972        return Ok("/genesis".into());
973    } else {
974        // check if we are one level inner
975        if let Some(top_level) = overrides.as_object() {
976            let k = top_level
977                .keys()
978                .next()
979                .ok_or_else(|| anyhow!("Invalid override value: {:?}", overrides))?;
980            match k.as_str() {
981                "runtimeGenesisConfigPatch" | "runtime" | "runtimeGenesis" => {
982                    return Ok(("/genesis").into())
983                },
984                "config" | "path" => {
985                    return Ok(("/genesis/runtimeGenesis").into());
986                },
987                "runtime_genesis_config" => {
988                    return Ok(("/genesis/runtime").into());
989                },
990                _ => {},
991            }
992        }
993    }
994
995    Err(anyhow!("Can not find the runtime pointer"))
996}
997
998// Merge `patch_section` with `overrides`.
999fn merge(patch_section: &mut serde_json::Value, overrides: &serde_json::Value) {
1000    trace!("patch: {:?}", patch_section);
1001    trace!("overrides: {:?}", overrides);
1002    if let (Some(genesis_obj), Some(overrides_obj)) =
1003        (patch_section.as_object_mut(), overrides.as_object())
1004    {
1005        for overrides_key in overrides_obj.keys() {
1006            trace!("overrides_key: {:?}", overrides_key);
1007            // we only want to override keys present in the genesis object
1008            if let Some(genesis_value) = genesis_obj.get_mut(overrides_key) {
1009                match (&genesis_value, overrides_obj.get(overrides_key)) {
1010                    // recurse if genesis value is an object
1011                    (serde_json::Value::Object(_), Some(overrides_value))
1012                        if overrides_value.is_object() =>
1013                    {
1014                        merge(genesis_value, overrides_value);
1015                    },
1016                    // override if genesis value not an object
1017                    (_, Some(overrides_value)) => {
1018                        trace!("overriding: {:?} / {:?}", genesis_value, overrides_value);
1019                        *genesis_value = overrides_value.clone();
1020                    },
1021                    _ => {
1022                        trace!("not match!");
1023                    },
1024                }
1025            } else {
1026                // Allow to add keys, see (https://github.com/paritytech/zombienet/issues/1614)
1027                warn!(
1028                    "key: {overrides_key} not present in genesis_obj: {:?} (adding key)",
1029                    genesis_obj
1030                );
1031                let overrides_value = overrides_obj.get(overrides_key).expect(&format!(
1032                    "overrides_key {overrides_key} should be present in the overrides obj. qed"
1033                ));
1034                genesis_obj.insert(overrides_key.clone(), overrides_value.clone());
1035            }
1036        }
1037    }
1038}
1039
1040fn clear_authorities(
1041    runtime_config_ptr: &str,
1042    chain_spec_json: &mut serde_json::Value,
1043    ctx: &Context,
1044) {
1045    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1046        // clear keys (session, aura, grandpa)
1047        if val.get("session").is_some() {
1048            val["session"]["keys"] = json!([]);
1049        }
1050
1051        if val.get("aura").is_some() {
1052            val["aura"]["authorities"] = json!([]);
1053        }
1054
1055        if val.get("grandpa").is_some() {
1056            val["grandpa"]["authorities"] = json!([]);
1057        }
1058
1059        // clear collatorSelector
1060        if val.get("collatorSelection").is_some() {
1061            val["collatorSelection"]["invulnerables"] = json!([]);
1062        }
1063
1064        // clear staking but not `validatorCount` if `devStakers` is set
1065        if val.get("staking").is_some() && ctx == &Context::Relay {
1066            val["staking"]["invulnerables"] = json!([]);
1067            val["staking"]["stakers"] = json!([]);
1068
1069            if val["staking"]["devStakers"] == json!(null) {
1070                val["staking"]["validatorCount"] = json!(0);
1071            }
1072        }
1073    } else {
1074        unreachable!("pointer to runtime config should be valid!")
1075    }
1076}
1077
1078fn get_staking_min(runtime_config_ptr: &str, chain_spec_json: &mut serde_json::Value) -> u128 {
1079    // get min staking
1080    let staking_ptr = format!("{runtime_config_ptr}/staking/stakers");
1081    if let Some(stakers) = chain_spec_json.pointer(&staking_ptr) {
1082        // stakers should be an array
1083        let min = stakers[0][2].clone();
1084        min.as_u64().unwrap_or(0).into()
1085    } else {
1086        0
1087    }
1088}
1089
1090fn add_balances(
1091    runtime_config_ptr: &str,
1092    chain_spec_json: &mut serde_json::Value,
1093    nodes: &Vec<NodeSpec>,
1094    token_decimals: u8,
1095    staking_min: u128,
1096) {
1097    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1098        let Some(balances) = val.pointer("/balances/balances") else {
1099            // should be a info log
1100            warn!("NO 'balances' key in runtime config, skipping...");
1101            return;
1102        };
1103
1104        // create a balance map
1105        let mut balances_map = generate_balance_map(balances);
1106        for node in nodes {
1107            if node.initial_balance.eq(&0) {
1108                continue;
1109            };
1110
1111            // TODO: handle error here and check the `accounts.accounts` design
1112            // Double down the minimal stake defined
1113            let balance = std::cmp::max(node.initial_balance, staking_min * 2);
1114            for k in ["sr", "sr_stash"] {
1115                let account = node.accounts.accounts.get(k).unwrap();
1116                balances_map.insert(account.address.clone(), balance);
1117            }
1118        }
1119
1120        // ensure zombie account (//Zombie) have funds
1121        // we will use for internal usage (e.g new validators)
1122        balances_map.insert(
1123            "5FTcLfwFc7ctvqp3RhbEig6UuHLHcHVRujuUm8r21wy4dAR8".to_string(),
1124            1000 * 10_u128.pow(token_decimals as u32),
1125        );
1126
1127        // convert the map and store again
1128        let new_balances: Vec<(&String, &u128)> =
1129            balances_map.iter().collect::<Vec<(&String, &u128)>>();
1130
1131        val["balances"]["balances"] = json!(new_balances);
1132    } else {
1133        unreachable!("pointer to runtime config should be valid!")
1134    }
1135}
1136
1137fn get_node_keys(
1138    node: &NodeSpec,
1139    session_key: SessionKeyType,
1140    asset_hub_polkadot: bool,
1141) -> GenesisNodeKey {
1142    let sr_account = node.accounts.accounts.get("sr").unwrap();
1143    let sr_stash = node.accounts.accounts.get("sr_stash").unwrap();
1144    let ed_account = node.accounts.accounts.get("ed").unwrap();
1145    let ec_account = node.accounts.accounts.get("ec").unwrap();
1146    let eth_account = node.accounts.accounts.get("eth").unwrap();
1147    let mut keys = HashMap::new();
1148    for k in [
1149        "babe",
1150        "im_online",
1151        "parachain_validator",
1152        "authority_discovery",
1153        "para_validator",
1154        "para_assignment",
1155        "aura",
1156        "nimbus",
1157        "vrf",
1158    ] {
1159        if k == "aura" && asset_hub_polkadot {
1160            keys.insert(k.to_string(), ed_account.address.clone());
1161            continue;
1162        }
1163        keys.insert(k.to_string(), sr_account.address.clone());
1164    }
1165
1166    keys.insert("grandpa".to_string(), ed_account.address.clone());
1167    keys.insert("beefy".to_string(), ec_account.address.clone());
1168    keys.insert("eth".to_string(), eth_account.public_key.clone());
1169
1170    let account_to_use = match session_key {
1171        SessionKeyType::Default => sr_account.address.clone(),
1172        SessionKeyType::Stash => sr_stash.address.clone(),
1173        SessionKeyType::Evm => format!("0x{}", eth_account.public_key),
1174    };
1175
1176    (account_to_use.clone(), account_to_use, keys)
1177}
1178fn add_authorities(
1179    runtime_config_ptr: &str,
1180    chain_spec_json: &mut serde_json::Value,
1181    nodes: &[&NodeSpec],
1182    session_key: SessionKeyType,
1183) {
1184    let asset_hub_polkadot = chain_spec_json
1185        .get("id")
1186        .and_then(|v| v.as_str())
1187        .map(|id| id.starts_with("asset-hub-polkadot"))
1188        .unwrap_or_default();
1189    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1190        if let Some(session_keys) = val.pointer_mut("/session/keys") {
1191            let keys: Vec<GenesisNodeKey> = nodes
1192                .iter()
1193                .map(|node| get_node_keys(node, session_key, asset_hub_polkadot))
1194                .collect();
1195            *session_keys = json!(keys);
1196        } else {
1197            warn!("⚠️  'session/keys' key not present in runtime config.");
1198        }
1199    } else {
1200        unreachable!("pointer to runtime config should be valid!")
1201    }
1202}
1203fn add_hrmp_channels(
1204    runtime_config_ptr: &str,
1205    chain_spec_json: &mut serde_json::Value,
1206    hrmp_channels: &[HrmpChannelConfig],
1207) {
1208    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1209        if let Some(preopen_hrmp_channels) = val.pointer_mut("/hrmp/preopenHrmpChannels") {
1210            let hrmp_channels = hrmp_channels
1211                .iter()
1212                .map(|c| {
1213                    (
1214                        c.sender(),
1215                        c.recipient(),
1216                        c.max_capacity(),
1217                        c.max_message_size(),
1218                    )
1219                })
1220                .collect::<Vec<_>>();
1221            *preopen_hrmp_channels = json!(hrmp_channels);
1222        } else {
1223            warn!("⚠️  'hrmp/preopenHrmpChannels' key not present in runtime config.");
1224        }
1225    } else {
1226        unreachable!("pointer to runtime config should be valid!")
1227    }
1228}
1229
1230fn add_aura_authorities(
1231    runtime_config_ptr: &str,
1232    chain_spec_json: &mut serde_json::Value,
1233    nodes: &[&NodeSpec],
1234    _key_type: KeyType,
1235) {
1236    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1237        if let Some(aura_authorities) = val.pointer_mut("/aura/authorities") {
1238            let keys: Vec<String> = nodes
1239                .iter()
1240                .map(|node| {
1241                    node.accounts
1242                        .accounts
1243                        .get("sr")
1244                        .expect(&format!(
1245                            "'sr' account should be set at spec computation {THIS_IS_A_BUG}"
1246                        ))
1247                        .address
1248                        .clone()
1249                })
1250                .collect();
1251            *aura_authorities = json!(keys);
1252        } else {
1253            warn!("⚠️  'aura/authorities' key not present in runtime config.");
1254        }
1255    } else {
1256        unreachable!("pointer to runtime config should be valid!")
1257    }
1258}
1259
1260fn add_grandpa_authorities(
1261    runtime_config_ptr: &str,
1262    chain_spec_json: &mut serde_json::Value,
1263    nodes: &[&NodeSpec],
1264    _key_type: KeyType,
1265) {
1266    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1267        if let Some(grandpa_authorities) = val.pointer_mut("/grandpa/authorities") {
1268            let keys: Vec<(String, usize)> = nodes
1269                .iter()
1270                .map(|node| {
1271                    (
1272                        node.accounts
1273                            .accounts
1274                            .get("ed")
1275                            .expect(&format!(
1276                                "'ed' account should be set at spec computation {THIS_IS_A_BUG}"
1277                            ))
1278                            .address
1279                            .clone(),
1280                        1,
1281                    )
1282                })
1283                .collect();
1284            *grandpa_authorities = json!(keys);
1285        } else {
1286            warn!("⚠️  'grandpa/authorities' key not present in runtime config.");
1287        }
1288    } else {
1289        unreachable!("pointer to runtime config should be valid!")
1290    }
1291}
1292
1293fn add_staking(
1294    runtime_config_ptr: &str,
1295    chain_spec_json: &mut serde_json::Value,
1296    nodes: &Vec<NodeSpec>,
1297    staking_min: u128,
1298) {
1299    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1300        let Some(_) = val.pointer("/staking") else {
1301            // should be a info log
1302            warn!("NO 'staking' key in runtime config, skipping...");
1303            return;
1304        };
1305
1306        let mut stakers = vec![];
1307        let mut invulnerables = vec![];
1308        for node in nodes {
1309            let sr_stash_addr = &node
1310                .accounts
1311                .accounts
1312                .get("sr_stash")
1313                .expect("'sr_stash account should be defined for the node. qed")
1314                .address;
1315            stakers.push(json!([
1316                sr_stash_addr,
1317                sr_stash_addr,
1318                staking_min,
1319                "Validator"
1320            ]));
1321
1322            if node.is_invulnerable {
1323                invulnerables.push(sr_stash_addr);
1324            }
1325        }
1326
1327        val["staking"]["validatorCount"] = json!(stakers.len());
1328        val["staking"]["stakers"] = json!(stakers);
1329        val["staking"]["invulnerables"] = json!(invulnerables);
1330    } else {
1331        unreachable!("pointer to runtime config should be valid!")
1332    }
1333}
1334
1335// TODO: (team)
1336// fn add_nominators() {}
1337
1338// // TODO: (team) we should think a better way to use the decorators from
1339// // current version (ts).
1340// fn para_custom() { todo!() }
1341fn override_parachain_info(
1342    runtime_config_ptr: &str,
1343    chain_spec_json: &mut serde_json::Value,
1344    para_id: u32,
1345) {
1346    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1347        if let Some(parachain_id) = val.pointer_mut("/parachainInfo/parachainId") {
1348            *parachain_id = json!(para_id)
1349        } else {
1350            // Add warning here!
1351        }
1352    } else {
1353        unreachable!("pointer to runtime config should be valid!")
1354    }
1355}
1356fn add_collator_selection(
1357    runtime_config_ptr: &str,
1358    chain_spec_json: &mut serde_json::Value,
1359    nodes: &[&NodeSpec],
1360    session_key: SessionKeyType,
1361) {
1362    if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
1363        let key_type = if let SessionKeyType::Evm = session_key {
1364            "eth"
1365        } else {
1366            "sr"
1367        };
1368        let keys: Vec<String> = nodes
1369            .iter()
1370            .map(|node| {
1371                node.accounts
1372                    .accounts
1373                    .get(key_type)
1374                    .expect(&format!(
1375                        "'sr' account should be set at spec computation {THIS_IS_A_BUG}"
1376                    ))
1377                    .address
1378                    .clone()
1379            })
1380            .collect();
1381
1382        // collatorSelection.invulnerables
1383        if let Some(invulnerables) = val.pointer_mut("/collatorSelection/invulnerables") {
1384            *invulnerables = json!(keys);
1385        } else {
1386            // TODO: add a nice warning here.
1387            debug!("⚠️  'invulnerables' not present in spec, will not be customized");
1388        }
1389    } else {
1390        unreachable!("pointer to runtime config should be valid!")
1391    }
1392}
1393
1394// Helpers
1395fn generate_balance_map(balances: &serde_json::Value) -> HashMap<String, u128> {
1396    // SAFETY: balances is always an array in chain-spec with items [k,v]
1397    let balances_map: HashMap<String, u128> =
1398        serde_json::from_value::<Vec<(String, u128)>>(balances.to_owned())
1399            .unwrap()
1400            .iter()
1401            .fold(HashMap::new(), |mut memo, balance| {
1402                memo.insert(balance.0.clone(), balance.1);
1403                memo
1404            });
1405    balances_map
1406}
1407
1408#[cfg(test)]
1409mod tests {
1410    use std::fs;
1411
1412    use configuration::HrmpChannelConfigBuilder;
1413
1414    use super::*;
1415    use crate::{generators, shared::types::NodeAccounts};
1416
1417    const ROCOCO_LOCAL_PLAIN_TESTING: &str = "./testing/rococo-local-plain.json";
1418
1419    fn chain_spec_test(file: &str) -> serde_json::Value {
1420        let content = fs::read_to_string(file).unwrap();
1421        serde_json::from_str(&content).unwrap()
1422    }
1423
1424    fn chain_spec_with_stake() -> serde_json::Value {
1425        json!({"genesis": {
1426            "runtimeGenesis" : {
1427                "patch": {
1428                    "staking": {
1429                        "forceEra": "NotForcing",
1430                        "invulnerables": [
1431                          "5GNJqTPyNqANBkUVMN1LPPrxXnFouWXoe2wNSmmEoLctxiZY",
1432                          "5HpG9w8EBLe5XCrbczpwq5TSXvedjrBGCwqxK1iQ7qUsSWFc"
1433                        ],
1434                        "minimumValidatorCount": 1,
1435                        "slashRewardFraction": 100000000,
1436                        "stakers": [
1437                          [
1438                            "5GNJqTPyNqANBkUVMN1LPPrxXnFouWXoe2wNSmmEoLctxiZY",
1439                            "5GNJqTPyNqANBkUVMN1LPPrxXnFouWXoe2wNSmmEoLctxiZY",
1440                            100000000000001_u128,
1441                            "Validator"
1442                          ],
1443                          [
1444                            "5HpG9w8EBLe5XCrbczpwq5TSXvedjrBGCwqxK1iQ7qUsSWFc",
1445                            "5HpG9w8EBLe5XCrbczpwq5TSXvedjrBGCwqxK1iQ7qUsSWFc",
1446                            100000000000000_u128,
1447                            "Validator"
1448                          ]
1449                        ],
1450                        "validatorCount": 2
1451                    },
1452                }
1453            }
1454        }})
1455    }
1456
1457    fn chain_spec_with_dev_stakers() -> serde_json::Value {
1458        json!({"genesis": {
1459            "runtimeGenesis" : {
1460                "patch": {
1461                    "staking": {
1462                        "activeEra": [
1463                            0,
1464                            0,
1465                            0
1466                        ],
1467                        "canceledPayout": 0,
1468                        "devStakers": [
1469                            2000,
1470                            25000
1471                        ],
1472                        "forceEra": "NotForcing",
1473                        "invulnerables": [],
1474                        "maxNominatorCount": null,
1475                        "maxValidatorCount": null,
1476                        "minNominatorBond": 0,
1477                        "minValidatorBond": 0,
1478                        "slashRewardFraction": 0,
1479                        "stakers": [],
1480                        "validatorCount": 500
1481                    },
1482                }
1483            }
1484        }})
1485    }
1486
1487    #[test]
1488    fn get_min_stake_works() {
1489        let mut chain_spec_json = chain_spec_with_stake();
1490
1491        let pointer = get_runtime_config_pointer(&chain_spec_json).unwrap();
1492        let min = get_staking_min(&pointer, &mut chain_spec_json);
1493
1494        assert_eq!(100000000000001, min);
1495    }
1496
1497    #[test]
1498    fn dev_stakers_not_override_count_works() {
1499        let mut chain_spec_json = chain_spec_with_dev_stakers();
1500
1501        let pointer = get_runtime_config_pointer(&chain_spec_json).unwrap();
1502        clear_authorities(&pointer, &mut chain_spec_json, &Context::Relay);
1503
1504        let validator_count = chain_spec_json
1505            .pointer(&format!("{pointer}/staking/validatorCount"))
1506            .unwrap();
1507        assert_eq!(validator_count, &json!(500));
1508    }
1509
1510    #[test]
1511    fn dev_stakers_override_count_works() {
1512        let mut chain_spec_json = chain_spec_with_stake();
1513
1514        let pointer = get_runtime_config_pointer(&chain_spec_json).unwrap();
1515        clear_authorities(&pointer, &mut chain_spec_json, &Context::Relay);
1516
1517        let validator_count = chain_spec_json
1518            .pointer(&format!("{pointer}/staking/validatorCount"))
1519            .unwrap();
1520        assert_eq!(validator_count, &json!(0));
1521    }
1522
1523    #[test]
1524    fn overrides_from_toml_works() {
1525        use serde::{Deserialize, Serialize};
1526
1527        #[derive(Debug, Serialize, Deserialize)]
1528        struct MockConfig {
1529            #[serde(rename = "genesis", skip_serializing_if = "Option::is_none")]
1530            genesis_overrides: Option<serde_json::Value>,
1531        }
1532
1533        let mut chain_spec_json = chain_spec_test(ROCOCO_LOCAL_PLAIN_TESTING);
1534        // Could also be  something like [genesis.runtimeGenesis.patch.balances]
1535        const TOML: &str = "[genesis.runtime.balances]
1536            devAccounts = [
1537            20000,
1538            1000000000000000000,
1539            \"//Sender//{}\"
1540        ]";
1541        let override_toml: MockConfig = toml::from_str(TOML).unwrap();
1542        let overrides = override_toml.genesis_overrides.unwrap();
1543        let pointer = get_runtime_config_pointer(&chain_spec_json).unwrap();
1544
1545        let percolated_overrides = percolate_overrides(&pointer, &overrides)
1546            .map_err(|e| GeneratorError::ChainSpecGeneration(e.to_string()))
1547            .unwrap();
1548        trace!("percolated_overrides: {:#?}", percolated_overrides);
1549        if let Some(genesis) = chain_spec_json.pointer_mut(&pointer) {
1550            merge(genesis, percolated_overrides);
1551        }
1552
1553        trace!("chain spec: {chain_spec_json:#?}");
1554        assert!(chain_spec_json
1555            .pointer("/genesis/runtime/balances/devAccounts")
1556            .is_some());
1557    }
1558
1559    #[test]
1560    fn add_balances_works() {
1561        let mut spec_plain = chain_spec_test(ROCOCO_LOCAL_PLAIN_TESTING);
1562        let mut name = String::from("luca");
1563        let initial_balance = 1_000_000_000_000_u128;
1564        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
1565        let accounts = NodeAccounts {
1566            accounts: generators::generate_node_keys(&seed).unwrap(),
1567            seed,
1568        };
1569        let node = NodeSpec {
1570            name,
1571            accounts,
1572            initial_balance,
1573            ..Default::default()
1574        };
1575
1576        let nodes = vec![node];
1577        add_balances("/genesis/runtime", &mut spec_plain, &nodes, 12, 0);
1578
1579        let new_balances = spec_plain
1580            .pointer("/genesis/runtime/balances/balances")
1581            .unwrap();
1582
1583        let balances_map = generate_balance_map(new_balances);
1584
1585        // sr and sr_stash keys exists
1586        let sr = nodes[0].accounts.accounts.get("sr").unwrap();
1587        let sr_stash = nodes[0].accounts.accounts.get("sr_stash").unwrap();
1588        assert_eq!(balances_map.get(&sr.address).unwrap(), &initial_balance);
1589        assert_eq!(
1590            balances_map.get(&sr_stash.address).unwrap(),
1591            &initial_balance
1592        );
1593    }
1594
1595    #[test]
1596    fn add_balances_ensure_zombie_account() {
1597        let mut spec_plain = chain_spec_test(ROCOCO_LOCAL_PLAIN_TESTING);
1598
1599        let balances = spec_plain
1600            .pointer("/genesis/runtime/balances/balances")
1601            .unwrap();
1602        let balances_map = generate_balance_map(balances);
1603
1604        let nodes: Vec<NodeSpec> = vec![];
1605        add_balances("/genesis/runtime", &mut spec_plain, &nodes, 12, 0);
1606
1607        let new_balances = spec_plain
1608            .pointer("/genesis/runtime/balances/balances")
1609            .unwrap();
1610
1611        let new_balances_map = generate_balance_map(new_balances);
1612
1613        // sr and sr_stash keys exists
1614        assert!(new_balances_map.contains_key("5FTcLfwFc7ctvqp3RhbEig6UuHLHcHVRujuUm8r21wy4dAR8"));
1615        assert_eq!(new_balances_map.len(), balances_map.len() + 1);
1616    }
1617
1618    #[test]
1619    fn add_balances_spec_without_balances() {
1620        let mut spec_plain = chain_spec_test(ROCOCO_LOCAL_PLAIN_TESTING);
1621
1622        {
1623            let balances = spec_plain.pointer_mut("/genesis/runtime/balances").unwrap();
1624            *balances = json!(serde_json::Value::Null);
1625        }
1626
1627        let mut name = String::from("luca");
1628        let initial_balance = 1_000_000_000_000_u128;
1629        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
1630        let accounts = NodeAccounts {
1631            accounts: generators::generate_node_keys(&seed).unwrap(),
1632            seed,
1633        };
1634        let node = NodeSpec {
1635            name,
1636            accounts,
1637            initial_balance,
1638            ..Default::default()
1639        };
1640
1641        let nodes = vec![node];
1642        add_balances("/genesis/runtime", &mut spec_plain, &nodes, 12, 0);
1643
1644        let new_balances = spec_plain.pointer("/genesis/runtime/balances/balances");
1645
1646        // assert 'balances' is not created
1647        assert_eq!(new_balances, None);
1648    }
1649
1650    #[test]
1651    fn add_staking_works() {
1652        let mut chain_spec_json = chain_spec_with_stake();
1653        let mut name = String::from("luca");
1654        let initial_balance = 1_000_000_000_000_u128;
1655        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
1656        let accounts = NodeAccounts {
1657            accounts: generators::generate_node_keys(&seed).unwrap(),
1658            seed,
1659        };
1660        let node = NodeSpec {
1661            name,
1662            accounts,
1663            initial_balance,
1664            ..Default::default()
1665        };
1666
1667        let pointer = get_runtime_config_pointer(&chain_spec_json).unwrap();
1668        let min = get_staking_min(&pointer, &mut chain_spec_json);
1669
1670        let nodes = vec![node];
1671        add_staking(&pointer, &mut chain_spec_json, &nodes, min);
1672
1673        let new_staking = chain_spec_json
1674            .pointer("/genesis/runtimeGenesis/patch/staking")
1675            .unwrap();
1676
1677        // stakers should be one (with the luca sr_stash accounts)
1678        let sr_stash = nodes[0].accounts.accounts.get("sr_stash").unwrap();
1679        assert_eq!(new_staking["stakers"][0][0], json!(sr_stash.address));
1680        // with the calculated minimal bound
1681        assert_eq!(new_staking["stakers"][0][2], json!(min));
1682        // and only one
1683        assert_eq!(new_staking["stakers"].as_array().unwrap().len(), 1);
1684    }
1685
1686    #[test]
1687    fn adding_hrmp_channels_works() {
1688        let mut spec_plain = chain_spec_test(ROCOCO_LOCAL_PLAIN_TESTING);
1689
1690        {
1691            let current_hrmp_channels = spec_plain
1692                .pointer("/genesis/runtime/hrmp/preopenHrmpChannels")
1693                .unwrap();
1694            // assert should be empty
1695            assert_eq!(current_hrmp_channels, &json!([]));
1696        }
1697
1698        let para_100_101 = HrmpChannelConfigBuilder::new()
1699            .with_sender(100)
1700            .with_recipient(101)
1701            .build();
1702        let para_101_100 = HrmpChannelConfigBuilder::new()
1703            .with_sender(101)
1704            .with_recipient(100)
1705            .build();
1706        let channels = vec![para_100_101, para_101_100];
1707
1708        add_hrmp_channels("/genesis/runtime", &mut spec_plain, &channels);
1709        let new_hrmp_channels = spec_plain
1710            .pointer("/genesis/runtime/hrmp/preopenHrmpChannels")
1711            .unwrap()
1712            .as_array()
1713            .unwrap();
1714
1715        assert_eq!(new_hrmp_channels.len(), 2);
1716        assert_eq!(new_hrmp_channels.first().unwrap()[0], 100);
1717        assert_eq!(new_hrmp_channels.first().unwrap()[1], 101);
1718        assert_eq!(new_hrmp_channels.last().unwrap()[0], 101);
1719        assert_eq!(new_hrmp_channels.last().unwrap()[1], 100);
1720    }
1721
1722    #[test]
1723    fn adding_hrmp_channels_to_an_spec_without_channels() {
1724        let mut spec_plain = chain_spec_test("./testing/rococo-local-plain.json");
1725
1726        {
1727            let hrmp = spec_plain.pointer_mut("/genesis/runtime/hrmp").unwrap();
1728            *hrmp = json!(serde_json::Value::Null);
1729        }
1730
1731        let para_100_101 = HrmpChannelConfigBuilder::new()
1732            .with_sender(100)
1733            .with_recipient(101)
1734            .build();
1735        let para_101_100 = HrmpChannelConfigBuilder::new()
1736            .with_sender(101)
1737            .with_recipient(100)
1738            .build();
1739        let channels = vec![para_100_101, para_101_100];
1740
1741        add_hrmp_channels("/genesis/runtime", &mut spec_plain, &channels);
1742        let new_hrmp_channels = spec_plain.pointer("/genesis/runtime/hrmp/preopenHrmpChannels");
1743
1744        // assert 'preopenHrmpChannels' is not created
1745        assert_eq!(new_hrmp_channels, None);
1746    }
1747
1748    #[test]
1749    fn get_node_keys_works() {
1750        let mut name = String::from("luca");
1751        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
1752        let accounts = NodeAccounts {
1753            accounts: generators::generate_node_keys(&seed).unwrap(),
1754            seed,
1755        };
1756        let node = NodeSpec {
1757            name,
1758            accounts,
1759            ..Default::default()
1760        };
1761
1762        let sr = &node.accounts.accounts["sr"];
1763        let keys = [
1764            ("babe".into(), sr.address.clone()),
1765            ("im_online".into(), sr.address.clone()),
1766            ("parachain_validator".into(), sr.address.clone()),
1767            ("authority_discovery".into(), sr.address.clone()),
1768            ("para_validator".into(), sr.address.clone()),
1769            ("para_assignment".into(), sr.address.clone()),
1770            ("aura".into(), sr.address.clone()),
1771            ("nimbus".into(), sr.address.clone()),
1772            ("vrf".into(), sr.address.clone()),
1773            (
1774                "grandpa".into(),
1775                node.accounts.accounts["ed"].address.clone(),
1776            ),
1777            ("beefy".into(), node.accounts.accounts["ec"].address.clone()),
1778            ("eth".into(), node.accounts.accounts["eth"].address.clone()),
1779        ]
1780        .into();
1781
1782        // Stash
1783        let sr_stash = &node.accounts.accounts["sr_stash"];
1784        let node_key = get_node_keys(&node, SessionKeyType::Stash, false);
1785        assert_eq!(node_key.0, sr_stash.address);
1786        assert_eq!(node_key.1, sr_stash.address);
1787        assert_eq!(node_key.2, keys);
1788        // Non-stash
1789        let node_key = get_node_keys(&node, SessionKeyType::Default, false);
1790        assert_eq!(node_key.0, sr.address);
1791        assert_eq!(node_key.1, sr.address);
1792        assert_eq!(node_key.2, keys);
1793    }
1794
1795    #[test]
1796    fn get_node_keys_supports_asset_hub_polkadot() {
1797        let mut name = String::from("luca");
1798        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
1799        let accounts = NodeAccounts {
1800            accounts: generators::generate_node_keys(&seed).unwrap(),
1801            seed,
1802        };
1803        let node = NodeSpec {
1804            name,
1805            accounts,
1806            ..Default::default()
1807        };
1808
1809        let node_key = get_node_keys(&node, SessionKeyType::default(), false);
1810        assert_eq!(node_key.2["aura"], node.accounts.accounts["sr"].address);
1811
1812        let node_key = get_node_keys(&node, SessionKeyType::default(), true);
1813        assert_eq!(node_key.2["aura"], node.accounts.accounts["ed"].address);
1814    }
1815}