Skip to main content

nxd_core/adapters/
process.rs

1use crate::config;
2use std::cell::RefCell;
3use std::collections::HashMap;
4use std::collections::HashSet;
5use std::process::{Command, Stdio};
6
7pub use crate::progress::stream::{format_output_line, stream_output_lines, write_output_line};
8pub use crate::progress::target::LogTarget;
9pub use crate::progress::target::Logger;
10
11pub fn sanitized_command(program: impl AsRef<std::ffi::OsStr>) -> Command {
12	let mut command = Command::new(program);
13	command.env_remove("GITHUB_TOKEN").env_remove("GH_TOKEN");
14	command
15}
16
17thread_local! {
18		static COMMAND_RECORDER: RefCell<Option<Vec<String>>> = const { RefCell::new(None) };
19		static COMMAND_MOCKS: RefCell<HashMap<String, String>> = RefCell::new(HashMap::new());
20		static COMMAND_MOCK_ERRORS: RefCell<HashSet<String>> = RefCell::new(HashSet::new());
21}
22
23pub struct CommandExecutor;
24
25impl CommandExecutor {
26	/// Starts recording executed commands and enables dry-run/mock mode.
27	pub fn start_recording() {
28		COMMAND_RECORDER.with(|r| {
29			*r.borrow_mut() = Some(Vec::new());
30		});
31	}
32
33	/// Stops recording executed commands and returns the list of recorded commands.
34	pub fn stop_recording() -> Option<Vec<String>> {
35		COMMAND_RECORDER.with(|r| r.borrow_mut().take())
36	}
37
38	/// Returns whether we are currently in recording/dry-run mode.
39	pub fn is_recording() -> bool {
40		COMMAND_RECORDER.with(|r| r.borrow().is_some())
41	}
42
43	/// Registers a mock stdout response for a given command.
44	pub fn register_mock(cmd: &str, output: &str) {
45		COMMAND_MOCKS.with(|m| {
46			m.borrow_mut().insert(cmd.to_string(), output.to_string());
47		});
48	}
49
50	/// Registers a deterministic failure for a matching recorded command.
51	pub fn register_mock_error(cmd: &str) {
52		COMMAND_MOCK_ERRORS.with(|errors| {
53			errors.borrow_mut().insert(cmd.to_string());
54		});
55	}
56
57	/// Clears all mock command registrations.
58	pub fn clear_mocks() {
59		COMMAND_MOCKS.with(|m| {
60			m.borrow_mut().clear();
61		});
62		COMMAND_MOCK_ERRORS.with(|errors| errors.borrow_mut().clear());
63	}
64
65	/// Executes a command locally, streaming stdout and stderr in real-time.
66	pub fn execute(
67		program: &str,
68		args: &[&str],
69		logger: Logger,
70	) -> Result<String, Box<dyn std::error::Error>> {
71		let cmd_str = format!("{} {}", program, args.join(" "));
72		let cmd_str = crate::config::redact_token(&cmd_str);
73		let is_rec = COMMAND_RECORDER.with(|r| {
74			if let Some(ref mut v) = *r.borrow_mut() {
75				v.push(cmd_str.clone());
76				true
77			} else {
78				false
79			}
80		});
81
82		if is_rec {
83			let fails = COMMAND_MOCK_ERRORS
84				.with(|errors| errors.borrow().iter().any(|pattern| cmd_str.contains(pattern)));
85			if fails {
86				return Err(format!("mocked command failure: {cmd_str}").into());
87			}
88			let mock_out = COMMAND_MOCKS
89				.with(|m| {
90					let map = m.borrow();
91					if let Some(out) = map.get(&cmd_str) {
92						return Some(out.clone());
93					}
94					for (key, val) in map.iter() {
95						if cmd_str.contains(key) {
96							return Some(val.clone());
97						}
98					}
99					None
100				})
101				.unwrap_or_default();
102			return Ok(mock_out);
103		}
104
105		let mut child_cmd = sanitized_command(program);
106		child_cmd.args(args).stdout(Stdio::piped()).stderr(Stdio::piped());
107
108		let mut child = child_cmd.spawn()?;
109
110		let stdout = child.stdout.take().ok_or("Failed to open stdout")?;
111		let stderr = child.stderr.take().ok_or("Failed to open stderr")?;
112
113		let filter_lock_warning = args.iter().any(|arg| arg.contains(config::SECRET_INPUT_NAME));
114
115		let stdout_handle = stream_output_lines(stdout, logger.clone(), false);
116		let stderr_handle = crate::progress::stream::stream_output_lines_filtered(
117			stderr,
118			logger.clone(),
119			true,
120			filter_lock_warning,
121		);
122
123		let stdout_output = stdout_handle.join().unwrap_or_default();
124		let stderr_output = stderr_handle.join().unwrap_or_default();
125
126		let status = child.wait()?;
127		if !status.success() {
128			let redacted_args = crate::config::redact_token(&args.join(" "));
129			let redacted_stderr = crate::config::redact_token(&stderr_output);
130			return Err(
131				format!(
132					"Command '{} {}' failed with status: {}\nError: {}",
133					program, redacted_args, status, redacted_stderr
134				)
135				.into(),
136			);
137		}
138
139		Ok(stdout_output)
140	}
141
142	/// Executes a command locally, feeding content to its stdin.
143	pub fn execute_with_stdin(
144		program: &str,
145		args: &[&str],
146		stdin_content: &str,
147		logger: Logger,
148	) -> Result<String, Box<dyn std::error::Error>> {
149		let safe_stdin =
150			if stdin_content.contains("PRIVATE KEY") { "<private key hidden>" } else { stdin_content };
151		let cmd_str = format!("{} {} [stdin: {}]", program, args.join(" "), safe_stdin);
152		let cmd_str = crate::config::redact_token(&cmd_str);
153		let is_rec = COMMAND_RECORDER.with(|r| {
154			if let Some(ref mut v) = *r.borrow_mut() {
155				v.push(cmd_str.clone());
156				true
157			} else {
158				false
159			}
160		});
161
162		if is_rec {
163			let mock_out = COMMAND_MOCKS
164				.with(|m| {
165					let map = m.borrow();
166					if let Some(out) = map.get(&cmd_str) {
167						return Some(out.clone());
168					}
169					for (key, val) in map.iter() {
170						if cmd_str.contains(key) {
171							return Some(val.clone());
172						}
173					}
174					None
175				})
176				.unwrap_or_default();
177			return Ok(mock_out);
178		}
179
180		let mut child = sanitized_command(program)
181			.args(args)
182			.stdin(Stdio::piped())
183			.stdout(Stdio::piped())
184			.stderr(Stdio::piped())
185			.spawn()?;
186
187		if let Some(mut stdin) = child.stdin.take() {
188			use std::io::Write as _;
189			stdin.write_all(stdin_content.as_bytes())?;
190			stdin.flush()?;
191		}
192
193		let stdout = child.stdout.take().ok_or("Failed to open stdout")?;
194		let stderr = child.stderr.take().ok_or("Failed to open stderr")?;
195
196		let filter_lock_warning =
197			program == "nix" && args.iter().any(|arg| arg.contains(config::SECRET_INPUT_NAME));
198
199		let stdout_handle = stream_output_lines(stdout, logger.clone(), false);
200		let stderr_handle = crate::progress::stream::stream_output_lines_filtered(
201			stderr,
202			logger.clone(),
203			true,
204			filter_lock_warning,
205		);
206
207		let stdout_output = stdout_handle.join().unwrap_or_default();
208		let stderr_output = stderr_handle.join().unwrap_or_default();
209
210		let status = child.wait()?;
211		if !status.success() {
212			let redacted_args = crate::config::redact_token(&args.join(" "));
213			let redacted_stderr = crate::config::redact_token(&stderr_output);
214			return Err(
215				format!(
216					"Command '{} {}' failed with status: {}\nError: {}",
217					program, redacted_args, status, redacted_stderr
218				)
219				.into(),
220			);
221		}
222
223		Ok(stdout_output)
224	}
225}
226
227pub fn shell_escape(arg: &str) -> String {
228	if arg.is_empty() {
229		return "''".to_string();
230	}
231
232	let is_safe = arg.chars().all(|c| {
233		c.is_ascii_alphanumeric()
234			|| c == '-'
235			|| c == '_'
236			|| c == '.'
237			|| c == '/'
238			|| c == ','
239			|| c == ':'
240			|| c == '='
241	});
242
243	if is_safe { arg.to_string() } else { format!("'{}'", arg.replace('\'', "'\\''")) }
244}
245
246#[cfg(test)]
247mod tests {
248	use super::*;
249	use std::fs::File;
250
251	#[test]
252	fn sanitized_commands_remove_ambient_github_credentials() {
253		let command = sanitized_command("true");
254		let removals = command
255			.get_envs()
256			.filter(|(_, value)| value.is_none())
257			.map(|(name, _)| name.to_string_lossy().into_owned())
258			.collect::<Vec<_>>();
259
260		assert!(removals.contains(&"GITHUB_TOKEN".to_string()));
261		assert!(removals.contains(&"GH_TOKEN".to_string()));
262	}
263
264	#[test]
265	fn test_shell_escape() {
266		assert_eq!(shell_escape(""), "''");
267		assert_eq!(shell_escape("safe-word_123.txt"), "safe-word_123.txt");
268		assert_eq!(shell_escape("unsafe word"), "'unsafe word'");
269		assert_eq!(shell_escape("don't"), "'don'\\''t'");
270	}
271
272	#[test]
273	fn test_command_recording_and_mocks() {
274		CommandExecutor::start_recording();
275		CommandExecutor::register_mock("echo hello", "mocked_hello");
276		CommandExecutor::register_mock("status", "mocked_status");
277
278		let log = Logger::silent();
279		let out = CommandExecutor::execute("echo", &["hello"], log.clone()).unwrap();
280		assert_eq!(out, "mocked_hello");
281
282		let out_substring = CommandExecutor::execute("qm", &["status", "101"], log.clone()).unwrap();
283		assert_eq!(out_substring, "mocked_status");
284
285		let out_unmocked = CommandExecutor::execute("echo", &["world"], log).unwrap();
286		assert_eq!(out_unmocked, "");
287
288		let recorded = CommandExecutor::stop_recording().unwrap();
289		assert_eq!(recorded.len(), 3);
290		assert_eq!(recorded[0], "echo hello");
291		assert_eq!(recorded[1], "qm status 101");
292		assert_eq!(recorded[2], "echo world");
293
294		CommandExecutor::clear_mocks();
295	}
296
297	#[test]
298	fn log_status_does_not_infer_level_from_message_text() {
299		let path = std::env::temp_dir().join(format!("nxd-log-status-test-{}", std::process::id()));
300		let file = File::create(&path).unwrap();
301		let mut target = LogTarget::File(file);
302
303		target.log_status("SUCCESS Target host SSH key staged successfully.");
304		drop(target);
305
306		let content = std::fs::read_to_string(&path).unwrap();
307		let _ = std::fs::remove_file(&path);
308
309		assert_eq!(content, "[INFO] SUCCESS Target host SSH key staged successfully.\n");
310	}
311}