commandspec/
lib.rs

1extern crate shlex;
2#[macro_use]
3extern crate failure;
4#[macro_use]
5extern crate lazy_static;
6#[macro_use]
7extern crate log;
8
9#[cfg(windows)]
10extern crate kernel32;
11#[cfg(unix)]
12extern crate nix;
13#[cfg(windows)]
14extern crate winapi;
15
16use std::process::Command;
17use std::fmt;
18use std::collections::HashMap;
19use std::sync::Arc;
20use std::sync::Mutex;
21use std::path::{Path, PathBuf};
22use std::process::Stdio;
23
24// Re-export for macros.
25pub use failure::Error;
26
27pub mod macros;
28mod process;
29mod signal;
30
31use process::Process;
32use signal::Signal;
33
34lazy_static! {
35    static ref PID_MAP: Arc<Mutex<HashMap<i32, Process>>> = Arc::new(Mutex::new(HashMap::new()));
36}
37
38pub fn disable_cleanup_on_ctrlc() {
39    signal::uninstall_handler();
40}
41
42pub fn cleanup_on_ctrlc() {
43    signal::install_handler(move |sig: Signal| {
44        match sig {
45            // SIGCHLD is special, initiate reap()
46            Signal::SIGCHLD => {
47                for (_pid, process) in PID_MAP.lock().unwrap().iter() {
48                    process.reap();
49                }
50            }
51            Signal::SIGINT => {
52                for (_pid, process) in PID_MAP.lock().unwrap().iter() {
53                    process.signal(sig);
54                }
55                ::std::process::exit(130);
56            }
57            _ => {
58                for (_pid, process) in PID_MAP.lock().unwrap().iter() {
59                    process.signal(sig);
60                }
61            }
62        }
63    });
64}
65
66pub struct SpawnGuard(i32);
67
68impl ::std::ops::Drop for SpawnGuard {
69    fn drop(&mut self) {
70        PID_MAP.lock().unwrap().remove(&self.0).map(|process| process.reap());
71    }
72}
73
74//---------------
75
76pub trait CommandSpecExt {
77    fn execute(self) -> Result<(), CommandError>;
78
79    fn scoped_spawn(self) -> Result<SpawnGuard, ::std::io::Error>;
80}
81
82#[derive(Debug, Fail)]
83pub enum CommandError {
84    #[fail(display = "Encountered an IO error: {:?}", _0)]
85    Io(#[cause] ::std::io::Error),
86
87    #[fail(display = "Command was interrupted.")]
88    Interrupt,
89
90    #[fail(display = "Command failed with error code {}.", _0)]
91    Code(i32),
92}
93
94impl CommandError {
95    /// Returns the error code this command failed with. Can panic if not a `Code`.
96    pub fn error_code(&self) -> i32 {
97        if let CommandError::Code(value) = *self {
98            value
99        } else {
100            panic!("Called error_code on a value that was not a CommandError::Code")
101        }
102    }
103}
104
105impl CommandSpecExt for Command {
106    // Executes the command, and returns a comprehensive error type
107    fn execute(mut self) -> Result<(), CommandError> {
108        self.stdout(Stdio::inherit());
109        self.stderr(Stdio::inherit());
110        match self.output() {
111            Ok(output) => {
112                if output.status.success() {
113                    Ok(())
114                } else if let Some(code) = output.status.code() {
115                    Err(CommandError::Code(code))
116                } else {
117                    Err(CommandError::Interrupt)
118                }
119            },
120            Err(err) => Err(CommandError::Io(err)),
121        }
122    }
123
124    fn scoped_spawn(self) -> Result<SpawnGuard, ::std::io::Error> {
125        let process = Process::new(self)?;
126        let id = process.id();
127        PID_MAP.lock().unwrap().insert(id, process);
128        Ok(SpawnGuard(id))
129    }
130}
131
132//---------------
133
134pub enum CommandArg {
135    Empty,
136    Literal(String),
137    List(Vec<String>),
138}
139
140fn shell_quote(value: &str) -> String {
141    shlex::quote(&format!("{}", value)).to_string()
142}
143
144impl fmt::Display for CommandArg {
145    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
146        use CommandArg::*;
147        match *self {
148            Empty => write!(f, ""),
149            Literal(ref value) => {
150                write!(f, "{}", shell_quote(&format!("{}", value)))
151            },
152            List(ref list) => {
153                write!(f, "{}", list
154                    .iter()
155                    .map(|x| shell_quote(&format!("{}", x)).to_string())
156                    .collect::<Vec<_>>()
157                    .join(" "))
158            }
159        }
160    }
161}
162
163impl<'a, 'b> From<&'a &'b str> for CommandArg {
164    fn from(value: &&str) -> Self {
165        CommandArg::Literal(value.to_string())
166    }
167}
168
169impl From<String> for CommandArg {
170    fn from(value: String) -> Self {
171        CommandArg::Literal(value)
172    }
173}
174
175impl<'a> From<&'a String> for CommandArg {
176    fn from(value: &String) -> Self {
177        CommandArg::Literal(value.to_string())
178    }
179}
180
181
182impl<'a> From<&'a str> for CommandArg {
183    fn from(value: &str) -> Self {
184        CommandArg::Literal(value.to_string())
185    }
186}
187
188impl<'a> From<&'a u64> for CommandArg {
189    fn from(value: &u64) -> Self {
190        CommandArg::Literal(value.to_string())
191    }
192}
193
194impl<'a> From<&'a f64> for CommandArg {
195    fn from(value: &f64) -> Self {
196        CommandArg::Literal(value.to_string())
197    }
198}
199
200impl<'a> From<&'a i32> for CommandArg {
201    fn from(value: &i32) -> Self {
202        CommandArg::Literal(value.to_string())
203    }
204}
205
206impl<'a> From<&'a i64> for CommandArg {
207    fn from(value: &i64) -> Self {
208        CommandArg::Literal(value.to_string())
209    }
210}
211
212impl<'a, T> From<&'a [T]> for CommandArg
213    where T: fmt::Display {
214    fn from(list: &[T]) -> Self {
215        CommandArg::List(
216            list
217                .iter()
218                .map(|x| format!("{}", x))
219                .collect()
220        )
221    }
222}
223
224impl<'a, T> From<&'a Vec<T>> for CommandArg
225    where T: fmt::Display {
226    fn from(list: &Vec<T>) -> Self {
227        CommandArg::from(list.as_slice())
228    }
229}
230
231impl<'a, T> From<&'a Option<T>> for CommandArg
232    where T: fmt::Display {
233    fn from(opt: &Option<T>) -> Self {
234        if let Some(ref value) = *opt {
235            CommandArg::Literal(format!("{}", value))
236        } else {
237            CommandArg::Empty
238        }
239    }
240}
241
242pub fn command_arg<'a, T>(value: &'a T) -> CommandArg
243    where CommandArg: std::convert::From<&'a T> {
244    CommandArg::from(value)
245}
246
247//---------------
248
249/// Represents the invocation specification used to generate a Command.
250#[derive(Debug)]
251struct CommandSpec {
252    binary: String,
253    args: Vec<String>,
254    env: HashMap<String, String>,
255    cd: Option<String>,
256}
257
258impl CommandSpec {
259    fn to_command(&self) -> Command {
260        let cd = if let Some(ref cd) = self.cd {
261            canonicalize_path(Path::new(cd)).unwrap()
262        } else {
263            ::std::env::current_dir().unwrap()
264        };
265        let mut binary = Path::new(&self.binary).to_owned();
266
267        // On Windows, current_dir takes place after binary name resolution.
268        // If current_dir is specified and the binary is referenced by a relative path,
269        // add the dir change to its relative path.
270        // https://github.com/rust-lang/rust/issues/37868
271        if cfg!(windows) && binary.is_relative() && binary.components().count() != 1 {
272            binary = cd.join(&binary);
273        }
274
275        // On windows, we run in cmd.exe by default. (This code is a naive way
276        // of accomplishing this and may contain errors.)
277        if cfg!(windows) {
278            let mut cmd = Command::new("cmd");
279            cmd.current_dir(cd);
280            let invoke_string = format!("{} {}", binary.as_path().to_string_lossy(), self.args.join(" "));
281            cmd.args(&["/C", &invoke_string]);
282            for (key, value) in &self.env {
283                cmd.env(key, value);
284            }
285            return cmd;
286        }
287
288        let mut cmd = Command::new(binary);
289        cmd.current_dir(cd);
290        cmd.args(&self.args);
291        for (key, value) in &self.env {
292            cmd.env(key, value);
293        }
294        cmd
295    }
296}
297
298// Strips UNC from canonicalized paths.
299// See https://github.com/rust-lang/rust/issues/42869 for why this is needed.
300#[cfg(windows)]
301fn canonicalize_path<'p, P>(path: P) -> Result<PathBuf, Error>
302where P: Into<&'p Path> {
303    use std::ffi::OsString;
304    use std::os::windows::prelude::*;
305
306    let canonical = path.into().canonicalize()?;
307    let vec_chars = canonical.as_os_str().encode_wide().collect::<Vec<u16>>();
308    if vec_chars[0..4] == [92, 92, 63, 92] {
309        return Ok(Path::new(&OsString::from_wide(&vec_chars[4..])).to_owned());
310    }
311
312    Ok(canonical)
313}
314
315#[cfg(not(windows))]
316fn canonicalize_path<'p, P>(path: P) -> Result<PathBuf, Error>
317where P: Into<&'p Path> {
318    Ok(path.into().canonicalize()?)
319}
320
321//---------------
322
323pub fn commandify(value: String) -> Result<Command, Error> {
324    let lines = value.trim().split("\n").map(String::from).collect::<Vec<_>>();
325
326    #[derive(Debug, PartialEq)]
327    enum SpecState {
328        Cd,
329        Env,
330        Cmd,
331    }
332
333    let mut env = HashMap::<String, String>::new();
334    let mut cd = None;
335
336    let mut state = SpecState::Cd;
337    let mut command_lines = vec![];
338    for raw_line in lines {
339        let mut line = shlex::split(&raw_line).unwrap_or(vec![]);
340        if state == SpecState::Cmd {
341            command_lines.push(raw_line);
342        } else {
343            if raw_line.trim().is_empty() {
344                continue;
345            }
346
347            match line.get(0).map(|x| x.as_ref()) {
348                Some("cd") => {
349                    if state != SpecState::Cd {
350                        bail!("cd should be the first line in your command! macro.");
351                    }
352                    ensure!(line.len() == 2, "Too many arguments in cd; expected 1, found {}", line.len() - 1);
353                    cd = Some(line.remove(1));
354                    state = SpecState::Env;
355                }
356                Some("export") => {
357                    if state != SpecState::Cd && state != SpecState::Env {
358                        bail!("exports should follow cd but precede your command in the command! macro.");
359                    }
360                    ensure!(line.len() >= 2, "Not enough arguments in export; expected at least 1, found {}", line.len() - 1);
361                    for item in &line[1..] {
362                        let mut items = item.splitn(2, "=").collect::<Vec<_>>();
363                        ensure!(items.len() > 0, "Expected export of the format NAME=VALUE");
364                        env.insert(items[0].to_string(), items[1].to_string());
365                    }
366                    state = SpecState::Env;
367                }
368                None | Some(_) => {
369                    command_lines.push(raw_line);
370                    state = SpecState::Cmd;
371                }
372            }
373        }
374    }
375    if state != SpecState::Cmd || command_lines.is_empty() {
376        bail!("Didn't find a command in your command! macro.");
377    }
378
379    // Join the command string and split out binary / args.
380    let command_string = command_lines.join("\n").replace("\\\n", "\n");
381    let mut command = shlex::split(&command_string).expect("Command string couldn't be parsed by shlex");
382    let binary = command.remove(0); 
383    let args = command;
384
385    // Generate the CommandSpec struct.
386    let spec = CommandSpec {
387        binary,
388        args,
389        env,
390        cd,
391    };
392
393    // DEBUG
394    // eprintln!("COMMAND: {:?}", spec);
395
396    Ok(spec.to_command())
397}