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

zombienet_orchestrator/generators/
command.rs
