zombienet_provider/shared/
types.rs

1use std::{
2    collections::HashMap,
3    path::{Path, PathBuf},
4    process::ExitStatus,
5};
6
7use configuration::{shared::resources::Resources, types::AssetLocation};
8use serde::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 (should be a tgz file)
50    /// Could be a local or remote asset
51    pub db_snapshot: Option<AssetLocation>,
52    pub port_mapping: Option<HashMap<Port, Port>>,
53}
54
55impl SpawnNodeOptions {
56    pub fn new<S>(name: S, program: S) -> Self
57    where
58        S: AsRef<str>,
59    {
60        Self {
61            name: name.as_ref().to_string(),
62            image: None,
63            resources: None,
64            program: program.as_ref().to_string(),
65            args: vec![],
66            env: vec![],
67            injected_files: vec![],
68            created_paths: vec![],
69            db_snapshot: None,
70            port_mapping: None,
71        }
72    }
73
74    pub fn image<S>(mut self, image: S) -> Self
75    where
76        S: AsRef<str>,
77    {
78        self.image = Some(image.as_ref().to_string());
79        self
80    }
81
82    pub fn resources(mut self, resources: Resources) -> Self {
83        self.resources = Some(resources);
84        self
85    }
86
87    pub fn db_snapshot(mut self, db_snap: Option<AssetLocation>) -> Self {
88        self.db_snapshot = db_snap;
89        self
90    }
91
92    pub fn args<S, I>(mut self, args: I) -> Self
93    where
94        S: AsRef<str>,
95        I: IntoIterator<Item = S>,
96    {
97        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
98        self
99    }
100
101    pub fn env<S, I>(mut self, env: I) -> Self
102    where
103        S: AsRef<str>,
104        I: IntoIterator<Item = (S, S)>,
105    {
106        self.env = env
107            .into_iter()
108            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
109            .collect();
110        self
111    }
112
113    pub fn injected_files<I>(mut self, injected_files: I) -> Self
114    where
115        I: IntoIterator<Item = TransferedFile>,
116    {
117        self.injected_files = injected_files.into_iter().collect();
118        self
119    }
120
121    pub fn created_paths<P, I>(mut self, created_paths: I) -> Self
122    where
123        P: AsRef<Path>,
124        I: IntoIterator<Item = P>,
125    {
126        self.created_paths = created_paths
127            .into_iter()
128            .map(|path| path.as_ref().into())
129            .collect();
130        self
131    }
132
133    pub fn port_mapping(mut self, ports: HashMap<Port, Port>) -> Self {
134        self.port_mapping = Some(ports);
135        self
136    }
137}
138
139#[derive(Debug)]
140pub struct GenerateFileCommand {
141    pub program: String,
142    pub args: Vec<String>,
143    pub env: Vec<(String, String)>,
144    pub local_output_path: PathBuf,
145}
146
147impl GenerateFileCommand {
148    pub fn new<S, P>(program: S, local_output_path: P) -> Self
149    where
150        S: AsRef<str>,
151        P: AsRef<Path>,
152    {
153        Self {
154            program: program.as_ref().to_string(),
155            args: vec![],
156            env: vec![],
157            local_output_path: local_output_path.as_ref().into(),
158        }
159    }
160
161    pub fn args<S, I>(mut self, args: I) -> Self
162    where
163        S: AsRef<str>,
164        I: IntoIterator<Item = S>,
165    {
166        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
167        self
168    }
169
170    pub fn env<S, I>(mut self, env: I) -> Self
171    where
172        S: AsRef<str>,
173        I: IntoIterator<Item = (S, S)>,
174    {
175        self.env = env
176            .into_iter()
177            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
178            .collect();
179        self
180    }
181}
182
183#[derive(Debug)]
184pub struct GenerateFilesOptions {
185    pub commands: Vec<GenerateFileCommand>,
186    pub image: Option<String>,
187    pub injected_files: Vec<TransferedFile>,
188    // Allow to control the name of the node used to create the files.
189    pub temp_name: Option<String>,
190}
191
192impl GenerateFilesOptions {
193    pub fn new<I>(commands: I, image: Option<String>) -> Self
194    where
195        I: IntoIterator<Item = GenerateFileCommand>,
196    {
197        Self {
198            commands: commands.into_iter().collect(),
199            injected_files: vec![],
200            image,
201            temp_name: None,
202        }
203    }
204
205    pub fn with_files<I>(
206        commands: I,
207        image: Option<String>,
208        injected_files: &[TransferedFile],
209    ) -> Self
210    where
211        I: IntoIterator<Item = GenerateFileCommand>,
212    {
213        Self {
214            commands: commands.into_iter().collect(),
215            injected_files: injected_files.into(),
216            image,
217            temp_name: None,
218        }
219    }
220
221    pub fn image<S>(mut self, image: S) -> Self
222    where
223        S: AsRef<str>,
224    {
225        self.image = Some(image.as_ref().to_string());
226        self
227    }
228
229    pub fn injected_files<I>(mut self, injected_files: I) -> Self
230    where
231        I: IntoIterator<Item = TransferedFile>,
232    {
233        self.injected_files = injected_files.into_iter().collect();
234        self
235    }
236
237    pub fn temp_name(mut self, name: impl Into<String>) -> Self {
238        self.temp_name = Some(name.into());
239        self
240    }
241}
242
243#[derive(Debug)]
244pub struct RunCommandOptions {
245    pub program: String,
246    pub args: Vec<String>,
247    pub env: Vec<(String, String)>,
248}
249
250impl RunCommandOptions {
251    pub fn new<S>(program: S) -> Self
252    where
253        S: AsRef<str>,
254    {
255        Self {
256            program: program.as_ref().to_string(),
257            args: vec![],
258            env: vec![],
259        }
260    }
261
262    pub fn args<S, I>(mut self, args: I) -> Self
263    where
264        S: AsRef<str>,
265        I: IntoIterator<Item = S>,
266    {
267        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
268        self
269    }
270
271    pub fn env<S, I>(mut self, env: I) -> Self
272    where
273        S: AsRef<str>,
274        I: IntoIterator<Item = (S, S)>,
275    {
276        self.env = env
277            .into_iter()
278            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
279            .collect();
280        self
281    }
282}
283
284pub struct RunScriptOptions {
285    pub local_script_path: PathBuf,
286    pub args: Vec<String>,
287    pub env: Vec<(String, String)>,
288}
289
290impl RunScriptOptions {
291    pub fn new<P>(local_script_path: P) -> Self
292    where
293        P: AsRef<Path>,
294    {
295        Self {
296            local_script_path: local_script_path.as_ref().into(),
297            args: vec![],
298            env: vec![],
299        }
300    }
301
302    pub fn args<S, I>(mut self, args: I) -> Self
303    where
304        S: AsRef<str>,
305        I: IntoIterator<Item = S>,
306    {
307        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
308        self
309    }
310
311    pub fn env<S, I>(mut self, env: I) -> Self
312    where
313        S: AsRef<str>,
314        I: IntoIterator<Item = (S, S)>,
315    {
316        self.env = env
317            .into_iter()
318            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
319            .collect();
320        self
321    }
322}
323
324// TODO(team): I think we can rename it to FileMap?
325#[derive(Debug, Clone, Serialize)]
326pub struct TransferedFile {
327    pub local_path: PathBuf,
328    pub remote_path: PathBuf,
329    // TODO: Can be narrowed to have strict typing on this?
330    pub mode: String,
331}
332
333impl TransferedFile {
334    pub fn new<P>(local_path: P, remote_path: P) -> Self
335    where
336        P: AsRef<Path>,
337    {
338        Self {
339            local_path: local_path.as_ref().into(),
340            remote_path: remote_path.as_ref().into(),
341            mode: "0644".to_string(), // default to rw-r--r--
342        }
343    }
344
345    pub fn mode<S>(mut self, mode: S) -> Self
346    where
347        S: AsRef<str>,
348    {
349        self.mode = mode.as_ref().to_string();
350        self
351    }
352}
353
354impl std::fmt::Display for TransferedFile {
355    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
356        write!(
357            f,
358            "File to transfer (local: {}, remote: {})",
359            self.local_path.display(),
360            self.remote_path.display()
361        )
362    }
363}