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 mut process = Command::new("/bin/sh");
147    process
148        .arg("-lc")
149        .arg(command)
150        .current_dir(cwd)
151        .env_remove(api_key_env)
152        .stdin(Stdio::null())
153        .stdout(Stdio::piped())
154        .stderr(Stdio::piped());
155
156    #[cfg(unix)]
157    {
158        use std::os::unix::process::CommandExt;
159
160        // Each command gets its own process group so a timed-out shell and its
161        // finite descendants can be cleaned up together.
162        unsafe {
163            process.pre_exec(|| {
164                if libc::setpgid(0, 0) == -1 {
165                    return Err(io::Error::last_os_error());
166                }
167                Ok(())
168            });
169        }
170    }
171
172    let mut child = match process.spawn() {
173        Ok(child) => child,
174        Err(_) => {
175            return CmdResult {
176                command: redact_secret(command, secret),
177                exit_code: None,
178                timed_out: false,
179                stdout: String::new(),
180                stderr: String::new(),
181                stdout_truncated: false,
182                stderr_truncated: false,
183                canceled: false,
184                error: Some("unable to start command".to_owned()),
185            }
186        }
187    };
188
189    let capture_stop = Arc::new(AtomicBool::new(false));
190    let stdout_reader = child
191        .stdout
192        .take()
193        .map(|stdout| spawn_capture(stdout, output_cap, Arc::clone(&capture_stop)));
194    let stderr_reader = child
195        .stderr
196        .take()
197        .map(|stderr| spawn_capture(stderr, output_cap, Arc::clone(&capture_stop)));
198
199    let child_id = child.id();
200    let deadline = Instant::now() + timeout;
201    let mut timed_out = false;
202    let mut canceled = false;
203    let status = loop {
204        if cancellation.is_some_and(|token| token.is_cancelled()) {
205            canceled = true;
206            kill_process_group(child_id);
207            let _ = child.kill();
208            break child.wait().ok();
209        }
210        match child.try_wait() {
211            Ok(Some(status)) => {
212                if cancellation.is_some_and(|token| token.is_cancelled()) {
213                    canceled = true;
214                }
215                break Some(status);
216            }
217            Ok(None) => {
218                if Instant::now() >= deadline {
219                    timed_out = true;
220                    kill_process_group(child_id);
221                    let _ = child.kill();
222                    break child.wait().ok();
223                }
224                thread::sleep(Duration::from_millis(10));
225            }
226            Err(_) => {
227                canceled = cancellation.is_some_and(|token| token.is_cancelled());
228                kill_process_group(child_id);
229                let _ = child.kill();
230                break child.wait().ok();
231            }
232        }
233    };
234
235    if !timed_out {
236        // A shell can finish while a background child remains in its process
237        // group. Lucy has no background-process API, so clean that group too.
238        kill_process_group(child_id);
239    }
240
241    capture_stop.store(true, Ordering::Release);
242    let stdout_capture = join_capture(stdout_reader);
243    let stderr_capture = join_capture(stderr_reader);
244    let (stdout, stdout_truncated) = bounded_output(&stdout_capture, output_cap, secret);
245    let (stderr, stderr_truncated) = bounded_output(&stderr_capture, output_cap, secret);
246
247    CmdResult {
248        command: redact_secret(command, secret),
249        exit_code: (!canceled)
250            .then(|| status.and_then(|status| status.code()))
251            .flatten(),
252        timed_out,
253        stdout,
254        stderr,
255        stdout_truncated,
256        stderr_truncated,
257        canceled,
258        error: canceled.then_some("command canceled".to_owned()),
259    }
260}
261
262#[cfg(unix)]
263fn spawn_capture<R>(mut reader: R, cap: usize, stop: Arc<AtomicBool>) -> JoinHandle<CapturedOutput>
264where
265    R: Read + Send + AsRawFd + 'static,
266{
267    let _ = set_nonblocking(reader.as_raw_fd());
268    thread::spawn(move || capture_output(&mut reader, cap, &stop))
269}
270
271#[cfg(not(unix))]
272fn spawn_capture<R>(mut reader: R, cap: usize, stop: Arc<AtomicBool>) -> JoinHandle<CapturedOutput>
273where
274    R: Read + Send + 'static,
275{
276    thread::spawn(move || capture_output(&mut reader, cap, &stop))
277}
278
279fn capture_output<R>(reader: &mut R, cap: usize, stop: &AtomicBool) -> CapturedOutput
280where
281    R: Read,
282{
283    let mut bytes = Vec::with_capacity(cap.min(8192));
284    let mut buffer = [0_u8; 8192];
285    let mut truncated = false;
286    let mut shutdown_incomplete = false;
287    let mut shutdown_deadline = None;
288    loop {
289        if stop.load(Ordering::Acquire) {
290            shutdown_deadline.get_or_insert_with(|| Instant::now() + CAPTURE_SHUTDOWN_GRACE);
291            if shutdown_deadline.is_some_and(|deadline| Instant::now() >= deadline) {
292                shutdown_incomplete = true;
293                break;
294            }
295        }
296
297        match reader.read(&mut buffer) {
298            Ok(0) => break,
299            Ok(read) => {
300                let remaining = cap.saturating_sub(bytes.len());
301                if remaining > 0 {
302                    bytes.extend_from_slice(&buffer[..read.min(remaining)]);
303                }
304                if read > remaining {
305                    truncated = true;
306                }
307            }
308            Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
309                thread::sleep(Duration::from_millis(1));
310            }
311            Err(_) => break,
312        }
313    }
314    CapturedOutput {
315        bytes,
316        truncated: truncated || shutdown_incomplete,
317    }
318}
319
320#[cfg(unix)]
321fn set_nonblocking(fd: std::os::fd::RawFd) -> io::Result<()> {
322    let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
323    if flags == -1 {
324        return Err(io::Error::last_os_error());
325    }
326    let result = unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) };
327    if result == -1 {
328        Err(io::Error::last_os_error())
329    } else {
330        Ok(())
331    }
332}
333
334fn bounded_output(
335    captured: &CapturedOutput,
336    output_cap: usize,
337    secret: Option<&str>,
338) -> (String, bool) {
339    let text = redact_secret(&String::from_utf8_lossy(&captured.bytes), secret);
340    let truncated =
341        captured.truncated || captured.bytes.len() > output_cap || text.len() > output_cap;
342    let mut end = text.len().min(output_cap);
343    while end > 0 && !text.is_char_boundary(end) {
344        end -= 1;
345    }
346    (text[..end].to_owned(), truncated)
347}
348
349fn join_capture(reader: Option<JoinHandle<CapturedOutput>>) -> CapturedOutput {
350    let Some(reader) = reader else {
351        return CapturedOutput {
352            bytes: Vec::new(),
353            truncated: false,
354        };
355    };
356
357    let deadline = Instant::now() + CAPTURE_SHUTDOWN_GRACE;
358    while !reader.is_finished() && Instant::now() < deadline {
359        thread::sleep(Duration::from_millis(1));
360    }
361    if reader.is_finished() {
362        reader.join().unwrap_or(CapturedOutput {
363            bytes: Vec::new(),
364            truncated: false,
365        })
366    } else {
367        // Non-blocking capture normally exits after the shutdown grace. If a
368        // platform refuses that setup, detach rather than blocking the command
369        // harness on a descendant-owned pipe forever.
370        CapturedOutput {
371            bytes: Vec::new(),
372            truncated: true,
373        }
374    }
375}
376
377fn kill_process_group(child_id: u32) {
378    #[cfg(unix)]
379    {
380        // The child is the process-group leader created above. Ignore errors:
381        // it may already have exited, and child.wait below remains authoritative.
382        unsafe {
383            let _ = libc::kill(-(child_id as libc::pid_t), libc::SIGKILL);
384        }
385    }
386    #[cfg(not(unix))]
387    let _ = child_id;
388}
389
390pub fn redact_secret(text: &str, secret: Option<&str>) -> String {
391    crate::redaction::redact_secret(text, secret)
392}
393
394pub(crate) fn canceled_result(arguments: &str, secret: &str) -> CmdResult {
395    let command = serde_json::from_str::<Value>(arguments)
396        .ok()
397        .and_then(|value| {
398            value
399                .as_object()
400                .filter(|object| object.len() == 1)
401                .and_then(|object| object.get("command"))
402                .and_then(Value::as_str)
403                .map(str::to_owned)
404        })
405        .map(|command| redact_secret(&command, Some(secret)))
406        .unwrap_or_else(|| "{}".to_owned());
407    CmdResult::canceled(command, "command canceled before execution")
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413    use std::fs;
414    use std::sync::atomic::{AtomicU64, Ordering};
415    use std::time::{SystemTime, UNIX_EPOCH};
416
417    static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
418
419    fn temporary_directory() -> std::path::PathBuf {
420        loop {
421            let stamp = SystemTime::now()
422                .duration_since(UNIX_EPOCH)
423                .expect("clock")
424                .as_nanos();
425            let counter = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
426            let path = std::env::temp_dir().join(format!(
427                "lucy-command-{stamp}-{}-{counter}",
428                std::process::id()
429            ));
430            match fs::create_dir(&path) {
431                Ok(()) => return path,
432                Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
433                Err(error) => panic!("temp directory: {error}"),
434            }
435        }
436    }
437
438    #[test]
439    fn captures_nonzero_exit_and_both_streams() {
440        let cwd = temporary_directory();
441        let result = execute_command(
442            "printf out; printf err >&2; exit 7",
443            &cwd,
444            "LUCY_API_KEY",
445            None,
446            Duration::from_secs(2),
447            COMMAND_OUTPUT_CAP,
448        );
449        assert_eq!(result.exit_code, Some(7));
450        assert!(!result.timed_out);
451        assert_eq!(result.stdout, "out");
452        assert_eq!(result.stderr, "err");
453        fs::remove_dir_all(cwd).expect("remove temp directory");
454    }
455
456    #[test]
457    fn caps_streams_independently_and_marks_truncation() {
458        let cwd = temporary_directory();
459        let result = execute_command(
460            "printf 123456789; printf abcdefghij >&2",
461            &cwd,
462            "LUCY_API_KEY",
463            None,
464            Duration::from_secs(2),
465            4,
466        );
467        assert_eq!(result.stdout, "1234");
468        assert_eq!(result.stderr, "abcd");
469        assert!(result.stdout_truncated);
470        assert!(result.stderr_truncated);
471        fs::remove_dir_all(cwd).expect("remove temp directory");
472    }
473
474    #[cfg(unix)]
475    #[test]
476    fn bounds_lossy_invalid_utf8_output_and_marks_truncation() {
477        let cwd = temporary_directory();
478        let result = execute_command(
479            r"printf '\377\376\375\374'; printf '\377\376\375\374' >&2",
480            &cwd,
481            "LUCY_API_KEY",
482            None,
483            Duration::from_secs(2),
484            4,
485        );
486        assert!(result.stdout.len() <= 4);
487        assert!(result.stderr.len() <= 4);
488        assert!(result.stdout_truncated);
489        assert!(result.stderr_truncated);
490        fs::remove_dir_all(cwd).expect("remove temp directory");
491    }
492
493    #[test]
494    fn timeout_kills_the_command_group() {
495        let cwd = temporary_directory();
496        let result = execute_command(
497            "sleep 30",
498            &cwd,
499            "LUCY_API_KEY",
500            None,
501            Duration::from_millis(80),
502            COMMAND_OUTPUT_CAP,
503        );
504        assert!(result.timed_out);
505        assert!(result.exit_code.is_none() || result.exit_code != Some(0));
506        fs::remove_dir_all(cwd).expect("remove temp directory");
507    }
508
509    #[test]
510    fn cancellation_kills_a_running_command_group() {
511        let cwd = temporary_directory();
512        let token = CancellationToken::new();
513        let worker_token = token.clone();
514        let worker_cwd = cwd.clone();
515        let started = Instant::now();
516        let worker = thread::spawn(move || {
517            execute_command_with_cancellation(
518                "sleep 30",
519                &worker_cwd,
520                "LUCY_API_KEY",
521                None,
522                Duration::from_secs(30),
523                COMMAND_OUTPUT_CAP,
524                Some(&worker_token),
525            )
526        });
527        thread::sleep(Duration::from_millis(80));
528        token.cancel();
529        let result = worker.join().expect("command worker");
530        assert!(result.canceled);
531        assert_eq!(result.error.as_deref(), Some("command canceled"));
532        assert!(started.elapsed() < Duration::from_secs(2));
533        fs::remove_dir_all(cwd).expect("remove temp directory");
534    }
535
536    #[cfg(unix)]
537    #[test]
538    fn timeout_capture_returns_when_a_descendant_escapes_the_process_group() {
539        let python_available = Command::new("python3")
540            .arg("--version")
541            .stdout(Stdio::null())
542            .stderr(Stdio::null())
543            .status()
544            .map(|status| status.success())
545            .unwrap_or(false);
546        if !python_available {
547            return;
548        }
549
550        let cwd = temporary_directory();
551        let started = Instant::now();
552        let result = execute_command(
553            "python3 -c 'import os,time; os.setsid(); open(\"ready\",\"w\").close(); time.sleep(1)' & while [ ! -f ready ]; do sleep 0.01; done; sleep 2",
554            &cwd,
555            "LUCY_API_KEY",
556            None,
557            Duration::from_millis(300),
558            COMMAND_OUTPUT_CAP,
559        );
560        assert!(result.timed_out);
561        assert!(result.stdout_truncated);
562        assert!(
563            started.elapsed() < Duration::from_secs(1),
564            "capture cleanup exceeded the bounded grace period: {:?}",
565            started.elapsed()
566        );
567        fs::remove_dir_all(cwd).expect("remove temp directory");
568    }
569
570    #[test]
571    fn output_redacts_the_provider_key() {
572        let cwd = temporary_directory();
573        let result = execute_command(
574            "printf secret-key",
575            &cwd,
576            "LUCY_API_KEY",
577            Some("secret-key"),
578            Duration::from_secs(2),
579            COMMAND_OUTPUT_CAP,
580        );
581        assert!(!result.stdout.contains("secret-key"));
582        assert_eq!(result.stdout, "[REDACTED]");
583        fs::remove_dir_all(cwd).expect("remove temp directory");
584    }
585
586    #[test]
587    fn redaction_stays_within_the_capture_byte_bound() {
588        let cwd = temporary_directory();
589        let result = execute_command(
590            "printf x",
591            &cwd,
592            "LUCY_API_KEY",
593            Some("x"),
594            Duration::from_secs(2),
595            1,
596        );
597        assert_eq!(result.stdout.len(), 1);
598        assert!(!result.stdout.contains('x'));
599        fs::remove_dir_all(cwd).expect("remove temp directory");
600    }
601
602    #[test]
603    fn collision_markers_do_not_reintroduce_the_provider_key() {
604        let cwd = temporary_directory();
605        for secret in ["REDACTED", "[REDACTED]"] {
606            let command = format!("printf '{secret}'");
607            let result = execute_command(
608                &command,
609                &cwd,
610                "LUCY_API_KEY",
611                Some(secret),
612                Duration::from_secs(2),
613                COMMAND_OUTPUT_CAP,
614            );
615            assert!(!result.stdout.contains(secret));
616            assert!(!result.command.contains(secret));
617        }
618        fs::remove_dir_all(cwd).expect("remove temp directory");
619    }
620
621    #[test]
622    fn rejects_extra_command_arguments() {
623        let cwd = temporary_directory();
624        let result = execute(
625            r#"{"command":"pwd","extra":true}"#,
626            &cwd,
627            "LUCY_API_KEY",
628            None,
629        );
630        assert_eq!(
631            result.error.as_deref(),
632            Some("cmd arguments must contain only command")
633        );
634        fs::remove_dir_all(cwd).expect("remove temp directory");
635    }
636}
lucy/command.rs

lucy/
command.rs
