zombienet_orchestrator/network_spec/
node.rs

1use std::path::PathBuf;
2
3use configuration::shared::{
4    node::{EnvVar, NodeConfig},
5    resources::Resources,
6    types::{Arg, AssetLocation, Command, Image},
7};
8use multiaddr::Multiaddr;
9use provider::types::Port;
10use serde::{Deserialize, Serialize};
11use support::constants::THIS_IS_A_BUG;
12
13use crate::{
14    errors::OrchestratorError,
15    generators,
16    network::AddNodeOptions,
17    shared::{
18        macros,
19        types::{ChainDefaultContext, NodeAccount, NodeAccounts, ParkedPort},
20    },
21    AddCollatorOptions,
22};
23
24macros::create_add_options!(AddNodeSpecOpts {
25     override_eth_key: Option<String>
26});
27
28macro_rules! impl_from_for_add_node_opts {
29    ($struct:ident) => {
30        impl From<$struct> for AddNodeSpecOpts {
31            fn from(value: $struct) -> Self {
32                Self {
33                    image: value.image,
34                    command: value.command,
35                    subcommand: value.subcommand,
36                    args: value.args,
37                    env: value.env,
38                    is_validator: value.is_validator,
39                    rpc_port: value.rpc_port,
40                    prometheus_port: value.prometheus_port,
41                    p2p_port: value.p2p_port,
42                    override_eth_key: value.override_eth_key,
43                }
44            }
45        }
46    };
47}
48
49impl_from_for_add_node_opts!(AddNodeOptions);
50impl_from_for_add_node_opts!(AddCollatorOptions);
51
52/// A node configuration, with fine-grained configuration options.
53#[derive(Debug, Clone, Default, Serialize, Deserialize)]
54pub struct NodeSpec {
55    // Node name (should be unique or an index will be appended).
56    pub(crate) name: String,
57
58    /// Node key, used for compute the p2p identity.
59    pub(crate) key: String,
60
61    // libp2p local identity
62    pub(crate) peer_id: String,
63
64    /// Accounts to be injected in the keystore.
65    pub(crate) accounts: NodeAccounts,
66
67    /// Image to run (only podman/k8s). Override the default.
68    pub(crate) image: Option<Image>,
69
70    /// Command to run the node. Override the default.
71    pub(crate) command: Command,
72
73    /// Optional subcommand for the node.
74    pub(crate) subcommand: Option<Command>,
75
76    /// Arguments to use for node. Appended to default.
77    pub(crate) args: Vec<Arg>,
78
79    // The help command output containing the available arguments.
80    pub(crate) available_args_output: Option<String>,
81
82    /// Wether the node is a validator.
83    pub(crate) is_validator: bool,
84
85    /// Whether the node keys must be added to invulnerables.
86    pub(crate) is_invulnerable: bool,
87
88    /// Whether the node is a bootnode.
89    pub(crate) is_bootnode: bool,
90
91    /// Node initial balance present in genesis.
92    pub(crate) initial_balance: u128,
93
94    /// Environment variables to set (inside pod for podman/k8s, inside shell for native).
95    pub(crate) env: Vec<EnvVar>,
96
97    /// List of node's bootnodes addresses to use. Appended to default.
98    pub(crate) bootnodes_addresses: Vec<Multiaddr>,
99
100    /// Default resources. Override the default.
101    pub(crate) resources: Option<Resources>,
102
103    /// Websocket port to use.
104    pub(crate) ws_port: ParkedPort,
105
106    /// RPC port to use.
107    pub(crate) rpc_port: ParkedPort,
108
109    /// Prometheus port to use.
110    pub(crate) prometheus_port: ParkedPort,
111
112    /// P2P port to use.
113    pub(crate) p2p_port: ParkedPort,
114
115    /// libp2p cert hash to use with `webrtc` transport.
116    pub(crate) p2p_cert_hash: Option<String>,
117
118    /// Database snapshot. Override the default.
119    pub(crate) db_snapshot: Option<AssetLocation>,
120
121    /// P2P port to use by full node if this is the case
122    pub(crate) full_node_p2p_port: Option<ParkedPort>,
123    /// Prometheus port to use by full node if this is the case
124    pub(crate) full_node_prometheus_port: Option<ParkedPort>,
125
126    /// Optionally specify a log path for the node
127    pub(crate) node_log_path: Option<PathBuf>,
128
129    /// Optionally specify a keystore path for the node
130    pub(crate) keystore_path: Option<PathBuf>,
131
132    /// Keystore key types to generate.
133    /// Supports short form (e.g., "audi") using predefined schemas,
134    /// or long form (e.g., "audi_sr") with explicit schema (sr, ed, ec).
135    pub(crate) keystore_key_types: Vec<String>,
136}
137
138impl NodeSpec {
139    pub fn from_config(
140        node_config: &NodeConfig,
141        chain_context: &ChainDefaultContext,
142        full_node_present: bool,
143        evm_based: bool,
144    ) -> Result<Self, OrchestratorError> {
145        // Check first if the image is set at node level, then try with the default
146        let image = node_config.image().or(chain_context.default_image).cloned();
147
148        // Check first if the command is set at node level, then try with the default
149        let command = if let Some(cmd) = node_config.command() {
150            cmd.clone()
151        } else if let Some(cmd) = chain_context.default_command {
152            cmd.clone()
153        } else {
154            return Err(OrchestratorError::InvalidNodeConfig(
155                node_config.name().into(),
156                "command".to_string(),
157            ));
158        };
159
160        let subcommand = node_config.subcommand().cloned();
161
162        // If `args` is set at `node` level use them
163        // otherwise use the default_args (can be empty).
164        let args: Vec<Arg> = if node_config.args().is_empty() {
165            chain_context
166                .default_args
167                .iter()
168                .map(|x| x.to_owned().clone())
169                .collect()
170        } else {
171            node_config.args().into_iter().cloned().collect()
172        };
173
174        let (key, peer_id) = generators::generate_node_identity(node_config.name())?;
175
176        let mut name = node_config.name().to_string();
177        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
178        let accounts = generators::generate_node_keys(&seed)?;
179        let mut accounts = NodeAccounts { seed, accounts };
180
181        if evm_based {
182            if let Some(session_key) = node_config.override_eth_key() {
183                accounts
184                    .accounts
185                    .insert("eth".into(), NodeAccount::new(session_key, session_key));
186            }
187        }
188
189        let db_snapshot = match (node_config.db_snapshot(), chain_context.default_db_snapshot) {
190            (Some(db_snapshot), _) => Some(db_snapshot),
191            (None, Some(db_snapshot)) => Some(db_snapshot),
192            _ => None,
193        };
194
195        let (full_node_p2p_port, full_node_prometheus_port) = if full_node_present {
196            (
197                Some(generators::generate_node_port(None)?),
198                Some(generators::generate_node_port(None)?),
199            )
200        } else {
201            (None, None)
202        };
203
204        Ok(Self {
205            name: node_config.name().to_string(),
206            key,
207            peer_id,
208            image,
209            command,
210            subcommand,
211            args,
212            available_args_output: None,
213            is_validator: node_config.is_validator(),
214            is_invulnerable: node_config.is_invulnerable(),
215            is_bootnode: node_config.is_bootnode(),
216            initial_balance: node_config.initial_balance(),
217            env: node_config.env().into_iter().cloned().collect(),
218            bootnodes_addresses: node_config
219                .bootnodes_addresses()
220                .into_iter()
221                .cloned()
222                .collect(),
223            resources: node_config.resources().cloned(),
224            p2p_cert_hash: node_config.p2p_cert_hash().map(str::to_string),
225            db_snapshot: db_snapshot.cloned(),
226            accounts,
227            ws_port: generators::generate_node_port(node_config.ws_port())?,
228            rpc_port: generators::generate_node_port(node_config.rpc_port())?,
229            prometheus_port: generators::generate_node_port(node_config.prometheus_port())?,
230            p2p_port: generators::generate_node_port(node_config.p2p_port())?,
231            full_node_p2p_port,
232            full_node_prometheus_port,
233            node_log_path: node_config.node_log_path().cloned(),
234            keystore_path: node_config.keystore_path().cloned(),
235            keystore_key_types: node_config
236                .keystore_key_types()
237                .into_iter()
238                .map(str::to_string)
239                .collect(),
240        })
241    }
242
243    pub fn from_ad_hoc(
244        name: impl Into<String>,
245        options: AddNodeSpecOpts,
246        chain_context: &ChainDefaultContext,
247        full_node_present: bool,
248        evm_based: bool,
249    ) -> Result<Self, OrchestratorError> {
250        // Check first if the image is set at node level, then try with the default
251        let image = if let Some(img) = options.image {
252            Some(img.clone())
253        } else {
254            chain_context.default_image.cloned()
255        };
256
257        let name = name.into();
258        // Check first if the command is set at node level, then try with the default
259        let command = if let Some(cmd) = options.command {
260            cmd.clone()
261        } else if let Some(cmd) = chain_context.default_command {
262            cmd.clone()
263        } else {
264            return Err(OrchestratorError::InvalidNodeConfig(
265                name,
266                "command".to_string(),
267            ));
268        };
269
270        let subcommand = options.subcommand.clone();
271
272        // If `args` is set at `node` level use them
273        // otherwise use the default_args (can be empty).
274        let args: Vec<Arg> = if options.args.is_empty() {
275            chain_context
276                .default_args
277                .iter()
278                .map(|x| x.to_owned().clone())
279                .collect()
280        } else {
281            options.args
282        };
283
284        let (key, peer_id) = generators::generate_node_identity(&name)?;
285
286        let mut name_capitalized = name.clone();
287        let seed = format!(
288            "//{}{name_capitalized}",
289            name_capitalized.remove(0).to_uppercase()
290        );
291        let accounts = generators::generate_node_keys(&seed)?;
292        let mut accounts = NodeAccounts { seed, accounts };
293
294        if evm_based {
295            if let Some(session_key) = options.override_eth_key.as_ref() {
296                accounts
297                    .accounts
298                    .insert("eth".into(), NodeAccount::new(session_key, session_key));
299            }
300        }
301
302        let (full_node_p2p_port, full_node_prometheus_port) = if full_node_present {
303            (
304                Some(generators::generate_node_port(None)?),
305                Some(generators::generate_node_port(None)?),
306            )
307        } else {
308            (None, None)
309        };
310
311        //
312        Ok(Self {
313            name,
314            key,
315            peer_id,
316            image,
317            command,
318            subcommand,
319            args,
320            available_args_output: None,
321            is_validator: options.is_validator,
322            is_invulnerable: false,
323            is_bootnode: false,
324            initial_balance: 0,
325            env: options.env,
326            bootnodes_addresses: vec![],
327            resources: None,
328            p2p_cert_hash: None,
329            db_snapshot: None,
330            accounts,
331            // should be deprecated now!
332            ws_port: generators::generate_node_port(None)?,
333            rpc_port: generators::generate_node_port(options.rpc_port)?,
334            prometheus_port: generators::generate_node_port(options.prometheus_port)?,
335            p2p_port: generators::generate_node_port(options.p2p_port)?,
336            full_node_p2p_port,
337            full_node_prometheus_port,
338            node_log_path: None,
339            keystore_path: None,
340            keystore_key_types: vec![],
341        })
342    }
343
344    pub(crate) fn supports_arg(&self, arg: impl AsRef<str>) -> bool {
345        self.available_args_output
346            .as_ref()
347            .expect(&format!(
348                "available args should be present at this point {THIS_IS_A_BUG}"
349            ))
350            .contains(arg.as_ref())
351    }
352
353    pub fn command(&self) -> &str {
354        self.command.as_str()
355    }
356}