zombienet_orchestrator/generators/
command.rs

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

zombienet_orchestrator/generators/
command.rs
