1use crate::default_network;
2use std::{error::Error, fmt, path::Path, process::Command};
3
4#[derive(Clone, Debug, Eq, PartialEq)]
9pub struct DfxRawOutput {
10 pub success: bool,
11 pub status: String,
12 pub stdout: Vec<u8>,
13 pub stderr: Vec<u8>,
14}
15
16#[derive(Debug)]
21pub enum DfxCommandError {
22 Io(std::io::Error),
23 Failed { command: String, stderr: String },
24 SnapshotIdUnavailable { output: String },
25}
26
27impl fmt::Display for DfxCommandError {
28 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
30 match self {
31 Self::Io(err) => write!(formatter, "{err}"),
32 Self::Failed { command, stderr } => {
33 write!(formatter, "dfx command failed: {command}\n{stderr}")
34 }
35 Self::SnapshotIdUnavailable { output } => {
36 write!(
37 formatter,
38 "could not parse snapshot id from dfx output: {output}"
39 )
40 }
41 }
42 }
43}
44
45impl Error for DfxCommandError {
46 fn source(&self) -> Option<&(dyn Error + 'static)> {
48 match self {
49 Self::Io(err) => Some(err),
50 Self::Failed { .. } | Self::SnapshotIdUnavailable { .. } => None,
51 }
52 }
53}
54
55impl From<std::io::Error> for DfxCommandError {
56 fn from(err: std::io::Error) -> Self {
58 Self::Io(err)
59 }
60}
61
62#[derive(Clone, Debug, Eq, PartialEq)]
67pub struct Dfx {
68 executable: String,
69 network: Option<String>,
70}
71
72impl Dfx {
73 #[must_use]
75 pub fn new(executable: impl Into<String>, network: Option<String>) -> Self {
76 Self {
77 executable: executable.into(),
78 network,
79 }
80 }
81
82 #[must_use]
84 pub fn network(&self) -> Option<&str> {
85 self.network.as_deref()
86 }
87
88 #[must_use]
90 pub fn canister_command(&self) -> Command {
91 let mut command = Command::new(&self.executable);
92 command.arg("canister");
93 add_network_args(&mut command, self.network());
94 command
95 }
96
97 pub fn ping(&self) -> Result<(), DfxCommandError> {
99 let mut command = Command::new(&self.executable);
100 command.arg("ping");
101 let network = self.network().map_or_else(default_network, str::to_string);
102 command.arg(network);
103 run_status(&mut command)
104 }
105
106 pub fn canister_id_optional(&self, name: &str) -> Result<Option<String>, DfxCommandError> {
108 let mut command = self.canister_command();
109 command.args(["id", name]);
110 match run_output(&mut command) {
111 Ok(output) => Ok(Some(output)),
112 Err(DfxCommandError::Failed { command, stderr }) if canister_id_missing(&stderr) => {
113 let _ = command;
114 Ok(None)
115 }
116 Err(err) => Err(err),
117 }
118 }
119
120 pub fn canister_id(&self, name: &str) -> Result<String, DfxCommandError> {
122 let mut command = self.canister_command();
123 command.args(["id", name]);
124 run_output(&mut command)
125 }
126
127 pub fn canister_call_output(
129 &self,
130 canister: &str,
131 method: &str,
132 output: Option<&str>,
133 ) -> Result<String, DfxCommandError> {
134 let mut command = self.canister_command();
135 command.args(["call", canister, method]);
136 if let Some(output) = output {
137 command.args(["--output", output]);
138 }
139 run_output(&mut command)
140 }
141
142 pub fn snapshot_list(&self, canister: &str) -> Result<String, DfxCommandError> {
144 let mut command = self.canister_command();
145 command.args(["snapshot", "list", canister]);
146 run_output(&mut command)
147 }
148
149 pub fn snapshot_create(&self, canister: &str) -> Result<String, DfxCommandError> {
151 let mut command = self.canister_command();
152 command.args(["snapshot", "create", canister]);
153 run_output_with_stderr(&mut command)
154 }
155
156 pub fn snapshot_create_id(&self, canister: &str) -> Result<String, DfxCommandError> {
158 let before = self.snapshot_list_ids(canister)?;
159 let output = self.snapshot_create(canister)?;
160 if let Some(snapshot_id) = parse_snapshot_id(&output) {
161 return Ok(snapshot_id);
162 }
163
164 let before = before
165 .into_iter()
166 .collect::<std::collections::BTreeSet<_>>();
167 let mut new_ids = self
168 .snapshot_list_ids(canister)?
169 .into_iter()
170 .filter(|snapshot_id| !before.contains(snapshot_id))
171 .collect::<Vec<_>>();
172 if new_ids.len() == 1 {
173 Ok(new_ids.remove(0))
174 } else {
175 Err(DfxCommandError::SnapshotIdUnavailable { output })
176 }
177 }
178
179 pub fn snapshot_list_ids(&self, canister: &str) -> Result<Vec<String>, DfxCommandError> {
181 let output = self.snapshot_list(canister)?;
182 Ok(parse_snapshot_list_ids(&output))
183 }
184
185 pub fn stop_canister(&self, canister: &str) -> Result<(), DfxCommandError> {
187 let mut command = self.canister_command();
188 command.args(["stop", canister]);
189 run_status(&mut command)
190 }
191
192 pub fn start_canister(&self, canister: &str) -> Result<(), DfxCommandError> {
194 let mut command = self.canister_command();
195 command.args(["start", canister]);
196 run_status(&mut command)
197 }
198
199 pub fn snapshot_download(
201 &self,
202 canister: &str,
203 snapshot_id: &str,
204 artifact_path: &Path,
205 ) -> Result<(), DfxCommandError> {
206 let mut command = self.canister_command();
207 command.args(["snapshot", "download", canister, snapshot_id, "--dir"]);
208 command.arg(artifact_path);
209 run_status(&mut command)
210 }
211
212 #[must_use]
214 pub fn snapshot_create_display(&self, canister: &str) -> String {
215 let mut command = self.canister_command();
216 command.args(["snapshot", "create", canister]);
217 command_display(&command)
218 }
219
220 #[must_use]
222 pub fn snapshot_download_display(
223 &self,
224 canister: &str,
225 snapshot_id: &str,
226 artifact_path: &Path,
227 ) -> String {
228 let mut command = self.canister_command();
229 command.args(["snapshot", "download", canister, snapshot_id, "--dir"]);
230 command.arg(artifact_path);
231 command_display(&command)
232 }
233
234 #[must_use]
236 pub fn stop_canister_display(&self, canister: &str) -> String {
237 let mut command = self.canister_command();
238 command.args(["stop", canister]);
239 command_display(&command)
240 }
241
242 #[must_use]
244 pub fn start_canister_display(&self, canister: &str) -> String {
245 let mut command = self.canister_command();
246 command.args(["start", canister]);
247 command_display(&command)
248 }
249}
250
251pub fn add_network_args(command: &mut Command, network: Option<&str>) {
253 if let Some(network) = network {
254 command.args(["--network", network]);
255 }
256}
257
258pub fn run_output(command: &mut Command) -> Result<String, DfxCommandError> {
260 let display = command_display(command);
261 let output = command.output()?;
262 if output.status.success() {
263 Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
264 } else {
265 Err(DfxCommandError::Failed {
266 command: display,
267 stderr: command_stderr(&output),
268 })
269 }
270}
271
272pub fn run_output_with_stderr(command: &mut Command) -> Result<String, DfxCommandError> {
274 let display = command_display(command);
275 let output = command.output()?;
276 if output.status.success() {
277 let mut text = String::from_utf8_lossy(&output.stdout).to_string();
278 text.push_str(&String::from_utf8_lossy(&output.stderr));
279 Ok(text.trim().to_string())
280 } else {
281 Err(DfxCommandError::Failed {
282 command: display,
283 stderr: command_stderr(&output),
284 })
285 }
286}
287
288pub fn run_status(command: &mut Command) -> Result<(), DfxCommandError> {
290 let display = command_display(command);
291 let output = command.output()?;
292 if output.status.success() {
293 Ok(())
294 } else {
295 Err(DfxCommandError::Failed {
296 command: display,
297 stderr: command_stderr(&output),
298 })
299 }
300}
301
302pub fn run_raw_output(program: &str, args: &[String]) -> Result<DfxRawOutput, std::io::Error> {
304 let output = Command::new(program).args(args).output()?;
305 Ok(DfxRawOutput {
306 success: output.status.success(),
307 status: exit_status_label(output.status),
308 stdout: output.stdout,
309 stderr: output.stderr,
310 })
311}
312
313#[must_use]
315pub fn command_display(command: &Command) -> String {
316 let mut parts = vec![command.get_program().to_string_lossy().to_string()];
317 parts.extend(
318 command
319 .get_args()
320 .map(|arg| arg.to_string_lossy().to_string()),
321 );
322 parts.join(" ")
323}
324
325#[must_use]
327pub fn canister_id_missing(stderr: &str) -> bool {
328 stderr.contains("Cannot find canister id")
329}
330
331#[must_use]
333pub fn parse_snapshot_id(output: &str) -> Option<String> {
334 output
335 .split(|c: char| c.is_whitespace() || matches!(c, '"' | '\'' | ':' | ','))
336 .filter(|part| !part.is_empty())
337 .rev()
338 .find(|part| {
339 part.chars()
340 .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
341 })
342 .map(str::to_string)
343}
344
345#[must_use]
347pub fn parse_snapshot_list_ids(output: &str) -> Vec<String> {
348 output
349 .lines()
350 .filter_map(|line| {
351 line.split_once(':')
352 .map(|(snapshot_id, _)| snapshot_id.trim())
353 })
354 .filter(|snapshot_id| !snapshot_id.is_empty())
355 .map(str::to_string)
356 .collect()
357}
358
359fn command_stderr(output: &std::process::Output) -> String {
361 let stderr = String::from_utf8_lossy(&output.stderr);
362 if stderr.trim().is_empty() {
363 String::from_utf8_lossy(&output.stdout).to_string()
364 } else {
365 stderr.to_string()
366 }
367}
368
369fn exit_status_label(status: std::process::ExitStatus) -> String {
371 status
372 .code()
373 .map_or_else(|| "signal".to_string(), |code| code.to_string())
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379
380 #[test]
382 fn parses_snapshot_id_from_output() {
383 let snapshot_id = parse_snapshot_id("Created snapshot: snap_abc-123\n");
384
385 assert_eq!(snapshot_id.as_deref(), Some("snap_abc-123"));
386 }
387
388 #[test]
390 fn parses_snapshot_ids_from_list_output() {
391 let snapshot_ids = parse_snapshot_list_ids(
392 "0000000000000000ffffffffff9000050101: size 10\n\
393 0000000000000000ffffffffff9000050102: size 12\n",
394 );
395
396 assert_eq!(
397 snapshot_ids,
398 vec![
399 "0000000000000000ffffffffff9000050101",
400 "0000000000000000ffffffffff9000050102"
401 ]
402 );
403 }
404}