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