Skip to main content

pop_chains/utils/
helpers.rs

1// SPDX-License-Identifier: GPL-3.0
2
3use crate::errors::Error;
4use sc_chain_spec::GenesisConfigBuilderRuntimeCaller;
5use std::{
6	fs::{self, OpenOptions},
7	io::{self, Write, stdin, stdout},
8	path::{Path, PathBuf},
9};
10
11pub(crate) type HostFunctions = (
12	sp_statement_store::runtime_api::HostFunctions,
13	cumulus_primitives_proof_size_hostfunction::storage_proof_size::HostFunctions,
14);
15
16pub(crate) fn sanitize(target: &Path) -> Result<(), Error> {
17	if target.exists() {
18		print!("\"{}\" directory exists. Do you want to clean it? [y/n]: ", target.display());
19		stdout().flush()?;
20
21		let mut input = String::new();
22		stdin().read_line(&mut input)?;
23
24		if input.trim().to_lowercase() == "y" {
25			fs::remove_dir_all(target).map_err(|_| Error::Aborted)?;
26		} else {
27			return Err(Error::Aborted);
28		}
29	}
30	Ok(())
31}
32
33/// Check if the initial endowment input by the user is a valid balance.
34///
35/// # Arguments
36///
37/// * `initial_endowment` - initial endowment amount to be checked for validity.
38pub fn is_initial_endowment_valid(initial_endowment: &str) -> bool {
39	initial_endowment.parse::<u128>().is_ok() ||
40		is_valid_bitwise_left_shift(initial_endowment).is_ok()
41}
42
43// Auxiliary method to check if the endowment input with a shift left (1u64 << 60) format is valid.
44// Parse the self << rhs format and check the shift left operation is valid.
45fn is_valid_bitwise_left_shift(initial_endowment: &str) -> Result<u128, Error> {
46	let v: Vec<&str> = initial_endowment.split(" << ").collect();
47	if v.len() < 2 {
48		return Err(Error::EndowmentError);
49	}
50	let left = v[0]
51		.split('u') // parse 1u64 characters
52		.take(1)
53		.collect::<String>()
54		.parse::<u128>()
55		.map_err(|_e| Error::EndowmentError)?;
56	let right = v[1]
57		.chars()
58		.filter(|c| c.is_numeric()) // parse 1u64 characters
59		.collect::<String>()
60		.parse::<u32>()
61		.map_err(|_e| Error::EndowmentError)?;
62	left.checked_shl(right).ok_or(Error::EndowmentError)
63}
64
65pub(crate) fn write_to_file(path: &Path, contents: &str) -> Result<(), Error> {
66	let mut file = OpenOptions::new()
67		.write(true)
68		.truncate(true)
69		.create(true)
70		.open(path)
71		.map_err(Error::RustfmtError)?;
72
73	file.write_all(contents.as_bytes()).map_err(Error::RustfmtError)?;
74
75	if path.extension().is_some_and(|ext| ext == "rs") {
76		let output = std::process::Command::new("rustfmt")
77			.arg(path.to_str().unwrap())
78			.output()
79			.map_err(Error::RustfmtError)?;
80
81		if !output.status.success() {
82			return Err(Error::RustfmtError(io::Error::other(
83				"rustfmt exited with non-zero status code",
84			)));
85		}
86	}
87
88	Ok(())
89}
90
91/// Get genesis builder preset names of the runtime.
92///
93/// # Arguments
94/// * `binary_path` - Path to the runtime binary.
95pub fn get_preset_names(binary_path: &PathBuf) -> Result<Vec<String>, Error> {
96	let binary = fs::read(binary_path)?;
97	let genesis_config_builder = GenesisConfigBuilderRuntimeCaller::<HostFunctions>::new(&binary);
98	genesis_config_builder
99		.preset_names()
100		.map_err(|e| Error::GenesisBuilderError(e.to_string()))
101}
102
103#[cfg(test)]
104mod tests {
105	use super::*;
106	use crate::{ChainTemplate, generator::chain::ChainSpec};
107	use askama::Template;
108	use tempfile::tempdir;
109
110	#[test]
111	fn test_write_to_file() -> Result<(), Box<dyn std::error::Error>> {
112		let temp_dir = tempdir()?;
113		let chainspec = ChainSpec {
114			token_symbol: "DOT".to_string(),
115			decimals: 6,
116			initial_endowment: "1000000".to_string(),
117			based_on: ChainTemplate::Standard.to_string(),
118		};
119		let file_path = temp_dir.path().join("file.rs");
120		let _ = fs::write(&file_path, "");
121		write_to_file(&file_path, chainspec.render().expect("infallible").as_ref())?;
122		let generated_file_content =
123			fs::read_to_string(temp_dir.path().join("file.rs")).expect("Failed to read file");
124		assert!(
125			generated_file_content
126				.contains("properties.insert(\"tokenSymbol\".into(), \"DOT\".into());")
127		);
128		assert!(
129			generated_file_content
130				.contains("properties.insert(\"tokenDecimals\".into(), 6.into());")
131		);
132		assert!(generated_file_content.contains("1000000"));
133		assert!(generated_file_content.contains(
134			"properties.insert(\"basedOn\".into(), \"r0gue-io/base-parachain\".into());"
135		));
136
137		Ok(())
138	}
139
140	#[test]
141	fn test_is_initial_endowment_valid() {
142		assert!(is_initial_endowment_valid("100000"));
143		assert!(is_initial_endowment_valid("1u64 << 60"));
144		assert!(!is_initial_endowment_valid("wrong"));
145		assert!(!is_initial_endowment_valid(" "));
146	}
147
148	#[test]
149	fn test_left_shift() {
150		// Values from https://stackoverflow.com/questions/56392875/how-can-i-initialize-a-users-balance-in-a-substrate-blockchain
151		assert_eq!(is_valid_bitwise_left_shift("1 << 60").unwrap(), 1152921504606846976);
152		let result = is_valid_bitwise_left_shift("wrong");
153		assert!(result.is_err());
154	}
155
156	#[test]
157	fn test_get_preset_names() -> Result<(), Box<dyn std::error::Error>> {
158		let path = PathBuf::from("../../tests/runtimes/base_parachain_benchmark.wasm");
159		assert!(path.is_file());
160		let presets = get_preset_names(&path)?;
161		assert_eq!(presets, vec!["development", "local_testnet"]);
162		Ok(())
163	}
164}