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 '{}' should be set for node {THIS_IS_A_BUG}",
76            account_key
77        ))
78        .public_key
79        .as_str();
80
81    format!("{}{}", encode(&key_type.key_type), pk)
82}
83
84#[cfg(test)]
85mod tests {
86    use std::{collections::HashMap, ffi::OsString, str::FromStr};
87
88    use support::fs::in_memory::{InMemoryFile, InMemoryFileSystem};
89
90    use super::*;
91    use crate::shared::types::{NodeAccount, NodeAccounts};
92
93    fn create_test_accounts() -> NodeAccounts {
94        let mut accounts = HashMap::new();
95        accounts.insert(
96            "sr".to_string(),
97            NodeAccount::new("sr_address", "sr_public_key"),
98        );
99        accounts.insert(
100            "ed".to_string(),
101            NodeAccount::new("ed_address", "ed_public_key"),
102        );
103        accounts.insert(
104            "ec".to_string(),
105            NodeAccount::new("ec_address", "ec_public_key"),
106        );
107        NodeAccounts {
108            seed: "//Alice".to_string(),
109            accounts,
110        }
111    }
112
113    fn create_test_fs() -> InMemoryFileSystem {
114        InMemoryFileSystem::new(HashMap::from([(
115            OsString::from_str("/").unwrap(),
116            InMemoryFile::dir(),
117        )]))
118    }
119
120    #[tokio::test]
121    async fn generate_creates_default_keystore_files_when_no_key_types_specified() {
122        let accounts = create_test_accounts();
123        let fs = create_test_fs();
124        let base_dir = "/tmp/test";
125
126        let scoped_fs = ScopedFilesystem { fs: &fs, base_dir };
127        let key_types: Vec<&str> = vec![];
128
129        let res = generate(&accounts, "node1", &scoped_fs, false, key_types).await;
130        assert!(res.is_ok());
131
132        let filenames = res.unwrap();
133
134        assert!(filenames.len() > 10);
135
136        let filename_strs: Vec<String> = filenames
137            .iter()
138            .map(|p| p.to_string_lossy().to_string())
139            .collect();
140
141        // Check that aura key is generated (hex of "aura" is 61757261)
142        assert!(filename_strs.iter().any(|f| f.starts_with("61757261")));
143        // Check that babe key is generated (hex of "babe" is 62616265)
144        assert!(filename_strs.iter().any(|f| f.starts_with("62616265")));
145        // Check that gran key is generated (hex of "gran" is 6772616e)
146        assert!(filename_strs.iter().any(|f| f.starts_with("6772616e")));
147    }
148
149    #[tokio::test]
150    async fn generate_creates_only_specified_keystore_files() {
151        let accounts = create_test_accounts();
152        let fs = create_test_fs();
153        let base_dir = "/tmp/test";
154
155        let scoped_fs = ScopedFilesystem { fs: &fs, base_dir };
156        let key_types = vec!["audi", "gran"];
157
158        let res = generate(&accounts, "node1", &scoped_fs, false, key_types).await;
159
160        assert!(res.is_ok());
161
162        let filenames = res.unwrap();
163        assert_eq!(filenames.len(), 2);
164
165        let filename_strs: Vec<String> = filenames
166            .iter()
167            .map(|p| p.to_string_lossy().to_string())
168            .collect();
169
170        // audi uses sr scheme by default
171        assert!(filename_strs
172            .iter()
173            .any(|f| f.starts_with("61756469") && f.contains("sr_public_key")));
174        // gran uses ed scheme by default
175        assert!(filename_strs
176            .iter()
177            .any(|f| f.starts_with("6772616e") && f.contains("ed_public_key")));
178    }
179
180    #[tokio::test]
181    async fn generate_produces_correct_keystore_files() {
182        struct TestCase {
183            name: &'static str,
184            key_types: Vec<&'static str>,
185            asset_hub_polkadot: bool,
186            expected_prefix: &'static str,
187            expected_public_key: &'static str,
188        }
189
190        let test_cases = vec![
191            TestCase {
192                name: "explicit scheme override (gran_sr)",
193                key_types: vec!["gran_sr"],
194                asset_hub_polkadot: false,
195                expected_prefix: "6772616e", // "gran" in hex
196                expected_public_key: "sr_public_key",
197            },
198            TestCase {
199                name: "aura with asset_hub_polkadot uses ed",
200                key_types: vec!["aura"],
201                asset_hub_polkadot: true,
202                expected_prefix: "61757261", // "aura" in hex
203                expected_public_key: "ed_public_key",
204            },
205            TestCase {
206                name: "aura without asset_hub_polkadot uses sr",
207                key_types: vec!["aura"],
208                asset_hub_polkadot: false,
209                expected_prefix: "61757261", // "aura" in hex
210                expected_public_key: "sr_public_key",
211            },
212            TestCase {
213                name: "custom key type with explicit ec scheme",
214                key_types: vec!["cust_ec"],
215                asset_hub_polkadot: false,
216                expected_prefix: "63757374", // "cust" in hex
217                expected_public_key: "ec_public_key",
218            },
219        ];
220
221        for tc in test_cases {
222            let accounts = create_test_accounts();
223            let fs = create_test_fs();
224            let scoped_fs = ScopedFilesystem {
225                fs: &fs,
226                base_dir: "/tmp/test",
227            };
228
229            let key_types: Vec<&str> = tc.key_types.clone();
230            let res = generate(
231                &accounts,
232                "node1",
233                &scoped_fs,
234                tc.asset_hub_polkadot,
235                key_types,
236            )
237            .await;
238
239            assert!(
240                res.is_ok(),
241                "[{}] Expected Ok but got: {:?}",
242                tc.name,
243                res.err()
244            );
245            let filenames = res.unwrap();
246
247            assert_eq!(filenames.len(), 1, "[{}] Expected 1 file", tc.name);
248
249            let filename = filenames[0].to_string_lossy().to_string();
250            assert!(
251                filename.starts_with(tc.expected_prefix),
252                "[{}] Expected prefix '{}', got '{}'",
253                tc.name,
254                tc.expected_prefix,
255                filename
256            );
257            assert!(
258                filename.contains(tc.expected_public_key),
259                "[{}] Expected public key '{}' in '{}'",
260                tc.name,
261                tc.expected_public_key,
262                filename
263            );
264        }
265    }
266
267    #[tokio::test]
268    async fn generate_ignores_invalid_key_specs_and_uses_defaults() {
269        let accounts = create_test_accounts();
270        let fs = create_test_fs();
271        let scoped_fs = ScopedFilesystem {
272            fs: &fs,
273            base_dir: "/tmp/test",
274        };
275
276        let key_types = vec![
277            "invalid", // Too long
278            "xxx",     // Too short
279            "audi_xx", // Invalid sceme
280        ];
281
282        let res = generate(&accounts, "node1", &scoped_fs, false, key_types).await;
283
284        assert!(res.is_ok());
285        let filenames = res.unwrap();
286
287        // Should fall back to defaults since all specs are invalid
288        assert!(filenames.len() > 10);
289    }
290}
zombienet_orchestrator/generators/keystore.rs

zombienet_orchestrator/generators/
keystore.rs
