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			logger.dump_suppressed_tail();
129			let redacted_args = crate::config::redact_token(&args.join(" "));
130			let redacted_stderr = crate::config::redact_token(&stderr_output);
131			return Err(
132				format!(
133					"Command '{} {}' failed with status: {}\nError: {}",
134					program, redacted_args, status, redacted_stderr
135				)
136				.into(),
137			);
138		}
139
140		Ok(stdout_output)
141	}
142
143	/// Executes a command locally, feeding content to its stdin.
144	pub fn execute_with_stdin(
145		program: &str,
146		args: &[&str],
147		stdin_content: &str,
148		logger: Logger,
149	) -> Result<String, Box<dyn std::error::Error>> {
150		let safe_stdin =
151			if stdin_content.contains("PRIVATE KEY") { "<private key hidden>" } else { stdin_content };
152		let cmd_str = format!("{} {} [stdin: {}]", program, args.join(" "), safe_stdin);
153		let cmd_str = crate::config::redact_token(&cmd_str);
154		let is_rec = COMMAND_RECORDER.with(|r| {
155			if let Some(ref mut v) = *r.borrow_mut() {
156				v.push(cmd_str.clone());
157				true
158			} else {
159				false
160			}
161		});
162
163		if is_rec {
164			let mock_out = COMMAND_MOCKS
165				.with(|m| {
166					let map = m.borrow();
167					if let Some(out) = map.get(&cmd_str) {
168						return Some(out.clone());
169					}
170					for (key, val) in map.iter() {
171						if cmd_str.contains(key) {
172							return Some(val.clone());
173						}
174					}
175					None
176				})
177				.unwrap_or_default();
178			return Ok(mock_out);
179		}
180
181		let mut child = sanitized_command(program)
182			.args(args)
183			.stdin(Stdio::piped())
184			.stdout(Stdio::piped())
185			.stderr(Stdio::piped())
186			.spawn()?;
187
188		if let Some(mut stdin) = child.stdin.take() {
189			use std::io::Write as _;
190			stdin.write_all(stdin_content.as_bytes())?;
191			stdin.flush()?;
192		}
193
194		let stdout = child.stdout.take().ok_or("Failed to open stdout")?;
195		let stderr = child.stderr.take().ok_or("Failed to open stderr")?;
196
197		let filter_lock_warning =
198			program == "nix" && args.iter().any(|arg| arg.contains(config::SECRET_INPUT_NAME));
199
200		let stdout_handle = stream_output_lines(stdout, logger.clone(), false);
201		let stderr_handle = crate::progress::stream::stream_output_lines_filtered(
202			stderr,
203			logger.clone(),
204			true,
205			filter_lock_warning,
206		);
207
208		let stdout_output = stdout_handle.join().unwrap_or_default();
209		let stderr_output = stderr_handle.join().unwrap_or_default();
210
211		let status = child.wait()?;
212		if !status.success() {
213			logger.dump_suppressed_tail();
214			let redacted_args = crate::config::redact_token(&args.join(" "));
215			let redacted_stderr = crate::config::redact_token(&stderr_output);
216			return Err(
217				format!(
218					"Command '{} {}' failed with status: {}\nError: {}",
219					program, redacted_args, status, redacted_stderr
220				)
221				.into(),
222			);
223		}
224
225		Ok(stdout_output)
226	}
227}
228
229pub fn shell_escape(arg: &str) -> String {
230	if arg.is_empty() {
231		return "''".to_string();
232	}
233
234	let is_safe = arg.chars().all(|c| {
235		c.is_ascii_alphanumeric()
236			|| c == '-'
237			|| c == '_'
238			|| c == '.'
239			|| c == '/'
240			|| c == ','
241			|| c == ':'
242			|| c == '='
243	});
244
245	if is_safe { arg.to_string() } else { format!("'{}'", arg.replace('\'', "'\\''")) }
246}
247
248#[cfg(test)]
249mod tests {
250	use super::*;
251	use std::fs::File;
252
253	#[test]
254	fn sanitized_commands_remove_ambient_github_credentials() {
255		let command = sanitized_command("true");
256		let removals = command
257			.get_envs()
258			.filter(|(_, value)| value.is_none())
259			.map(|(name, _)| name.to_string_lossy().into_owned())
260			.collect::<Vec<_>>();
261
262		assert!(removals.contains(&"GITHUB_TOKEN".to_string()));
263		assert!(removals.contains(&"GH_TOKEN".to_string()));
264	}
265
266	#[test]
267	fn test_shell_escape() {
268		assert_eq!(shell_escape(""), "''");
269		assert_eq!(shell_escape("safe-word_123.txt"), "safe-word_123.txt");
270		assert_eq!(shell_escape("unsafe word"), "'unsafe word'");
271		assert_eq!(shell_escape("don't"), "'don'\\''t'");
272	}
273
274	#[test]
275	fn test_command_recording_and_mocks() {
276		CommandExecutor::start_recording();
277		CommandExecutor::register_mock("echo hello", "mocked_hello");
278		CommandExecutor::register_mock("status", "mocked_status");
279
280		let log = Logger::silent();
281		let out = CommandExecutor::execute("echo", &["hello"], log.clone()).unwrap();
282		assert_eq!(out, "mocked_hello");
283
284		let out_substring = CommandExecutor::execute("qm", &["status", "101"], log.clone()).unwrap();
285		assert_eq!(out_substring, "mocked_status");
286
287		let out_unmocked = CommandExecutor::execute("echo", &["world"], log).unwrap();
288		assert_eq!(out_unmocked, "");
289
290		let recorded = CommandExecutor::stop_recording().unwrap();
291		assert_eq!(recorded.len(), 3);
292		assert_eq!(recorded[0], "echo hello");
293		assert_eq!(recorded[1], "qm status 101");
294		assert_eq!(recorded[2], "echo world");
295
296		CommandExecutor::clear_mocks();
297	}
298
299	#[test]
300	fn log_status_does_not_infer_level_from_message_text() {
301		let path = std::env::temp_dir().join(format!("nxd-log-status-test-{}", std::process::id()));
302		let file = File::create(&path).unwrap();
303		let mut target = LogTarget::File(file);
304
305		target.log_status("SUCCESS Target host SSH key staged successfully.");
306		drop(target);
307
308		let content = std::fs::read_to_string(&path).unwrap();
309		let _ = std::fs::remove_file(&path);
310
311		assert_eq!(content, "[INFO] SUCCESS Target host SSH key staged successfully.\n");
312	}
313}