Skip to main content

soma_codemode/runner/
jail.rs

1pub(super) fn reset_execution_jail() {
2    let Some(base) = jail_base_path() else {
3        return;
4    };
5    EXECUTION_JAIL.with(|cell| {
6        let mut cell = cell.borrow_mut();
7        if let Some(previous) = cell.take() {
8            let _ = std::env::set_current_dir(&base);
9            let _ = std::fs::remove_dir_all(previous);
10        }
11        let jail = base.join(format!("exec-{}-{}", std::process::id(), next_jail_seq()));
12        if std::fs::create_dir(&jail).is_ok() && std::env::set_current_dir(&jail).is_ok() {
13            *cell = Some(jail);
14        } else {
15            let _ = std::fs::remove_dir_all(&jail);
16            let _ = std::env::set_current_dir(&base);
17        }
18    });
19}
20
21pub(super) fn cleanup_execution_jail(drop_base: bool) {
22    let base = JAIL_BASE.with(|cell| cell.borrow().as_ref().map(|dir| dir.path().to_path_buf()));
23    EXECUTION_JAIL.with(|cell| {
24        if let Some(jail) = cell.borrow_mut().take() {
25            if let Some(base) = &base {
26                let _ = std::env::set_current_dir(base);
27            }
28            let _ = std::fs::remove_dir_all(jail);
29        }
30    });
31    if drop_base {
32        let _ = std::env::set_current_dir(std::env::temp_dir());
33        JAIL_BASE.with(|cell| {
34            drop(cell.borrow_mut().take());
35        });
36    }
37}
38
39fn jail_base_path() -> Option<std::path::PathBuf> {
40    JAIL_BASE.with(|cell| {
41        let mut cell = cell.borrow_mut();
42        if cell.is_none() {
43            *cell = tempfile::Builder::new()
44                .prefix("soma-codemode-")
45                .tempdir()
46                .ok();
47        }
48        cell.as_ref().map(|dir| dir.path().to_path_buf())
49    })
50}
51
52fn next_jail_seq() -> u64 {
53    JAIL_SEQ.with(|seq| {
54        let next = seq.get();
55        seq.set(next.saturating_add(1));
56        next
57    })
58}
59
60thread_local! {
61    static JAIL_BASE: std::cell::RefCell<Option<tempfile::TempDir>> =
62        const { std::cell::RefCell::new(None) };
63    static EXECUTION_JAIL: std::cell::RefCell<Option<std::path::PathBuf>> =
64        const { std::cell::RefCell::new(None) };
65    static JAIL_SEQ: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
66}