Skip to main content

lucy/
command.rs

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        // Each command gets its own process group so a timed-out shell and its
163        // finite descendants can be cleaned up together.
164        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        // A shell can finish while a background child remains in its process
239        // group. Lucy has no background-process API, so clean that group too.
240        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        // Non-blocking capture normally exits after the shutdown grace. If a
370        // platform refuses that setup, detach rather than blocking the command
371        // harness on a descendant-owned pipe forever.
372        CapturedOutput {
373            bytes: Vec::new(),
374            truncated: true,
375        }
376    }
377}
378
379fn kill_process_group(child_id: u32) {
380    #[cfg(unix)]
381    {
382        // The child is the process-group leader created above. Ignore errors:
383        // it may already have exited, and child.wait below remains authoritative.
384        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}