Skip to main content

lean_ctx/core/
sandbox.rs

1use std::collections::HashMap;
2use std::process::Command;
3
4#[derive(Debug, Clone)]
5pub struct SandboxResult {
6    pub stdout: String,
7    pub stderr: String,
8    pub exit_code: i32,
9    pub language: String,
10    pub duration_ms: u64,
11}
12
13const TIMEOUT_SECS: u64 = 30;
14const MAX_OUTPUT_BYTES: usize = 32_768;
15/// Upper bound on the `code` payload. Generous for real scripts while preventing an
16/// agent from forcing a multi-megabyte temp-file write / interpreter argv → memory abuse.
17const MAX_CODE_BYTES: usize = 256 * 1024;
18
19pub fn execute(language: &str, code: &str, timeout_secs: Option<u64>) -> SandboxResult {
20    if code.len() > MAX_CODE_BYTES {
21        return SandboxResult {
22            stdout: String::new(),
23            stderr: format!(
24                "Code exceeds the {MAX_CODE_BYTES}-byte limit ({} bytes). Split it into smaller scripts.",
25                code.len()
26            ),
27            exit_code: 1,
28            language: language.to_string(),
29            duration_ms: 0,
30        };
31    }
32
33    let timeout = timeout_secs.unwrap_or(TIMEOUT_SECS);
34    let start = std::time::Instant::now();
35
36    let Some(runtime) = resolve_runtime(language) else {
37        return SandboxResult {
38            stdout: String::new(),
39            stderr: format!(
40                "Unsupported language: {language}. Supported: javascript, typescript, python, shell, ruby, go, rust, php, perl, r, elixir"
41            ),
42            exit_code: 1,
43            language: language.to_string(),
44            duration_ms: 0,
45        };
46    };
47
48    let sandbox_level = std::env::var("LEAN_CTX_SANDBOX_LEVEL")
49        .ok()
50        .and_then(|v| v.parse::<u8>().ok())
51        .unwrap_or_else(|| crate::core::config::Config::load().sandbox_level);
52
53    if sandbox_level >= 1 && cfg!(target_os = "macos") {
54        let result = seatbelt_execute(&runtime, code, timeout);
55        let duration_ms = start.elapsed().as_millis() as u64;
56        return match result {
57            Ok((stdout, stderr, exit_code)) => SandboxResult {
58                stdout: truncate_output(&stdout),
59                stderr: truncate_smart(&stderr, 2048),
60                exit_code,
61                language: language.to_string(),
62                duration_ms,
63            },
64            Err(e) => SandboxResult {
65                stdout: String::new(),
66                stderr: format!("Seatbelt execution error: {e}"),
67                exit_code: 1,
68                language: language.to_string(),
69                duration_ms,
70            },
71        };
72    } else if sandbox_level >= 1 {
73        #[cfg(target_os = "linux")]
74        {
75            let result = landlock_execute(&runtime, code, timeout);
76            let duration_ms = start.elapsed().as_millis() as u64;
77            return match result {
78                Ok((stdout, stderr, exit_code)) => SandboxResult {
79                    stdout: truncate_output(&stdout),
80                    stderr: truncate_smart(&stderr, 2048),
81                    exit_code,
82                    language: language.to_string(),
83                    duration_ms,
84                },
85                Err(e) => SandboxResult {
86                    stdout: String::new(),
87                    stderr: format!("Landlock execution error: {e}"),
88                    exit_code: 1,
89                    language: language.to_string(),
90                    duration_ms,
91                },
92            };
93        }
94
95        #[cfg(not(any(target_os = "macos", target_os = "linux")))]
96        eprintln!(
97            "[lean-ctx] sandbox_level=1 requested but sandboxing not available on this platform; falling back to Level 0"
98        );
99    }
100
101    let result = if runtime.needs_temp_file {
102        execute_with_file(&runtime, code, timeout)
103    } else {
104        execute_with_stdin(&runtime, code, timeout)
105    };
106
107    let duration_ms = start.elapsed().as_millis() as u64;
108
109    match result {
110        Ok((stdout, stderr, code)) => SandboxResult {
111            stdout: truncate_output(&stdout),
112            stderr: truncate_smart(&stderr, 2048),
113            exit_code: code,
114            language: language.to_string(),
115            duration_ms,
116        },
117        Err(e) => SandboxResult {
118            stdout: String::new(),
119            stderr: format!("Execution error: {e}"),
120            exit_code: 1,
121            language: language.to_string(),
122            duration_ms,
123        },
124    }
125}
126
127pub fn batch_execute(items: &[(String, String)]) -> Vec<SandboxResult> {
128    items
129        .iter()
130        .map(|(lang, code)| execute(lang, code, None))
131        .collect()
132}
133
134struct RuntimeConfig {
135    command: String,
136    args: Vec<String>,
137    needs_temp_file: bool,
138    file_extension: String,
139    env: HashMap<String, String>,
140}
141
142fn resolve_runtime(language: &str) -> Option<RuntimeConfig> {
143    let lang = language.to_lowercase();
144    let lang = lang.as_str();
145
146    match lang {
147        "javascript" | "js" | "node" => Some(RuntimeConfig {
148            command: find_binary(&["bun", "node"])?,
149            args: vec!["-e".to_string()],
150            needs_temp_file: false,
151            file_extension: "js".to_string(),
152            env: HashMap::new(),
153        }),
154        "typescript" | "ts" => Some(RuntimeConfig {
155            command: find_binary(&["bun", "npx"])?,
156            args: if which_exists("bun") {
157                vec!["-e".to_string()]
158            } else {
159                vec!["tsx".to_string(), "-e".to_string()]
160            },
161            needs_temp_file: false,
162            file_extension: "ts".to_string(),
163            env: HashMap::new(),
164        }),
165        "python" | "py" => Some(RuntimeConfig {
166            command: find_binary(&["python3", "python"])?,
167            args: vec!["-c".to_string()],
168            needs_temp_file: false,
169            file_extension: "py".to_string(),
170            env: HashMap::from([("PYTHONDONTWRITEBYTECODE".into(), "1".into())]),
171        }),
172        "shell" | "bash" | "sh" => {
173            #[cfg(target_os = "windows")]
174            {
175                Some(RuntimeConfig {
176                    command: "cmd".to_string(),
177                    args: vec!["/C".to_string()],
178                    needs_temp_file: false,
179                    file_extension: "bat".to_string(),
180                    env: HashMap::new(),
181                })
182            }
183            #[cfg(not(target_os = "windows"))]
184            {
185                Some(RuntimeConfig {
186                    command: find_binary(&["bash", "sh"])?,
187                    args: vec!["-c".to_string()],
188                    needs_temp_file: false,
189                    file_extension: "sh".to_string(),
190                    env: HashMap::new(),
191                })
192            }
193        }
194        "ruby" | "rb" => Some(RuntimeConfig {
195            command: find_binary(&["ruby"])?,
196            args: vec!["-e".to_string()],
197            needs_temp_file: false,
198            file_extension: "rb".to_string(),
199            env: HashMap::new(),
200        }),
201        "go" | "golang" => Some(RuntimeConfig {
202            command: find_binary(&["go"])?,
203            args: vec!["run".to_string()],
204            needs_temp_file: true,
205            file_extension: "go".to_string(),
206            env: HashMap::new(),
207        }),
208        "rust" | "rs" => Some(RuntimeConfig {
209            command: "rustc_script".to_string(),
210            args: vec![],
211            needs_temp_file: true,
212            file_extension: "rs".to_string(),
213            env: HashMap::new(),
214        }),
215        "php" => Some(RuntimeConfig {
216            command: find_binary(&["php"])?,
217            args: vec!["-r".to_string()],
218            needs_temp_file: false,
219            file_extension: "php".to_string(),
220            env: HashMap::new(),
221        }),
222        "perl" | "pl" => Some(RuntimeConfig {
223            command: find_binary(&["perl"])?,
224            args: vec!["-e".to_string()],
225            needs_temp_file: false,
226            file_extension: "pl".to_string(),
227            env: HashMap::new(),
228        }),
229        "r" => Some(RuntimeConfig {
230            command: find_binary(&["Rscript"])?,
231            args: vec!["-e".to_string()],
232            needs_temp_file: false,
233            file_extension: "R".to_string(),
234            env: HashMap::new(),
235        }),
236        "elixir" | "ex" => Some(RuntimeConfig {
237            command: find_binary(&["elixir"])?,
238            args: vec!["-e".to_string()],
239            needs_temp_file: false,
240            file_extension: "exs".to_string(),
241            env: HashMap::new(),
242        }),
243        _ => None,
244    }
245}
246
247fn seatbelt_execute(
248    runtime: &RuntimeConfig,
249    code: &str,
250    timeout: u64,
251) -> Result<(String, String, i32), String> {
252    let tmp_dir = std::env::temp_dir().join("lean-ctx-sandbox");
253    let _ = std::fs::create_dir_all(&tmp_dir);
254
255    let env_pairs: Vec<(String, String)> = runtime
256        .env
257        .iter()
258        .map(|(k, v)| (k.clone(), v.clone()))
259        .collect();
260
261    if runtime.needs_temp_file {
262        let suffix = format!(".{}", runtime.file_extension);
263        let tmp = tempfile::Builder::new()
264            .prefix("exec_")
265            .suffix(&suffix)
266            .tempfile_in(&tmp_dir)
267            .map_err(|e| format!("Failed to create temp file: {e}"))?;
268        let file_path = tmp.into_temp_path();
269        std::fs::write(&file_path, code).map_err(|e| format!("Failed to write temp file: {e}"))?;
270
271        let allowed = [file_path.to_path_buf()];
272        let allowed_refs: Vec<&std::path::Path> =
273            allowed.iter().map(std::path::PathBuf::as_path).collect();
274        let file_str = file_path.to_string_lossy().to_string();
275
276        let mut args: Vec<&str> = runtime
277            .args
278            .iter()
279            .map(std::string::String::as_str)
280            .collect();
281        args.push(&file_str);
282
283        let result = super::sandbox_seatbelt::execute_sandboxed(
284            &runtime.command,
285            &args,
286            &allowed_refs,
287            &env_pairs,
288            timeout,
289        );
290        let _ = std::fs::remove_file(&file_path);
291        result
292    } else {
293        let mut args: Vec<&str> = runtime
294            .args
295            .iter()
296            .map(std::string::String::as_str)
297            .collect();
298        args.push(code);
299        super::sandbox_seatbelt::execute_sandboxed(
300            &runtime.command,
301            &args,
302            &[],
303            &env_pairs,
304            timeout,
305        )
306    }
307}
308
309#[cfg(target_os = "linux")]
310fn landlock_execute(
311    runtime: &RuntimeConfig,
312    code: &str,
313    timeout: u64,
314) -> Result<(String, String, i32), String> {
315    let tmp_dir = std::env::temp_dir().join("lean-ctx-sandbox");
316    let _ = std::fs::create_dir_all(&tmp_dir);
317
318    let env_pairs: Vec<(String, String)> = runtime
319        .env
320        .iter()
321        .map(|(k, v)| (k.clone(), v.clone()))
322        .collect();
323
324    if runtime.needs_temp_file {
325        let suffix = format!(".{}", runtime.file_extension);
326        let tmp = tempfile::Builder::new()
327            .prefix("exec_")
328            .suffix(&suffix)
329            .tempfile_in(&tmp_dir)
330            .map_err(|e| format!("Failed to create temp file: {e}"))?;
331        let file_path = tmp.into_temp_path();
332        std::fs::write(&file_path, code).map_err(|e| format!("Failed to write temp file: {e}"))?;
333
334        let allowed = [file_path.to_path_buf()];
335        let allowed_refs: Vec<&std::path::Path> =
336            allowed.iter().map(std::path::PathBuf::as_path).collect();
337        let file_str = file_path.to_string_lossy().to_string();
338
339        let mut args: Vec<&str> = runtime
340            .args
341            .iter()
342            .map(std::string::String::as_str)
343            .collect();
344        args.push(&file_str);
345
346        let result = super::sandbox_landlock::execute_sandboxed(
347            &runtime.command,
348            &args,
349            &allowed_refs,
350            &env_pairs,
351            timeout,
352        );
353        let _ = std::fs::remove_file(&file_path);
354        result
355    } else {
356        let mut args: Vec<&str> = runtime
357            .args
358            .iter()
359            .map(std::string::String::as_str)
360            .collect();
361        args.push(code);
362        super::sandbox_landlock::execute_sandboxed(
363            &runtime.command,
364            &args,
365            &[],
366            &env_pairs,
367            timeout,
368        )
369    }
370}
371
372const SANDBOX_ENV_ALLOWLIST: &[&str] = &[
373    "PATH",
374    "HOME",
375    "USER",
376    "LANG",
377    "LC_ALL",
378    "TERM",
379    "TMPDIR",
380    "TMP",
381    "TEMP",
382    "SYSTEMROOT",
383    "WINDIR",
384];
385
386fn apply_sandbox_env(cmd: &mut Command, runtime: &RuntimeConfig) {
387    cmd.env_clear();
388    for key in SANDBOX_ENV_ALLOWLIST {
389        if let Ok(val) = std::env::var(key) {
390            cmd.env(key, val);
391        }
392    }
393    for (k, v) in &runtime.env {
394        cmd.env(k, v);
395    }
396    cmd.env("LEAN_CTX_SANDBOX", "1");
397}
398
399fn execute_with_stdin(
400    runtime: &RuntimeConfig,
401    code: &str,
402    timeout: u64,
403) -> Result<(String, String, i32), String> {
404    let mut cmd = Command::new(&runtime.command);
405    for arg in &runtime.args {
406        cmd.arg(arg);
407    }
408    cmd.arg(code);
409    apply_sandbox_env(&mut cmd, runtime);
410    cmd.stdout(std::process::Stdio::piped());
411    cmd.stderr(std::process::Stdio::piped());
412
413    let child = cmd
414        .spawn()
415        .map_err(|e| format!("Failed to spawn {}: {e}", runtime.command))?;
416
417    let output = wait_with_timeout(child, timeout)?;
418    Ok((
419        crate::shell::decode_output(&output.stdout),
420        crate::shell::decode_output(&output.stderr),
421        output.status.code().unwrap_or(1),
422    ))
423}
424
425fn execute_with_file(
426    runtime: &RuntimeConfig,
427    code: &str,
428    timeout: u64,
429) -> Result<(String, String, i32), String> {
430    let tmp_dir = std::env::temp_dir().join("lean-ctx-sandbox");
431    let _ = std::fs::create_dir_all(&tmp_dir);
432
433    let suffix = format!(".{}", runtime.file_extension);
434    let tmp = tempfile::Builder::new()
435        .prefix("exec_")
436        .suffix(&suffix)
437        .tempfile_in(&tmp_dir)
438        .map_err(|e| format!("Failed to create temp file: {e}"))?;
439    let file_path = tmp.into_temp_path();
440
441    std::fs::write(&file_path, code).map_err(|e| format!("Failed to write temp file: {e}"))?;
442
443    let result = if runtime.command == "rustc_script" {
444        execute_rust(&file_path, timeout)
445    } else {
446        let mut cmd = Command::new(&runtime.command);
447        for arg in &runtime.args {
448            cmd.arg(arg);
449        }
450        cmd.arg(&file_path);
451        apply_sandbox_env(&mut cmd, runtime);
452        cmd.stdout(std::process::Stdio::piped());
453        cmd.stderr(std::process::Stdio::piped());
454
455        let child = cmd
456            .spawn()
457            .map_err(|e| format!("Failed to spawn {}: {e}", runtime.command))?;
458        let output = wait_with_timeout(child, timeout)?;
459        Ok((
460            crate::shell::decode_output(&output.stdout),
461            crate::shell::decode_output(&output.stderr),
462            output.status.code().unwrap_or(1),
463        ))
464    };
465
466    let _ = std::fs::remove_file(&file_path);
467    result
468}
469
470fn execute_rust(
471    source_path: &std::path::Path,
472    timeout: u64,
473) -> Result<(String, String, i32), String> {
474    let binary_path = source_path.with_extension("");
475
476    let mut compile_cmd = Command::new("rustc");
477    compile_cmd.arg(source_path).arg("-o").arg(&binary_path);
478    compile_cmd.env_clear();
479    for key in SANDBOX_ENV_ALLOWLIST {
480        if let Ok(val) = std::env::var(key) {
481            compile_cmd.env(key, val);
482        }
483    }
484    compile_cmd.env("LEAN_CTX_SANDBOX", "1");
485
486    let compile = compile_cmd
487        .output()
488        .map_err(|e| format!("rustc not found: {e}"))?;
489
490    if !compile.status.success() {
491        let stderr = crate::shell::decode_output(&compile.stderr);
492        let _ = std::fs::remove_file(&binary_path);
493        return Ok((String::new(), stderr, compile.status.code().unwrap_or(1)));
494    }
495
496    let mut run_cmd = Command::new(&binary_path);
497    run_cmd.env_clear();
498    for key in SANDBOX_ENV_ALLOWLIST {
499        if let Ok(val) = std::env::var(key) {
500            run_cmd.env(key, val);
501        }
502    }
503    run_cmd.env("LEAN_CTX_SANDBOX", "1");
504    run_cmd.stdout(std::process::Stdio::piped());
505    run_cmd.stderr(std::process::Stdio::piped());
506
507    let child = run_cmd
508        .spawn()
509        .map_err(|e| format!("Failed to run compiled binary: {e}"))?;
510
511    let output = wait_with_timeout(child, timeout)?;
512    let _ = std::fs::remove_file(&binary_path);
513
514    Ok((
515        crate::shell::decode_output(&output.stdout),
516        crate::shell::decode_output(&output.stderr),
517        output.status.code().unwrap_or(1),
518    ))
519}
520
521fn wait_with_timeout(
522    child: std::process::Child,
523    timeout_secs: u64,
524) -> Result<std::process::Output, String> {
525    let mut child = child;
526    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(timeout_secs);
527
528    loop {
529        match child.try_wait() {
530            Ok(Some(_)) => return child.wait_with_output().map_err(|e| e.to_string()),
531            Ok(None) => {
532                if std::time::Instant::now() > deadline {
533                    let _ = child.kill();
534                    return Err(format!("Execution timed out after {timeout_secs}s"));
535                }
536                std::thread::sleep(std::time::Duration::from_millis(50));
537            }
538            Err(e) => return Err(e.to_string()),
539        }
540    }
541}
542
543fn find_binary(candidates: &[&str]) -> Option<String> {
544    for name in candidates {
545        if which_exists(name) {
546            return Some(name.to_string());
547        }
548    }
549    None
550}
551
552fn which_exists(name: &str) -> bool {
553    #[cfg(target_os = "windows")]
554    let check_cmd = Command::new("where")
555        .arg(name)
556        .stdout(std::process::Stdio::null())
557        .stderr(std::process::Stdio::null())
558        .status();
559
560    #[cfg(not(target_os = "windows"))]
561    let check_cmd = Command::new("which")
562        .arg(name)
563        .stdout(std::process::Stdio::null())
564        .stderr(std::process::Stdio::null())
565        .status();
566
567    check_cmd.is_ok_and(|s| s.success())
568}
569
570fn truncate_output(output: &str) -> String {
571    if output.len() <= MAX_OUTPUT_BYTES {
572        return output.to_string();
573    }
574    truncate_smart(output, MAX_OUTPUT_BYTES)
575}
576
577fn truncate_smart(output: &str, max_bytes: usize) -> String {
578    if output.len() <= max_bytes {
579        return output.to_string();
580    }
581
582    let lines: Vec<&str> = output.lines().collect();
583    let total_lines = lines.len();
584
585    let head_count = (total_lines * 60) / 100;
586    let tail_count = total_lines - head_count;
587
588    let head: Vec<&str> = lines.iter().take(head_count).copied().collect();
589    let tail: Vec<&str> = lines
590        .iter()
591        .skip(total_lines - tail_count)
592        .copied()
593        .collect();
594
595    let head_text = head.join("\n");
596    let tail_text = tail.join("\n");
597
598    if head_text.len() + tail_text.len() + 100 > max_bytes {
599        let half = max_bytes / 2;
600        let h = &output[..output.floor_char_boundary(half.min(output.len()))];
601        let t_start = output.ceil_char_boundary(output.len().saturating_sub(half));
602        let t = &output[t_start..];
603        let skipped = output.len() - h.len() - t.len();
604        return format!("{h}\n\n... [{skipped} bytes truncated — showing head + tail] ...\n\n{t}");
605    }
606
607    let skipped_lines = total_lines - head_count - tail_count;
608    let skipped_bytes = output.len() - head_text.len() - tail_text.len();
609    format!(
610        "{head_text}\n\n... [{skipped_lines} lines / {skipped_bytes} bytes truncated — showing first {head_count} + last {tail_count} lines] ...\n\n{tail_text}"
611    )
612}
613
614pub fn supported_languages() -> &'static [&'static str] {
615    &[
616        "javascript",
617        "typescript",
618        "python",
619        "shell",
620        "ruby",
621        "go",
622        "rust",
623        "php",
624        "perl",
625        "r",
626        "elixir",
627    ]
628}
629
630#[cfg(test)]
631mod tests {
632    use super::*;
633
634    fn python_available() -> bool {
635        find_binary(&["python3", "python"]).is_some()
636    }
637
638    #[test]
639    fn execute_python_hello() {
640        if !python_available() {
641            return;
642        }
643        let result = execute("python", "print('hello sandbox')", None);
644        assert_eq!(result.exit_code, 0);
645        assert!(result.stdout.contains("hello sandbox"));
646    }
647
648    #[test]
649    #[cfg(not(target_os = "windows"))]
650    fn execute_shell_echo() {
651        let result = execute("shell", "echo 'test output'", None);
652        assert_eq!(result.exit_code, 0);
653        assert!(result.stdout.contains("test output"));
654    }
655
656    #[test]
657    fn execute_unsupported_language() {
658        let result = execute("brainfuck", "++++", None);
659        assert_eq!(result.exit_code, 1);
660        assert!(result.stderr.contains("Unsupported language"));
661    }
662
663    #[test]
664    fn execute_rejects_oversized_code() {
665        let huge = "a".repeat(MAX_CODE_BYTES + 1);
666        let result = execute("python", &huge, None);
667        assert_eq!(result.exit_code, 1);
668        assert!(result.stderr.contains("exceeds the"));
669    }
670
671    #[test]
672    fn execute_python_error() {
673        if !python_available() {
674            return;
675        }
676        let result = execute("python", "raise ValueError('test error')", None);
677        assert_ne!(result.exit_code, 0);
678        assert!(result.stderr.contains("ValueError"));
679    }
680
681    #[test]
682    fn execute_with_timeout() {
683        if !python_available() {
684            return;
685        }
686        let result = execute("python", "import time; time.sleep(60)", Some(1));
687        assert_ne!(result.exit_code, 0);
688    }
689
690    #[test]
691    fn truncate_preserves_head_and_tail() {
692        let lines: Vec<String> = (0..100)
693            .map(|i| format!("line {i}: some content here"))
694            .collect();
695        let output = lines.join("\n");
696        let truncated = truncate_smart(&output, 500);
697        assert!(truncated.contains("line 0:"));
698        assert!(truncated.contains("line 99:"));
699        assert!(truncated.contains("truncated"));
700    }
701
702    #[test]
703    fn supported_languages_list() {
704        let langs = supported_languages();
705        assert!(langs.contains(&"python"));
706        assert!(langs.contains(&"javascript"));
707        assert!(langs.contains(&"rust"));
708        assert_eq!(langs.len(), 11);
709    }
710
711    #[test]
712    fn sandbox_env_is_set() {
713        if !python_available() {
714            return;
715        }
716        let result = execute(
717            "python",
718            "import os; print(os.environ.get('LEAN_CTX_SANDBOX', 'missing'))",
719            None,
720        );
721        assert_eq!(result.exit_code, 0);
722        assert!(result.stdout.contains('1'));
723    }
724
725    #[test]
726    #[cfg(not(target_os = "windows"))]
727    fn batch_execute_multiple() {
728        let items = vec![
729            ("python".to_string(), "print(1+1)".to_string()),
730            ("shell".to_string(), "echo hello".to_string()),
731        ];
732        let results = batch_execute(&items);
733        assert_eq!(results.len(), 2);
734        assert!(results[0].stdout.contains('2'));
735        assert!(results[1].stdout.contains("hello"));
736    }
737}