1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

use std::{
    ffi::OsStr,
    fmt::{self, Write},
    process::{self, Child, Command, Output},
};

use crate::{utf8_to_string_lossy, CommandFailed, InstallError};

pub trait CommandExt {
    fn formattable(&self) -> CommandFormattable<'_>;

    fn output_checked_status(&mut self) -> Result<Output, InstallError>;

    fn spawn_with_cmd(self) -> Result<ChildWithCommand, InstallError>;
}

impl CommandExt for Command {
    fn formattable(&self) -> CommandFormattable<'_> {
        CommandFormattable(self)
    }

    fn output_checked_status(&mut self) -> Result<Output, InstallError> {
        self.output()
            .map_err(InstallError::from)
            .and_then(|output| check_status(self, output))
    }

    fn spawn_with_cmd(mut self) -> Result<ChildWithCommand, InstallError> {
        self.spawn()
            .map_err(InstallError::from)
            .map(move |child| ChildWithCommand { child, cmd: self })
    }
}

pub struct CommandFormattable<'a>(&'a Command);

fn needs_escape(s: &str) -> bool {
    s.contains(|ch| !matches!(ch, 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' | '=' | '/' | ',' | '.' | '+'))
}

fn write_shell_arg_escaped(f: &mut fmt::Formatter<'_>, os_str: &OsStr) -> fmt::Result {
    let s = os_str.to_string_lossy();

    if needs_escape(&s) {
        // There is some ascii whitespace (' ', '\n', '\t'),
        // or non-ascii characters need to quote them using `"`.
        //
        // But then, it is possible for the `s` to contains `"`,
        // so they needs to be escaped.
        f.write_str("\"")?;

        for ch in s.chars() {
            if ch == '"' {
                // Escape it with `\`.
                f.write_char('\\')?;
            }

            f.write_char(ch)?;
        }

        f.write_str("\"")
    } else {
        f.write_str(&s)
    }
}

impl fmt::Display for CommandFormattable<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let cmd = self.0;

        write_shell_arg_escaped(f, cmd.get_program())?;

        for arg in cmd.get_args() {
            f.write_str(" ")?;
            write_shell_arg_escaped(f, arg)?;
        }

        Ok(())
    }
}

pub struct ChildWithCommand {
    cmd: Command,
    child: Child,
}

fn check_status(cmd: &Command, output: Output) -> Result<Output, InstallError> {
    if output.status.success() {
        Ok(output)
    } else {
        Err(CommandFailed {
            command: cmd.formattable().to_string(),
            stdout: utf8_to_string_lossy(output.stdout),
            stderr: utf8_to_string_lossy(output.stderr),
        }
        .into())
    }
}

impl ChildWithCommand {
    pub fn wait_with_output_checked_status(self) -> Result<Output, InstallError> {
        let cmd = self.cmd;

        self.child
            .wait_with_output()
            .map_err(InstallError::from)
            .and_then(|output| check_status(&cmd, output))
    }

    pub fn stdin(&mut self) -> &mut Option<process::ChildStdin> {
        &mut self.child.stdin
    }

    pub fn stdout(&mut self) -> &mut Option<process::ChildStdout> {
        &mut self.child.stdout
    }

    pub fn stderr(&mut self) -> &mut Option<process::ChildStderr> {
        &mut self.child.stderr
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cmd_format() {
        assert_eq!(
            Command::new("cargo")
                .args(["binstall", "-V"])
                .formattable()
                .to_string(),
            "cargo binstall -V"
        );
    }
}