zombienet_orchestrator/generators/
command.rs

1use configuration::types::Arg;
2use support::constants::THIS_IS_A_BUG;
3
4use crate::{network_spec::node::NodeSpec, shared::constants::*};
5
6pub struct GenCmdOptions<'a> {
7    pub relay_chain_name: &'a str,
8    pub cfg_path: &'a str,
9    pub data_path: &'a str,
10    pub relay_data_path: &'a str,
11    pub use_wrapper: bool,
12    pub bootnode_addr: Vec<String>,
13    pub use_default_ports_in_cmd: bool,
14    pub is_native: bool,
15}
16
17impl Default for GenCmdOptions<'_> {
18    fn default() -> Self {
19        Self {
20            relay_chain_name: "rococo-local",
21            cfg_path: "/cfg",
22            data_path: "/data",
23            relay_data_path: "/relay-data",
24            use_wrapper: true,
25            bootnode_addr: vec![],
26            use_default_ports_in_cmd: false,
27            is_native: true,
28        }
29    }
30}
31
32const FLAGS_ADDED_BY_US: [&str; 3] = ["--no-telemetry", "--collator", "--"];
33const OPS_ADDED_BY_US: [&str; 6] = [
34    "--chain",
35    "--name",
36    "--rpc-cors",
37    "--rpc-methods",
38    "--parachain-id",
39    "--node-key",
40];
41
42// TODO: can we abstract this and use only one fn (or at least split and reuse in small fns)
43pub fn generate_for_cumulus_node(
44    node: &NodeSpec,
45    options: GenCmdOptions,
46    para_id: u32,
47) -> (String, Vec<String>) {
48    let NodeSpec {
49        key,
50        args,
51        is_validator,
52        bootnodes_addresses,
53        ..
54    } = node;
55
56    let mut tmp_args: Vec<String> = vec!["--node-key".into(), key.clone()];
57
58    if !args.contains(&Arg::Flag("--prometheus-external".into())) {
59        tmp_args.push("--prometheus-external".into())
60    }
61
62    if *is_validator && !args.contains(&Arg::Flag("--validator".into())) {
63        tmp_args.push("--collator".into())
64    }
65
66    if !bootnodes_addresses.is_empty() {
67        tmp_args.push("--bootnodes".into());
68        let bootnodes = bootnodes_addresses
69            .iter()
70            .map(|m| m.to_string())
71            .collect::<Vec<String>>()
72            .join(" ");
73        tmp_args.push(bootnodes)
74    }
75
76    // ports
77    let (prometheus_port, rpc_port, p2p_port) =
78        resolve_ports(node, options.use_default_ports_in_cmd);
79
80    tmp_args.push("--prometheus-port".into());
81    tmp_args.push(prometheus_port.to_string());
82
83    tmp_args.push("--rpc-port".into());
84    tmp_args.push(rpc_port.to_string());
85
86    tmp_args.push("--listen-addr".into());
87    tmp_args.push(format!("/ip4/0.0.0.0/tcp/{p2p_port}/ws"));
88
89    let mut collator_args: &[Arg] = &[];
90    let mut full_node_args: &[Arg] = &[];
91    if !args.is_empty() {
92        if let Some(index) = args.iter().position(|arg| match arg {
93            Arg::Flag(flag) => flag.eq("--"),
94            Arg::Option(..) => false,
95            Arg::Array(..) => false,
96        }) {
97            (collator_args, full_node_args) = args.split_at(index);
98        } else {
99            // Assume args are those specified for collator only
100            collator_args = args;
101        }
102    }
103
104    // set our base path
105    tmp_args.push("--base-path".into());
106    tmp_args.push(options.data_path.into());
107
108    let node_specific_bootnodes: Vec<String> = node
109        .bootnodes_addresses
110        .iter()
111        .map(|b| b.to_string())
112        .collect();
113    let full_bootnodes = [node_specific_bootnodes, options.bootnode_addr].concat();
114    if !full_bootnodes.is_empty() {
115        tmp_args.push("--bootnodes".into());
116        tmp_args.push(full_bootnodes.join(" "));
117    }
118
119    let mut full_node_p2p_needs_to_be_injected = true;
120    let mut full_node_prometheus_needs_to_be_injected = true;
121    let mut full_node_args_filtered = full_node_args
122        .iter()
123        .filter_map(|arg| match arg {
124            Arg::Flag(flag) => {
125                if FLAGS_ADDED_BY_US.contains(&flag.as_str()) {
126                    None
127                } else {
128                    Some(vec![flag.to_owned()])
129                }
130            },
131            Arg::Option(k, v) => {
132                if OPS_ADDED_BY_US.contains(&k.as_str()) {
133                    None
134                } else if k.eq(&"port") {
135                    if v.eq(&"30333") {
136                        full_node_p2p_needs_to_be_injected = true;
137                        None
138                    } else {
139                        // non default
140                        full_node_p2p_needs_to_be_injected = false;
141                        Some(vec![k.to_owned(), v.to_owned()])
142                    }
143                } else if k.eq(&"--prometheus-port") {
144                    if v.eq(&"9616") {
145                        full_node_prometheus_needs_to_be_injected = true;
146                        None
147                    } else {
148                        // non default
149                        full_node_prometheus_needs_to_be_injected = false;
150                        Some(vec![k.to_owned(), v.to_owned()])
151                    }
152                } else {
153                    Some(vec![k.to_owned(), v.to_owned()])
154                }
155            },
156            Arg::Array(k, v) => {
157                let mut args = vec![k.to_owned()];
158                args.extend(v.to_owned());
159                Some(args)
160            },
161        })
162        .flatten()
163        .collect::<Vec<String>>();
164
165    let full_p2p_port = node
166        .full_node_p2p_port
167        .as_ref()
168        .expect(&format!(
169            "full node p2p_port should be specifed: {THIS_IS_A_BUG}"
170        ))
171        .0;
172    let full_prometheus_port = node
173        .full_node_prometheus_port
174        .as_ref()
175        .expect(&format!(
176            "full node prometheus_port should be specifed: {THIS_IS_A_BUG}"
177        ))
178        .0;
179
180    // full_node: change p2p port if is the default
181    if full_node_p2p_needs_to_be_injected {
182        full_node_args_filtered.push("--port".into());
183        full_node_args_filtered.push(full_p2p_port.to_string());
184    }
185
186    // full_node: change prometheus port if is the default
187    if full_node_prometheus_needs_to_be_injected {
188        full_node_args_filtered.push("--prometheus-port".into());
189        full_node_args_filtered.push(full_prometheus_port.to_string());
190    }
191
192    let mut args_filtered = collator_args
193        .iter()
194        .filter_map(|arg| match arg {
195            Arg::Flag(flag) => {
196                if FLAGS_ADDED_BY_US.contains(&flag.as_str()) {
197                    None
198                } else {
199                    Some(vec![flag.to_owned()])
200                }
201            },
202            Arg::Option(k, v) => {
203                if OPS_ADDED_BY_US.contains(&k.as_str()) {
204                    None
205                } else {
206                    Some(vec![k.to_owned(), v.to_owned()])
207                }
208            },
209            Arg::Array(k, v) => {
210                let mut args = vec![k.to_owned()];
211                args.extend(v.to_owned());
212                Some(args)
213            },
214        })
215        .flatten()
216        .collect::<Vec<String>>();
217
218    tmp_args.append(&mut args_filtered);
219
220    let parachain_spec_path = format!("{}/{}.json", options.cfg_path, para_id);
221    let mut final_args = vec![
222        node.command.as_str().to_string(),
223        "--chain".into(),
224        parachain_spec_path,
225        "--name".into(),
226        node.name.clone(),
227        "--rpc-cors".into(),
228        "all".into(),
229        "--rpc-methods".into(),
230        "unsafe".into(),
231    ];
232
233    // The `--unsafe-rpc-external` option spawns an additional RPC server on a random port,
234    // which can conflict with reserved ports, causing an "Address already in use" error
235    // when using the `native` provider. Since this option isn't needed for `native`,
236    // it should be omitted in that case.
237    if !options.is_native {
238        final_args.push("--unsafe-rpc-external".into());
239    }
240
241    final_args.append(&mut tmp_args);
242
243    let relaychain_spec_path = format!("{}/{}.json", options.cfg_path, options.relay_chain_name);
244    let mut full_node_injected: Vec<String> = vec![
245        "--".into(),
246        "--base-path".into(),
247        options.relay_data_path.into(),
248        "--chain".into(),
249        relaychain_spec_path,
250        "--execution".into(),
251        "wasm".into(),
252    ];
253
254    final_args.append(&mut full_node_injected);
255    final_args.append(&mut full_node_args_filtered);
256
257    if options.use_wrapper {
258        ("/cfg/zombie-wrapper.sh".to_string(), final_args)
259    } else {
260        (final_args.remove(0), final_args)
261    }
262}
263
264pub fn generate_for_node(
265    node: &NodeSpec,
266    options: GenCmdOptions,
267    para_id: Option<u32>,
268) -> (String, Vec<String>) {
269    let NodeSpec {
270        key,
271        args,
272        is_validator,
273        bootnodes_addresses,
274        ..
275    } = node;
276    let mut tmp_args: Vec<String> = vec![
277        "--node-key".into(),
278        key.clone(),
279        // TODO:(team) we should allow to set the telemetry url from config
280        "--no-telemetry".into(),
281    ];
282
283    if !args.contains(&Arg::Flag("--prometheus-external".into())) {
284        tmp_args.push("--prometheus-external".into())
285    }
286
287    if let Some(para_id) = para_id {
288        tmp_args.push("--parachain-id".into());
289        tmp_args.push(para_id.to_string());
290    }
291
292    if *is_validator && !args.contains(&Arg::Flag("--validator".into())) {
293        tmp_args.push("--validator".into());
294        if node.supports_arg("--insecure-validator-i-know-what-i-do") {
295            tmp_args.push("--insecure-validator-i-know-what-i-do".into());
296        }
297    }
298
299    if !bootnodes_addresses.is_empty() {
300        tmp_args.push("--bootnodes".into());
301        let bootnodes = bootnodes_addresses
302            .iter()
303            .map(|m| m.to_string())
304            .collect::<Vec<String>>()
305            .join(" ");
306        tmp_args.push(bootnodes)
307    }
308
309    // ports
310    let (prometheus_port, rpc_port, p2p_port) =
311        resolve_ports(node, options.use_default_ports_in_cmd);
312
313    // Prometheus
314    tmp_args.push("--prometheus-port".into());
315    tmp_args.push(prometheus_port.to_string());
316
317    // RPC
318    // TODO (team): do we want to support old --ws-port?
319    tmp_args.push("--rpc-port".into());
320    tmp_args.push(rpc_port.to_string());
321
322    let listen_value = if let Some(listen_val) = args.iter().find_map(|arg| match arg {
323        Arg::Flag(_) => None,
324        Arg::Option(k, v) => {
325            if k.eq("--listen-addr") {
326                Some(v)
327            } else {
328                None
329            }
330        },
331        Arg::Array(..) => None,
332    }) {
333        let mut parts = listen_val.split('/').collect::<Vec<&str>>();
334        // TODO: move this to error
335        let port_part = parts
336            .get_mut(4)
337            .expect(&format!("should have at least 5 parts {THIS_IS_A_BUG}"));
338        let port_to_use = p2p_port.to_string();
339        *port_part = port_to_use.as_str();
340        parts.join("/")
341    } else {
342        format!("/ip4/0.0.0.0/tcp/{p2p_port}/ws")
343    };
344
345    tmp_args.push("--listen-addr".into());
346    tmp_args.push(listen_value);
347
348    // set our base path
349    tmp_args.push("--base-path".into());
350    tmp_args.push(options.data_path.into());
351
352    let node_specific_bootnodes: Vec<String> = node
353        .bootnodes_addresses
354        .iter()
355        .map(|b| b.to_string())
356        .collect();
357    let full_bootnodes = [node_specific_bootnodes, options.bootnode_addr].concat();
358    if !full_bootnodes.is_empty() {
359        tmp_args.push("--bootnodes".into());
360        tmp_args.push(full_bootnodes.join(" "));
361    }
362
363    // add the rest of the args
364    let mut args_filtered = args
365        .iter()
366        .filter_map(|arg| match arg {
367            Arg::Flag(flag) => {
368                if FLAGS_ADDED_BY_US.contains(&flag.as_str()) {
369                    None
370                } else {
371                    Some(vec![flag.to_owned()])
372                }
373            },
374            Arg::Option(k, v) => {
375                if OPS_ADDED_BY_US.contains(&k.as_str()) {
376                    None
377                } else {
378                    Some(vec![k.to_owned(), v.to_owned()])
379                }
380            },
381            Arg::Array(k, v) => {
382                let mut args = vec![k.to_owned()];
383                args.extend(v.to_owned());
384                Some(args)
385            },
386        })
387        .flatten()
388        .collect::<Vec<String>>();
389
390    tmp_args.append(&mut args_filtered);
391
392    let chain_spec_path = format!("{}/{}.json", options.cfg_path, options.relay_chain_name);
393    let mut final_args = vec![
394        node.command.as_str().to_string(),
395        "--chain".into(),
396        chain_spec_path,
397        "--name".into(),
398        node.name.clone(),
399        "--rpc-cors".into(),
400        "all".into(),
401        "--rpc-methods".into(),
402        "unsafe".into(),
403    ];
404
405    // The `--unsafe-rpc-external` option spawns an additional RPC server on a random port,
406    // which can conflict with reserved ports, causing an "Address already in use" error
407    // when using the `native` provider. Since this option isn't needed for `native`,
408    // it should be omitted in that case.
409    if !options.is_native {
410        final_args.push("--unsafe-rpc-external".into());
411    }
412
413    final_args.append(&mut tmp_args);
414
415    if let Some(ref subcommand) = node.subcommand {
416        final_args.insert(1, subcommand.as_str().to_string());
417    }
418
419    if options.use_wrapper {
420        ("/cfg/zombie-wrapper.sh".to_string(), final_args)
421    } else {
422        (final_args.remove(0), final_args)
423    }
424}
425
426/// Returns (prometheus, rpc, p2p) ports to use in the command
427fn resolve_ports(node: &NodeSpec, use_default_ports_in_cmd: bool) -> (u16, u16, u16) {
428    if use_default_ports_in_cmd {
429        (PROMETHEUS_PORT, RPC_PORT, P2P_PORT)
430    } else {
431        (node.prometheus_port.0, node.rpc_port.0, node.p2p_port.0)
432    }
433}
434
435#[cfg(test)]
436mod tests {
437    use super::*;
438    use crate::{generators, shared::types::NodeAccounts};
439
440    fn get_node_spec(full_node_present: bool) -> NodeSpec {
441        let mut name = String::from("luca");
442        let initial_balance = 1_000_000_000_000_u128;
443        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
444        let accounts = NodeAccounts {
445            accounts: generators::generate_node_keys(&seed).unwrap(),
446            seed,
447        };
448        let (full_node_p2p_port, full_node_prometheus_port) = if full_node_present {
449            (
450                Some(generators::generate_node_port(None).unwrap()),
451                Some(generators::generate_node_port(None).unwrap()),
452            )
453        } else {
454            (None, None)
455        };
456        NodeSpec {
457            name,
458            accounts,
459            initial_balance,
460            full_node_p2p_port,
461            full_node_prometheus_port,
462            ..Default::default()
463        }
464    }
465
466    #[test]
467    fn generate_for_native_cumulus_node_works() {
468        let node = get_node_spec(true);
469        let opts = GenCmdOptions {
470            use_wrapper: false,
471            is_native: true,
472            ..GenCmdOptions::default()
473        };
474
475        let (program, args) = generate_for_cumulus_node(&node, opts, 1000);
476        assert_eq!(program.as_str(), "polkadot");
477
478        let divider_flag = args.iter().position(|x| x == "--").unwrap();
479
480        // ensure full node ports
481        let i = args[divider_flag..]
482            .iter()
483            .position(|x| {
484                x == node
485                    .full_node_p2p_port
486                    .as_ref()
487                    .unwrap()
488                    .0
489                    .to_string()
490                    .as_str()
491            })
492            .unwrap();
493        assert_eq!(&args[divider_flag + i - 1], "--port");
494
495        let i = args[divider_flag..]
496            .iter()
497            .position(|x| {
498                x == node
499                    .full_node_prometheus_port
500                    .as_ref()
501                    .unwrap()
502                    .0
503                    .to_string()
504                    .as_str()
505            })
506            .unwrap();
507        assert_eq!(&args[divider_flag + i - 1], "--prometheus-port");
508
509        assert!(!args.iter().any(|arg| arg == "--unsafe-rpc-external"));
510    }
511
512    #[test]
513    fn generate_for_native_cumulus_node_rpc_external_is_not_removed_if_is_set_by_user() {
514        let mut node = get_node_spec(true);
515        node.args.push("--unsafe-rpc-external".into());
516        let opts = GenCmdOptions {
517            use_wrapper: false,
518            is_native: true,
519            ..GenCmdOptions::default()
520        };
521
522        let (_, args) = generate_for_cumulus_node(&node, opts, 1000);
523
524        assert!(args.iter().any(|arg| arg == "--unsafe-rpc-external"));
525    }
526
527    #[test]
528    fn generate_for_non_native_cumulus_node_works() {
529        let node = get_node_spec(true);
530        let opts = GenCmdOptions {
531            use_wrapper: false,
532            is_native: false,
533            ..GenCmdOptions::default()
534        };
535
536        let (program, args) = generate_for_cumulus_node(&node, opts, 1000);
537        assert_eq!(program.as_str(), "polkadot");
538
539        let divider_flag = args.iter().position(|x| x == "--").unwrap();
540
541        // ensure full node ports
542        let i = args[divider_flag..]
543            .iter()
544            .position(|x| {
545                x == node
546                    .full_node_p2p_port
547                    .as_ref()
548                    .unwrap()
549                    .0
550                    .to_string()
551                    .as_str()
552            })
553            .unwrap();
554        assert_eq!(&args[divider_flag + i - 1], "--port");
555
556        let i = args[divider_flag..]
557            .iter()
558            .position(|x| {
559                x == node
560                    .full_node_prometheus_port
561                    .as_ref()
562                    .unwrap()
563                    .0
564                    .to_string()
565                    .as_str()
566            })
567            .unwrap();
568        assert_eq!(&args[divider_flag + i - 1], "--prometheus-port");
569
570        // we expect to find this arg in collator node part
571        assert!(&args[0..divider_flag]
572            .iter()
573            .any(|arg| arg == "--unsafe-rpc-external"));
574    }
575
576    #[test]
577    fn generate_for_native_node_rpc_external_works() {
578        let node = get_node_spec(false);
579        let opts = GenCmdOptions {
580            use_wrapper: false,
581            is_native: true,
582            ..GenCmdOptions::default()
583        };
584
585        let (program, args) = generate_for_node(&node, opts, Some(1000));
586        assert_eq!(program.as_str(), "polkadot");
587
588        assert!(!args.iter().any(|arg| arg == "--unsafe-rpc-external"));
589    }
590
591    #[test]
592    fn generate_for_non_native_node_rpc_external_works() {
593        let node = get_node_spec(false);
594        let opts = GenCmdOptions {
595            use_wrapper: false,
596            is_native: false,
597            ..GenCmdOptions::default()
598        };
599
600        let (program, args) = generate_for_node(&node, opts, Some(1000));
601        assert_eq!(program.as_str(), "polkadot");
602
603        assert!(args.iter().any(|arg| arg == "--unsafe-rpc-external"));
604    }
605}
zombienet_orchestrator/generators/command.rs

zombienet_orchestrator/generators/
command.rs
