zombienet_orchestrator/generators/
keystore.rs

1use std::{
2    path::{Path, PathBuf},
3    vec,
4};
5
6use hex::encode;
7use support::{constants::THIS_IS_A_BUG, fs::FileSystem};
8
9use super::errors::GeneratorError;
10use crate::{
11    generators::keystore_key_types::{parse_keystore_key_types, KeystoreKeyType},
12    shared::types::NodeAccounts,
13    ScopedFilesystem,
14};
15
16/// Generates keystore files for a node.
17///
18/// # Arguments
19/// * `acc` - The node accounts containing the seed and public keys
20/// * `node_files_path` - The path where keystore files will be created
21/// * `scoped_fs` - The scoped filesystem for file operations
22/// * `asset_hub_polkadot` - Whether this is for asset-hub-polkadot (affects aura key scheme)
23/// * `keystore_key_types` - Optional list of key type specifications
24///
25/// If `keystore_key_types` is empty, all default key types will be generated.
26/// Otherwise, only the specified key types will be generated.
27pub async fn generate<'a, T>(
28    acc: &NodeAccounts,
29    node_files_path: impl AsRef<Path>,
30    scoped_fs: &ScopedFilesystem<'a, T>,
31    asset_hub_polkadot: bool,
32    keystore_key_types: Vec<&str>,
33) -> Result<Vec<PathBuf>, GeneratorError>
34where
35    T: FileSystem,
36{
37    // Create local keystore
38    scoped_fs.create_dir_all(node_files_path.as_ref()).await?;
39    let mut filenames = vec![];
40
41    // Parse the key type specifications
42    let key_types = parse_keystore_key_types(&keystore_key_types, asset_hub_polkadot);
43
44    let futures: Vec<_> = key_types
45        .iter()
46        .map(|key_type| {
47            let filename = generate_keystore_filename(key_type, acc);
48            let file_path = PathBuf::from(format!(
49                "{}/{}",
50                node_files_path.as_ref().to_string_lossy(),
51                filename
52            ));
53            let content = format!("\"{}\"", acc.seed);
54            (filename, scoped_fs.write(file_path, content))
55        })
56        .collect();
57
58    for (filename, future) in futures {
59        future.await?;
60        filenames.push(PathBuf::from(filename));
61    }
62
63    Ok(filenames)
64}
65
66/// Generates the keystore filename for a given key type.
67///
68/// The filename format is: `{hex_encoded_key_type}{public_key}`
69fn generate_keystore_filename(key_type: &KeystoreKeyType, acc: &NodeAccounts) -> String {
70    let account_key = key_type.scheme.account_key();
71    let pk = acc
72        .accounts
73        .get(account_key)
74        .expect(&format!(
75            "Key '{account_key}' should be set for node {THIS_IS_A_BUG}"
76        ))
77        .public_key
78        .as_str();
79
80    format!("{}{}", encode(&key_type.key_type), pk)
81}
82
83#[cfg(test)]
84mod tests {
85    use std::{collections::HashMap, ffi::OsString, str::FromStr};
86
87    use support::fs::in_memory::{InMemoryFile, InMemoryFileSystem};
88
89    use super::*;
90    use crate::shared::types::{NodeAccount, NodeAccounts};
91
92    fn create_test_accounts() -> NodeAccounts {
93        let mut accounts = HashMap::new();
94        accounts.insert(
95            "sr".to_string(),
96            NodeAccount::new("sr_address", "sr_public_key"),
97        );
98        accounts.insert(
99            "ed".to_string(),
100            NodeAccount::new("ed_address", "ed_public_key"),
101        );
102        accounts.insert(
103            "ec".to_string(),
104            NodeAccount::new("ec_address", "ec_public_key"),
105        );
106        NodeAccounts {
107            seed: "//Alice".to_string(),
108            accounts,
109        }
110    }
111
112    fn create_test_fs() -> InMemoryFileSystem {
113        InMemoryFileSystem::new(HashMap::from([(
114            OsString::from_str("/").unwrap(),
115            InMemoryFile::dir(),
116        )]))
117    }
118
119    #[tokio::test]
120    async fn generate_creates_default_keystore_files_when_no_key_types_specified() {
121        let accounts = create_test_accounts();
122        let fs = create_test_fs();
123        let base_dir = "/tmp/test";
124
125        let scoped_fs = ScopedFilesystem { fs: &fs, base_dir };
126        let key_types: Vec<&str> = vec![];
127
128        let res = generate(&accounts, "node1", &scoped_fs, false, key_types).await;
129        assert!(res.is_ok());
130
131        let filenames = res.unwrap();
132
133        assert!(filenames.len() > 10);
134
135        let filename_strs: Vec<String> = filenames
136            .iter()
137            .map(|p| p.to_string_lossy().to_string())
138            .collect();
139
140        // Check that aura key is generated (hex of "aura" is 61757261)
141        assert!(filename_strs.iter().any(|f| f.starts_with("61757261")));
142        // Check that babe key is generated (hex of "babe" is 62616265)
143        assert!(filename_strs.iter().any(|f| f.starts_with("62616265")));
144        // Check that gran key is generated (hex of "gran" is 6772616e)
145        assert!(filename_strs.iter().any(|f| f.starts_with("6772616e")));
146    }
147
148    #[tokio::test]
149    async fn generate_creates_only_specified_keystore_files() {
150        let accounts = create_test_accounts();
151        let fs = create_test_fs();
152        let base_dir = "/tmp/test";
153
154        let scoped_fs = ScopedFilesystem { fs: &fs, base_dir };
155        let key_types = vec!["audi", "gran"];
156
157        let res = generate(&accounts, "node1", &scoped_fs, false, key_types).await;
158
159        assert!(res.is_ok());
160
161        let filenames = res.unwrap();
162        assert_eq!(filenames.len(), 2);
163
164        let filename_strs: Vec<String> = filenames
165            .iter()
166            .map(|p| p.to_string_lossy().to_string())
167            .collect();
168
169        // audi uses sr scheme by default
170        assert!(filename_strs
171            .iter()
172            .any(|f| f.starts_with("61756469") && f.contains("sr_public_key")));
173        // gran uses ed scheme by default
174        assert!(filename_strs
175            .iter()
176            .any(|f| f.starts_with("6772616e") && f.contains("ed_public_key")));
177    }
178
179    #[tokio::test]
180    async fn generate_produces_correct_keystore_files() {
181        struct TestCase {
182            name: &'static str,
183            key_types: Vec<&'static str>,
184            asset_hub_polkadot: bool,
185            expected_prefix: &'static str,
186            expected_public_key: &'static str,
187        }
188
189        let test_cases = vec![
190            TestCase {
191                name: "explicit scheme override (gran_sr)",
192                key_types: vec!["gran_sr"],
193                asset_hub_polkadot: false,
194                expected_prefix: "6772616e", // "gran" in hex
195                expected_public_key: "sr_public_key",
196            },
197            TestCase {
198                name: "aura with asset_hub_polkadot uses ed",
199                key_types: vec!["aura"],
200                asset_hub_polkadot: true,
201                expected_prefix: "61757261", // "aura" in hex
202                expected_public_key: "ed_public_key",
203            },
204            TestCase {
205                name: "aura without asset_hub_polkadot uses sr",
206                key_types: vec!["aura"],
207                asset_hub_polkadot: false,
208                expected_prefix: "61757261", // "aura" in hex
209                expected_public_key: "sr_public_key",
210            },
211            TestCase {
212                name: "custom key type with explicit ec scheme",
213                key_types: vec!["cust_ec"],
214                asset_hub_polkadot: false,
215                expected_prefix: "63757374", // "cust" in hex
216                expected_public_key: "ec_public_key",
217            },
218        ];
219
220        for tc in test_cases {
221            let accounts = create_test_accounts();
222            let fs = create_test_fs();
223            let scoped_fs = ScopedFilesystem {
224                fs: &fs,
225                base_dir: "/tmp/test",
226            };
227
228            let key_types: Vec<&str> = tc.key_types.clone();
229            let res = generate(
230                &accounts,
231                "node1",
232                &scoped_fs,
233                tc.asset_hub_polkadot,
234                key_types,
235            )
236            .await;
237
238            assert!(
239                res.is_ok(),
240                "[{}] Expected Ok but got: {:?}",
241                tc.name,
242                res.err()
243            );
244            let filenames = res.unwrap();
245
246            assert_eq!(filenames.len(), 1, "[{}] Expected 1 file", tc.name);
247
248            let filename = filenames[0].to_string_lossy().to_string();
249            assert!(
250                filename.starts_with(tc.expected_prefix),
251                "[{}] Expected prefix '{}', got '{}'",
252                tc.name,
253                tc.expected_prefix,
254                filename
255            );
256            assert!(
257                filename.contains(tc.expected_public_key),
258                "[{}] Expected public key '{}' in '{}'",
259                tc.name,
260                tc.expected_public_key,
261                filename
262            );
263        }
264    }
265
266    #[tokio::test]
267    async fn generate_ignores_invalid_key_specs_and_uses_defaults() {
268        let accounts = create_test_accounts();
269        let fs = create_test_fs();
270        let scoped_fs = ScopedFilesystem {
271            fs: &fs,
272            base_dir: "/tmp/test",
273        };
274
275        let key_types = vec![
276            "invalid", // Too long
277            "xxx",     // Too short
278            "audi_xx", // Invalid sceme
279        ];
280
281        let res = generate(&accounts, "node1", &scoped_fs, false, key_types).await;
282
283        assert!(res.is_ok());
284        let filenames = res.unwrap();
285
286        // Should fall back to defaults since all specs are invalid
287        assert!(filenames.len() > 10);
288    }
289}
zombienet_orchestrator/generators/keystore.rs

zombienet_orchestrator/generators/
keystore.rs
