zombienet_orchestrator/generators/
command.rs

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

zombienet_orchestrator/generators/
command.rs
