zombienet_orchestrator/generators/
command.rs

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

zombienet_orchestrator/generators/
command.rs
