Skip to main content

zombienet_provider/shared/
types.rs

1use std::{
2    collections::HashMap,
3    path::{Path, PathBuf},
4    process::ExitStatus,
5};
6
7use configuration::shared::resources::Resources;
8use serde::{Deserialize, Serialize};
9
10pub type Port = u16;
11
12pub type ExecutionResult = Result<String, (ExitStatus, String)>;
13
14#[derive(Debug, Clone, PartialEq)]
15pub struct ProviderCapabilities {
16    // default ports internal
17    /// Ensure that we have an image for each node (k8s/podman/docker)
18    pub requires_image: bool,
19    /// Allow to customize the resources through manifest (k8s).
20    pub has_resources: bool,
21    /// Used in native to prefix filepath with fullpath
22    pub prefix_with_full_path: bool,
23    /// Use default ports in node cmd/args.
24    /// NOTE: generally used in k8s/dockers since the images expose those ports.
25    pub use_default_ports_in_cmd: bool,
26}
27
28#[derive(Debug, Clone)]
29pub struct SpawnNodeOptions {
30    /// Name of the node
31    pub name: String,
32    /// Image of the node (IFF is supported by the provider)
33    pub image: Option<String>,
34    /// Resources to apply to the node (IFF is supported by the provider)
35    pub resources: Option<Resources>,
36    /// Main command to execute
37    pub program: String,
38    /// Arguments to pass to the main command
39    pub args: Vec<String>,
40    /// Environment to set when running the `program`
41    pub env: Vec<(String, String)>,
42    // TODO: rename startup_files
43    /// Files to inject at startup
44    pub injected_files: Vec<TransferedFile>,
45    /// Paths to create before start the node (e.g keystore)
46    /// should be created with `create_dir_all` in order
47    /// to create the full path even when we have missing parts
48    pub created_paths: Vec<PathBuf>,
49    /// Database snapshot to be injected (a local `.tgz` resolved by the
50    /// orchestrator before spawn — providers extract it directly, no
51    /// download). See `orchestrator::generators::resolve_db_snapshots`.
52    pub db_snapshot: Option<PathBuf>,
53    pub port_mapping: Option<HashMap<Port, Port>>,
54    /// Optionally specify a log path for the node
55    pub node_log_path: Option<PathBuf>,
56}
57
58impl SpawnNodeOptions {
59    pub fn new<S>(name: S, program: S) -> Self
60    where
61        S: AsRef<str>,
62    {
63        Self {
64            name: name.as_ref().to_string(),
65            image: None,
66            resources: None,
67            program: program.as_ref().to_string(),
68            args: vec![],
69            env: vec![],
70            injected_files: vec![],
71            created_paths: vec![],
72            db_snapshot: None,
73            port_mapping: None,
74            node_log_path: None,
75        }
76    }
77
78    pub fn image<S>(mut self, image: S) -> Self
79    where
80        S: AsRef<str>,
81    {
82        self.image = Some(image.as_ref().to_string());
83        self
84    }
85
86    pub fn resources(mut self, resources: Resources) -> Self {
87        self.resources = Some(resources);
88        self
89    }
90
91    pub fn db_snapshot(mut self, db_snap: Option<PathBuf>) -> Self {
92        self.db_snapshot = db_snap;
93        self
94    }
95
96    pub fn args<S, I>(mut self, args: I) -> Self
97    where
98        S: AsRef<str>,
99        I: IntoIterator<Item = S>,
100    {
101        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
102        self
103    }
104
105    pub fn env<S, I>(mut self, env: I) -> Self
106    where
107        S: AsRef<str>,
108        I: IntoIterator<Item = (S, S)>,
109    {
110        self.env = env
111            .into_iter()
112            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
113            .collect();
114        self
115    }
116
117    pub fn injected_files<I>(mut self, injected_files: I) -> Self
118    where
119        I: IntoIterator<Item = TransferedFile>,
120    {
121        self.injected_files = injected_files.into_iter().collect();
122        self
123    }
124
125    pub fn created_paths<P, I>(mut self, created_paths: I) -> Self
126    where
127        P: AsRef<Path>,
128        I: IntoIterator<Item = P>,
129    {
130        self.created_paths = created_paths
131            .into_iter()
132            .map(|path| path.as_ref().into())
133            .collect();
134        self
135    }
136
137    pub fn port_mapping(mut self, ports: HashMap<Port, Port>) -> Self {
138        self.port_mapping = Some(ports);
139        self
140    }
141
142    pub fn node_log_path(mut self, path: Option<PathBuf>) -> Self {
143        self.node_log_path = path;
144        self
145    }
146}
147
148#[derive(Debug)]
149pub struct GenerateFileCommand {
150    pub program: String,
151    pub args: Vec<String>,
152    pub env: Vec<(String, String)>,
153    pub local_output_path: PathBuf,
154}
155
156impl GenerateFileCommand {
157    pub fn new<S, P>(program: S, local_output_path: P) -> Self
158    where
159        S: AsRef<str>,
160        P: AsRef<Path>,
161    {
162        Self {
163            program: program.as_ref().to_string(),
164            args: vec![],
165            env: vec![],
166            local_output_path: local_output_path.as_ref().into(),
167        }
168    }
169
170    pub fn args<S, I>(mut self, args: I) -> Self
171    where
172        S: AsRef<str>,
173        I: IntoIterator<Item = S>,
174    {
175        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
176        self
177    }
178
179    pub fn env<S, I>(mut self, env: I) -> Self
180    where
181        S: AsRef<str>,
182        I: IntoIterator<Item = (S, S)>,
183    {
184        self.env = env
185            .into_iter()
186            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
187            .collect();
188        self
189    }
190}
191
192#[derive(Debug)]
193pub struct GenerateFilesOptions {
194    pub commands: Vec<GenerateFileCommand>,
195    pub image: Option<String>,
196    pub injected_files: Vec<TransferedFile>,
197    // Allow to control the name of the node used to create the files.
198    pub temp_name: Option<String>,
199    pub expected_path: Option<PathBuf>,
200}
201
202impl GenerateFilesOptions {
203    pub fn new<I>(commands: I, image: Option<String>, expected_path: Option<PathBuf>) -> Self
204    where
205        I: IntoIterator<Item = GenerateFileCommand>,
206    {
207        Self {
208            commands: commands.into_iter().collect(),
209            injected_files: vec![],
210            image,
211            temp_name: None,
212            expected_path,
213        }
214    }
215
216    pub fn with_files<I>(
217        commands: I,
218        image: Option<String>,
219        injected_files: &[TransferedFile],
220        expected_path: Option<PathBuf>,
221    ) -> Self
222    where
223        I: IntoIterator<Item = GenerateFileCommand>,
224    {
225        Self {
226            commands: commands.into_iter().collect(),
227            injected_files: injected_files.into(),
228            image,
229            temp_name: None,
230            expected_path,
231        }
232    }
233
234    pub fn image<S>(mut self, image: S) -> Self
235    where
236        S: AsRef<str>,
237    {
238        self.image = Some(image.as_ref().to_string());
239        self
240    }
241
242    pub fn injected_files<I>(mut self, injected_files: I) -> Self
243    where
244        I: IntoIterator<Item = TransferedFile>,
245    {
246        self.injected_files = injected_files.into_iter().collect();
247        self
248    }
249
250    pub fn temp_name(mut self, name: impl Into<String>) -> Self {
251        self.temp_name = Some(name.into());
252        self
253    }
254}
255
256#[derive(Debug)]
257pub struct RunCommandOptions {
258    pub program: String,
259    pub args: Vec<String>,
260    pub env: Vec<(String, String)>,
261}
262
263impl RunCommandOptions {
264    pub fn new<S>(program: S) -> Self
265    where
266        S: AsRef<str>,
267    {
268        Self {
269            program: program.as_ref().to_string(),
270            args: vec![],
271            env: vec![],
272        }
273    }
274
275    pub fn args<S, I>(mut self, args: I) -> Self
276    where
277        S: AsRef<str>,
278        I: IntoIterator<Item = S>,
279    {
280        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
281        self
282    }
283
284    pub fn env<S, I>(mut self, env: I) -> Self
285    where
286        S: AsRef<str>,
287        I: IntoIterator<Item = (S, S)>,
288    {
289        self.env = env
290            .into_iter()
291            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
292            .collect();
293        self
294    }
295}
296
297pub struct RunScriptOptions {
298    pub local_script_path: PathBuf,
299    pub args: Vec<String>,
300    pub env: Vec<(String, String)>,
301}
302
303impl RunScriptOptions {
304    pub fn new<P>(local_script_path: P) -> Self
305    where
306        P: AsRef<Path>,
307    {
308        Self {
309            local_script_path: local_script_path.as_ref().into(),
310            args: vec![],
311            env: vec![],
312        }
313    }
314
315    pub fn args<S, I>(mut self, args: I) -> Self
316    where
317        S: AsRef<str>,
318        I: IntoIterator<Item = S>,
319    {
320        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
321        self
322    }
323
324    pub fn env<S, I>(mut self, env: I) -> Self
325    where
326        S: AsRef<str>,
327        I: IntoIterator<Item = (S, S)>,
328    {
329        self.env = env
330            .into_iter()
331            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
332            .collect();
333        self
334    }
335}
336
337// TODO(team): I think we can rename it to FileMap?
338#[derive(Debug, Clone, Serialize, Deserialize)]
339pub struct TransferedFile {
340    pub local_path: PathBuf,
341    pub remote_path: PathBuf,
342    // TODO: Can be narrowed to have strict typing on this?
343    pub mode: String,
344}
345
346impl TransferedFile {
347    pub fn new<P>(local_path: P, remote_path: P) -> Self
348    where
349        P: AsRef<Path>,
350    {
351        Self {
352            local_path: local_path.as_ref().into(),
353            remote_path: remote_path.as_ref().into(),
354            mode: "0644".to_string(), // default to rw-r--r--
355        }
356    }
357
358    pub fn mode<S>(mut self, mode: S) -> Self
359    where
360        S: AsRef<str>,
361    {
362        self.mode = mode.as_ref().to_string();
363        self
364    }
365}
366
367impl std::fmt::Display for TransferedFile {
368    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
369        write!(
370            f,
371            "File to transfer (local: {}, remote: {})",
372            self.local_path.display(),
373            self.remote_path.display()
374        )
375    }
376}
377
378/// Result of [`ProviderNode::snapshot_db`].
379/// return the filename, hash and bytes of the generated `{node.name}.tgz`
380pub struct InnerSnapshotDb {
381    /// local filename (scoped in node base_dir)
382    pub filename: String,
383    /// Hex-encoded SHA-256 of the archive contents.
384    pub sha256: String,
385    /// Size of the archive in bytes.
386    pub size: u64,
387}