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