1use std::io::{self, Read};
2use std::path::Path;
3use std::process::{Command, Stdio};
4use std::sync::{
5 atomic::{AtomicBool, Ordering},
6 Arc,
7};
8use std::thread::{self, JoinHandle};
9use std::time::{Duration, Instant};
10
11#[cfg(unix)]
12use std::os::fd::AsRawFd;
13
14use serde::{Deserialize, Serialize};
15use serde_json::Value;
16
17use crate::cancellation::CancellationToken;
18
19pub const COMMAND_TIMEOUT: Duration = Duration::from_secs(10 * 60);
20pub const COMMAND_OUTPUT_CAP: usize = 64 * 1024;
21const CAPTURE_SHUTDOWN_GRACE: Duration = Duration::from_millis(100);
22
23#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
24pub struct CmdResult {
25 pub command: String,
26 pub exit_code: Option<i32>,
27 pub timed_out: bool,
28 pub stdout: String,
29 pub stderr: String,
30 pub stdout_truncated: bool,
31 pub stderr_truncated: bool,
32 #[serde(default, skip_serializing_if = "is_false")]
33 pub canceled: bool,
34 #[serde(skip_serializing_if = "Option::is_none")]
35 pub error: Option<String>,
36}
37
38impl CmdResult {
39 fn error(arguments: &str, message: impl Into<String>) -> Self {
40 Self {
41 command: arguments.to_owned(),
42 exit_code: None,
43 timed_out: false,
44 stdout: String::new(),
45 stderr: String::new(),
46 stdout_truncated: false,
47 stderr_truncated: false,
48 canceled: false,
49 error: Some(message.into()),
50 }
51 }
52
53 pub(crate) fn canceled(command: impl Into<String>, message: impl Into<String>) -> Self {
54 Self {
55 command: command.into(),
56 exit_code: None,
57 timed_out: false,
58 stdout: String::new(),
59 stderr: String::new(),
60 stdout_truncated: false,
61 stderr_truncated: false,
62 canceled: true,
63 error: Some(message.into()),
64 }
65 }
66}
67
68fn is_false(value: &bool) -> bool {
69 !*value
70}
71
72#[derive(Debug)]
73struct CapturedOutput {
74 bytes: Vec<u8>,
75 truncated: bool,
76}
77
78pub fn execute(arguments: &str, cwd: &Path, api_key_env: &str, secret: Option<&str>) -> CmdResult {
79 execute_with_cancellation(arguments, cwd, api_key_env, secret, None)
80}
81
82pub(crate) fn execute_with_cancellation(
83 arguments: &str,
84 cwd: &Path,
85 api_key_env: &str,
86 secret: Option<&str>,
87 cancellation: Option<&CancellationToken>,
88) -> CmdResult {
89 let value: Value = match serde_json::from_str(arguments) {
90 Ok(value) => value,
91 Err(_) => return CmdResult::error("{}", "cmd arguments must be a JSON object"),
92 };
93 let Some(object) = value.as_object() else {
94 return CmdResult::error("{}", "cmd arguments must be a JSON object");
95 };
96 if object.len() != 1 || !object.contains_key("command") {
97 return CmdResult::error("{}", "cmd arguments must contain only command");
98 }
99 let Some(command) = object.get("command").and_then(Value::as_str) else {
100 return CmdResult::error("{}", "cmd command must be a string");
101 };
102 if cancellation.is_some_and(|token| token.is_cancelled()) {
103 return CmdResult::canceled(
104 redact_secret(command, secret),
105 "command canceled before execution",
106 );
107 }
108 execute_command_with_cancellation(
109 command,
110 cwd,
111 api_key_env,
112 secret,
113 COMMAND_TIMEOUT,
114 COMMAND_OUTPUT_CAP,
115 cancellation,
116 )
117}
118
119pub fn execute_command(
120 command: &str,
121 cwd: &Path,
122 api_key_env: &str,
123 secret: Option<&str>,
124 timeout: Duration,
125 output_cap: usize,
126) -> CmdResult {
127 execute_command_with_cancellation(command, cwd, api_key_env, secret, timeout, output_cap, None)
128}
129
130pub(crate) fn execute_command_with_cancellation(
131 command: &str,
132 cwd: &Path,
133 _api_key_env: &str,
134 secret: Option<&str>,
135 timeout: Duration,
136 output_cap: usize,
137 cancellation: Option<&CancellationToken>,
138) -> CmdResult {
139 if cancellation.is_some_and(|token| token.is_cancelled()) {
140 return CmdResult::canceled(
141 redact_secret(command, secret),
142 "command canceled before execution",
143 );
144 }
145
146 let shell = std::env::var_os("SHELL")
147 .filter(|shell| !shell.is_empty())
148 .unwrap_or_else(|| "/bin/sh".into());
149 let mut process = Command::new(shell);
150 process
151 .arg("-lc")
152 .arg(command)
153 .current_dir(cwd)
154 .stdin(Stdio::null())
155 .stdout(Stdio::piped())
156 .stderr(Stdio::piped());
157
158 #[cfg(unix)]
159 {
160 use std::os::unix::process::CommandExt;
161
162 unsafe {
165 process.pre_exec(|| {
166 if libc::setpgid(0, 0) == -1 {
167 return Err(io::Error::last_os_error());
168 }
169 Ok(())
170 });
171 }
172 }
173
174 let mut child = match process.spawn() {
175 Ok(child) => child,
176 Err(_) => {
177 return CmdResult {
178 command: redact_secret(command, secret),
179 exit_code: None,
180 timed_out: false,
181 stdout: String::new(),
182 stderr: String::new(),
183 stdout_truncated: false,
184 stderr_truncated: false,
185 canceled: false,
186 error: Some("unable to start command".to_owned()),
187 }
188 }
189 };
190
191 let capture_stop = Arc::new(AtomicBool::new(false));
192 let stdout_reader = child
193 .stdout
194 .take()
195 .map(|stdout| spawn_capture(stdout, output_cap, Arc::clone(&capture_stop)));
196 let stderr_reader = child
197 .stderr
198 .take()
199 .map(|stderr| spawn_capture(stderr, output_cap, Arc::clone(&capture_stop)));
200
201 let child_id = child.id();
202 let deadline = Instant::now() + timeout;
203 let mut timed_out = false;
204 let mut canceled = false;
205 let status = loop {
206 if cancellation.is_some_and(|token| token.is_cancelled()) {
207 canceled = true;
208 kill_process_group(child_id);
209 let _ = child.kill();
210 break child.wait().ok();
211 }
212 match child.try_wait() {
213 Ok(Some(status)) => {
214 if cancellation.is_some_and(|token| token.is_cancelled()) {
215 canceled = true;
216 }
217 break Some(status);
218 }
219 Ok(None) => {
220 if Instant::now() >= deadline {
221 timed_out = true;
222 kill_process_group(child_id);
223 let _ = child.kill();
224 break child.wait().ok();
225 }
226 thread::sleep(Duration::from_millis(10));
227 }
228 Err(_) => {
229 canceled = cancellation.is_some_and(|token| token.is_cancelled());
230 kill_process_group(child_id);
231 let _ = child.kill();
232 break child.wait().ok();
233 }
234 }
235 };
236
237 if !timed_out {
238 kill_process_group(child_id);
241 }
242
243 capture_stop.store(true, Ordering::Release);
244 let stdout_capture = join_capture(stdout_reader);
245 let stderr_capture = join_capture(stderr_reader);
246 let (stdout, stdout_truncated) = bounded_output(&stdout_capture, output_cap, secret);
247 let (stderr, stderr_truncated) = bounded_output(&stderr_capture, output_cap, secret);
248
249 CmdResult {
250 command: redact_secret(command, secret),
251 exit_code: (!canceled)
252 .then(|| status.and_then(|status| status.code()))
253 .flatten(),
254 timed_out,
255 stdout,
256 stderr,
257 stdout_truncated,
258 stderr_truncated,
259 canceled,
260 error: canceled.then_some("command canceled".to_owned()),
261 }
262}
263
264#[cfg(unix)]
265fn spawn_capture<R>(mut reader: R, cap: usize, stop: Arc<AtomicBool>) -> JoinHandle<CapturedOutput>
266where
267 R: Read + Send + AsRawFd + 'static,
268{
269 let _ = set_nonblocking(reader.as_raw_fd());
270 thread::spawn(move || capture_output(&mut reader, cap, &stop))
271}
272
273#[cfg(not(unix))]
274fn spawn_capture<R>(mut reader: R, cap: usize, stop: Arc<AtomicBool>) -> JoinHandle<CapturedOutput>
275where
276 R: Read + Send + 'static,
277{
278 thread::spawn(move || capture_output(&mut reader, cap, &stop))
279}
280
281fn capture_output<R>(reader: &mut R, cap: usize, stop: &AtomicBool) -> CapturedOutput
282where
283 R: Read,
284{
285 let mut bytes = Vec::with_capacity(cap.min(8192));
286 let mut buffer = [0_u8; 8192];
287 let mut truncated = false;
288 let mut shutdown_incomplete = false;
289 let mut shutdown_deadline = None;
290 loop {
291 if stop.load(Ordering::Acquire) {
292 shutdown_deadline.get_or_insert_with(|| Instant::now() + CAPTURE_SHUTDOWN_GRACE);
293 if shutdown_deadline.is_some_and(|deadline| Instant::now() >= deadline) {
294 shutdown_incomplete = true;
295 break;
296 }
297 }
298
299 match reader.read(&mut buffer) {
300 Ok(0) => break,
301 Ok(read) => {
302 let remaining = cap.saturating_sub(bytes.len());
303 if remaining > 0 {
304 bytes.extend_from_slice(&buffer[..read.min(remaining)]);
305 }
306 if read > remaining {
307 truncated = true;
308 }
309 }
310 Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
311 thread::sleep(Duration::from_millis(1));
312 }
313 Err(_) => break,
314 }
315 }
316 CapturedOutput {
317 bytes,
318 truncated: truncated || shutdown_incomplete,
319 }
320}
321
322#[cfg(unix)]
323fn set_nonblocking(fd: std::os::fd::RawFd) -> io::Result<()> {
324 let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
325 if flags == -1 {
326 return Err(io::Error::last_os_error());
327 }
328 let result = unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) };
329 if result == -1 {
330 Err(io::Error::last_os_error())
331 } else {
332 Ok(())
333 }
334}
335
336fn bounded_output(
337 captured: &CapturedOutput,
338 output_cap: usize,
339 secret: Option<&str>,
340) -> (String, bool) {
341 let text = redact_secret(&String::from_utf8_lossy(&captured.bytes), secret);
342 let truncated =
343 captured.truncated || captured.bytes.len() > output_cap || text.len() > output_cap;
344 let mut end = text.len().min(output_cap);
345 while end > 0 && !text.is_char_boundary(end) {
346 end -= 1;
347 }
348 (text[..end].to_owned(), truncated)
349}
350
351fn join_capture(reader: Option<JoinHandle<CapturedOutput>>) -> CapturedOutput {
352 let Some(reader) = reader else {
353 return CapturedOutput {
354 bytes: Vec::new(),
355 truncated: false,
356 };
357 };
358
359 let deadline = Instant::now() + CAPTURE_SHUTDOWN_GRACE;
360 while !reader.is_finished() && Instant::now() < deadline {
361 thread::sleep(Duration::from_millis(1));
362 }
363 if reader.is_finished() {
364 reader.join().unwrap_or(CapturedOutput {
365 bytes: Vec::new(),
366 truncated: false,
367 })
368 } else {
369 CapturedOutput {
373 bytes: Vec::new(),
374 truncated: true,
375 }
376 }
377}
378
379fn kill_process_group(child_id: u32) {
380 #[cfg(unix)]
381 {
382 unsafe {
385 let _ = libc::kill(-(child_id as libc::pid_t), libc::SIGKILL);
386 }
387 }
388 #[cfg(not(unix))]
389 let _ = child_id;
390}
391
392pub fn redact_secret(text: &str, secret: Option<&str>) -> String {
393 crate::redaction::redact_secret(text, secret)
394}
395
396pub(crate) fn canceled_result(arguments: &str, secret: &str) -> CmdResult {
397 let command = serde_json::from_str::<Value>(arguments)
398 .ok()
399 .and_then(|value| {
400 value
401 .as_object()
402 .filter(|object| object.len() == 1)
403 .and_then(|object| object.get("command"))
404 .and_then(Value::as_str)
405 .map(str::to_owned)
406 })
407 .map(|command| redact_secret(&command, Some(secret)))
408 .unwrap_or_else(|| "{}".to_owned());
409 CmdResult::canceled(command, "command canceled before execution")
410}
411
412#[cfg(test)]
413mod tests {
414 use super::*;
415 use std::fs;
416 use std::sync::atomic::{AtomicU64, Ordering};
417 use std::sync::Mutex;
418 use std::time::{SystemTime, UNIX_EPOCH};
419
420 static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
421 static COMMAND_TEST_LOCK: Mutex<()> = Mutex::new(());
422
423 fn temporary_directory() -> std::path::PathBuf {
424 loop {
425 let stamp = SystemTime::now()
426 .duration_since(UNIX_EPOCH)
427 .expect("clock")
428 .as_nanos();
429 let counter = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
430 let path = std::env::temp_dir().join(format!(
431 "lucy-command-{stamp}-{}-{counter}",
432 std::process::id()
433 ));
434 match fs::create_dir(&path) {
435 Ok(()) => return path,
436 Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
437 Err(error) => panic!("temp directory: {error}"),
438 }
439 }
440 }
441
442 #[test]
443 fn captures_nonzero_exit_and_both_streams() {
444 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
445 let cwd = temporary_directory();
446 let result = execute_command(
447 "printf out; printf err >&2; exit 7",
448 &cwd,
449 "LUCY_API_KEY",
450 None,
451 Duration::from_secs(2),
452 COMMAND_OUTPUT_CAP,
453 );
454 assert_eq!(result.exit_code, Some(7));
455 assert!(!result.timed_out);
456 assert_eq!(result.stdout, "out");
457 assert_eq!(result.stderr, "err");
458 fs::remove_dir_all(cwd).expect("remove temp directory");
459 }
460
461 #[test]
462 fn caps_streams_independently_and_marks_truncation() {
463 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
464 let cwd = temporary_directory();
465 let result = execute_command(
466 "printf 123456789; printf abcdefghij >&2",
467 &cwd,
468 "LUCY_API_KEY",
469 None,
470 Duration::from_secs(2),
471 4,
472 );
473 assert_eq!(result.stdout, "1234");
474 assert_eq!(result.stderr, "abcd");
475 assert!(result.stdout_truncated);
476 assert!(result.stderr_truncated);
477 fs::remove_dir_all(cwd).expect("remove temp directory");
478 }
479
480 #[cfg(unix)]
481 #[test]
482 fn bounds_lossy_invalid_utf8_output_and_marks_truncation() {
483 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
484 let cwd = temporary_directory();
485 let result = execute_command(
486 r"printf '\377\376\375\374'; printf '\377\376\375\374' >&2",
487 &cwd,
488 "LUCY_API_KEY",
489 None,
490 Duration::from_secs(2),
491 4,
492 );
493 assert!(result.stdout.len() <= 4);
494 assert!(result.stderr.len() <= 4);
495 assert!(result.stdout_truncated);
496 assert!(result.stderr_truncated);
497 fs::remove_dir_all(cwd).expect("remove temp directory");
498 }
499
500 #[test]
501 fn timeout_kills_the_command_group() {
502 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
503 let cwd = temporary_directory();
504 let result = execute_command(
505 "sleep 30",
506 &cwd,
507 "LUCY_API_KEY",
508 None,
509 Duration::from_millis(80),
510 COMMAND_OUTPUT_CAP,
511 );
512 assert!(result.timed_out);
513 assert!(result.exit_code.is_none() || result.exit_code != Some(0));
514 fs::remove_dir_all(cwd).expect("remove temp directory");
515 }
516
517 #[test]
518 fn cancellation_kills_a_running_command_group() {
519 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
520 let cwd = temporary_directory();
521 let token = CancellationToken::new();
522 let worker_token = token.clone();
523 let worker_cwd = cwd.clone();
524 let started = Instant::now();
525 let worker = thread::spawn(move || {
526 execute_command_with_cancellation(
527 "sleep 30",
528 &worker_cwd,
529 "LUCY_API_KEY",
530 None,
531 Duration::from_secs(30),
532 COMMAND_OUTPUT_CAP,
533 Some(&worker_token),
534 )
535 });
536 thread::sleep(Duration::from_millis(80));
537 token.cancel();
538 let result = worker.join().expect("command worker");
539 assert!(result.canceled);
540 assert_eq!(result.error.as_deref(), Some("command canceled"));
541 assert!(started.elapsed() < Duration::from_secs(2));
542 fs::remove_dir_all(cwd).expect("remove temp directory");
543 }
544
545 #[cfg(unix)]
546 #[test]
547 fn timeout_capture_returns_when_a_descendant_escapes_the_process_group() {
548 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
549 let python_available = Command::new("python3")
550 .arg("--version")
551 .stdout(Stdio::null())
552 .stderr(Stdio::null())
553 .status()
554 .map(|status| status.success())
555 .unwrap_or(false);
556 if !python_available {
557 return;
558 }
559
560 let cwd = temporary_directory();
561 let started = Instant::now();
562 let result = execute_command(
563 "python3 -c 'import os,time; os.setsid(); open(\"ready\",\"w\").close(); time.sleep(1)' & while [ ! -f ready ]; do sleep 0.01; done; sleep 2",
564 &cwd,
565 "LUCY_API_KEY",
566 None,
567 Duration::from_millis(300),
568 COMMAND_OUTPUT_CAP,
569 );
570 assert!(result.timed_out);
571 assert!(result.stdout_truncated);
572 assert!(
573 started.elapsed() < Duration::from_secs(1),
574 "capture cleanup exceeded the bounded grace period: {:?}",
575 started.elapsed()
576 );
577 fs::remove_dir_all(cwd).expect("remove temp directory");
578 }
579
580 #[test]
581 fn output_redacts_the_provider_key() {
582 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
583 let cwd = temporary_directory();
584 let result = execute_command(
585 "printf secret-key",
586 &cwd,
587 "LUCY_API_KEY",
588 Some("secret-key"),
589 Duration::from_secs(2),
590 COMMAND_OUTPUT_CAP,
591 );
592 assert!(!result.stdout.contains("secret-key"));
593 assert_eq!(result.stdout, "[REDACTED]");
594 fs::remove_dir_all(cwd).expect("remove temp directory");
595 }
596
597 #[test]
598 fn redaction_stays_within_the_capture_byte_bound() {
599 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
600 let cwd = temporary_directory();
601 let result = execute_command(
602 "printf x",
603 &cwd,
604 "LUCY_API_KEY",
605 Some("x"),
606 Duration::from_secs(2),
607 1,
608 );
609 assert_eq!(result.stdout.len(), 1);
610 assert!(!result.stdout.contains('x'));
611 fs::remove_dir_all(cwd).expect("remove temp directory");
612 }
613
614 #[test]
615 fn collision_markers_do_not_reintroduce_the_provider_key() {
616 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
617 let cwd = temporary_directory();
618 for secret in ["REDACTED", "[REDACTED]"] {
619 let command = format!("printf '{secret}'");
620 let result = execute_command(
621 &command,
622 &cwd,
623 "LUCY_API_KEY",
624 Some(secret),
625 Duration::from_secs(2),
626 COMMAND_OUTPUT_CAP,
627 );
628 assert!(!result.stdout.contains(secret));
629 assert!(!result.command.contains(secret));
630 }
631 fs::remove_dir_all(cwd).expect("remove temp directory");
632 }
633
634 #[test]
635 fn rejects_extra_command_arguments() {
636 let _test_lock = COMMAND_TEST_LOCK.lock().expect("command test lock");
637 let cwd = temporary_directory();
638 let result = execute(
639 r#"{"command":"pwd","extra":true}"#,
640 &cwd,
641 "LUCY_API_KEY",
642 None,
643 );
644 assert_eq!(
645 result.error.as_deref(),
646 Some("cmd arguments must contain only command")
647 );
648 fs::remove_dir_all(cwd).expect("remove temp directory");
649 }
650}