cli/util/
command.rs

1/*---------------------------------------------------------------------------------------------
2 *  Copyright (c) Microsoft Corporation. All rights reserved.
3 *  Licensed under the MIT License. See License.txt in the project root for license information.
4 *--------------------------------------------------------------------------------------------*/
5use super::errors::CodeError;
6use std::{
7	borrow::Cow,
8	ffi::OsStr,
9	process::{Output, Stdio},
10};
11use tokio::process::Command;
12
13pub async fn capture_command_and_check_status(
14	command_str: impl AsRef<OsStr>,
15	args: &[impl AsRef<OsStr>],
16) -> Result<std::process::Output, CodeError> {
17	let output = capture_command(&command_str, args).await?;
18
19	check_output_status(output, || {
20		format!(
21			"{} {}",
22			command_str.as_ref().to_string_lossy(),
23			args.iter()
24				.map(|a| a.as_ref().to_string_lossy())
25				.collect::<Vec<Cow<'_, str>>>()
26				.join(" ")
27		)
28	})
29}
30
31pub fn check_output_status(
32	output: Output,
33	cmd_str: impl FnOnce() -> String,
34) -> Result<std::process::Output, CodeError> {
35	if !output.status.success() {
36		return Err(CodeError::CommandFailed {
37			command: cmd_str(),
38			code: output.status.code().unwrap_or(-1),
39			output: String::from_utf8_lossy(if output.stderr.is_empty() {
40				&output.stdout
41			} else {
42				&output.stderr
43			})
44			.into(),
45		});
46	}
47
48	Ok(output)
49}
50
51pub async fn capture_command<A, I, S>(
52	command_str: A,
53	args: I,
54) -> Result<std::process::Output, CodeError>
55where
56	A: AsRef<OsStr>,
57	I: IntoIterator<Item = S>,
58	S: AsRef<OsStr>,
59{
60	new_tokio_command(&command_str)
61		.args(args)
62		.stdin(Stdio::null())
63		.stdout(Stdio::piped())
64		.output()
65		.await
66		.map_err(|e| CodeError::CommandFailed {
67			command: command_str.as_ref().to_string_lossy().to_string(),
68			code: -1,
69			output: e.to_string(),
70		})
71}
72
73/// Makes a new Command, setting flags to avoid extra windows on win32
74#[cfg(windows)]
75pub fn new_tokio_command(exe: impl AsRef<OsStr>) -> Command {
76	let mut p = tokio::process::Command::new(exe);
77	p.creation_flags(winapi::um::winbase::CREATE_NO_WINDOW);
78	p
79}
80
81/// Makes a new Command, setting flags to avoid extra windows on win32
82#[cfg(not(windows))]
83pub fn new_tokio_command(exe: impl AsRef<OsStr>) -> Command {
84	tokio::process::Command::new(exe)
85}
86
87/// Makes a new command to run the target script. For windows, ensures it's run
88/// in a cmd.exe context.
89#[cfg(windows)]
90pub fn new_script_command(script: impl AsRef<OsStr>) -> Command {
91	let mut cmd = new_tokio_command("cmd");
92	cmd.arg("/Q");
93	cmd.arg("/C");
94	cmd.arg(script);
95	cmd
96}
97
98/// Makes a new command to run the target script. For windows, ensures it's run
99/// in a cmd.exe context.
100#[cfg(not(windows))]
101pub fn new_script_command(script: impl AsRef<OsStr>) -> Command {
102	new_tokio_command(script) // it's assumed scripts are already +x and don't need extra handling
103}
104
105/// Makes a new Command, setting flags to avoid extra windows on win32
106#[cfg(windows)]
107pub fn new_std_command(exe: impl AsRef<OsStr>) -> std::process::Command {
108	let mut p = std::process::Command::new(exe);
109	std::os::windows::process::CommandExt::creation_flags(
110		&mut p,
111		winapi::um::winbase::CREATE_NO_WINDOW,
112	);
113	p
114}
115
116/// Makes a new Command, setting flags to avoid extra windows on win32
117#[cfg(not(windows))]
118pub fn new_std_command(exe: impl AsRef<OsStr>) -> std::process::Command {
119	std::process::Command::new(exe)
120}
121
122/// Kills and processes and all of its children.
123#[cfg(windows)]
124pub async fn kill_tree(process_id: u32) -> Result<(), CodeError> {
125	capture_command("taskkill", &["/t", "/pid", &process_id.to_string()]).await?;
126	Ok(())
127}
128
129/// Kills and processes and all of its children.
130#[cfg(not(windows))]
131pub async fn kill_tree(process_id: u32) -> Result<(), CodeError> {
132	use futures::future::join_all;
133	use tokio::io::{AsyncBufReadExt, BufReader};
134
135	async fn kill_single_pid(process_id_str: String) {
136		capture_command("kill", &[&process_id_str]).await.ok();
137	}
138
139	// Rusty version of https://github.com/microsoft/vscode-js-debug/blob/main/src/targets/node/terminateProcess.sh
140
141	let parent_id = process_id.to_string();
142	let mut prgrep_cmd = Command::new("pgrep")
143		.arg("-P")
144		.arg(&parent_id)
145		.stdin(Stdio::null())
146		.stdout(Stdio::piped())
147		.spawn()
148		.map_err(|e| CodeError::CommandFailed {
149			command: format!("pgrep -P {}", parent_id),
150			code: -1,
151			output: e.to_string(),
152		})?;
153
154	let mut kill_futures = vec![tokio::spawn(
155		async move { kill_single_pid(parent_id).await },
156	)];
157
158	if let Some(stdout) = prgrep_cmd.stdout.take() {
159		let mut reader = BufReader::new(stdout).lines();
160		while let Some(line) = reader.next_line().await.unwrap_or(None) {
161			kill_futures.push(tokio::spawn(async move { kill_single_pid(line).await }))
162		}
163	}
164
165	join_all(kill_futures).await;
166	prgrep_cmd.kill().await.ok();
167	Ok(())
168}