owlen/crates/owlen-core/src/sandbox.rs

use std::path::PathBuf;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

use anyhow::{Context, Result, bail};
use tempfile::TempDir;

/// Configuration options for sandboxed process execution.
#[derive(Clone, Debug)]
pub struct SandboxConfig {
    pub allow_network: bool,
    pub allow_paths: Vec<PathBuf>,
    pub readonly_paths: Vec<PathBuf>,
    pub timeout_seconds: u64,
    pub max_memory_mb: u64,
}

impl Default for SandboxConfig {
    fn default() -> Self {
        Self {
            allow_network: false,
            allow_paths: Vec::new(),
            readonly_paths: Vec::new(),
            timeout_seconds: 30,
            max_memory_mb: 512,
        }
    }
}

/// Wrapper around a bubblewrap sandbox instance.
///
/// Memory limits are enforced via:
/// - bwrap's --rlimit-as (version >= 0.12.0)
/// - prlimit wrapper (fallback for older bwrap versions)
/// - timeout mechanism (always enforced as last resort)
pub struct SandboxedProcess {
    temp_dir: TempDir,
    config: SandboxConfig,
}

impl SandboxedProcess {
    pub fn new(config: SandboxConfig) -> Result<Self> {
        let temp_dir = TempDir::new().context("Failed to create temp directory")?;

        which::which("bwrap")
            .context("bubblewrap not found. Install with: sudo apt install bubblewrap")?;

        Ok(Self { temp_dir, config })
    }

    pub fn execute(&self, command: &str, args: &[&str]) -> Result<SandboxResult> {
        let supports_rlimit = self.supports_rlimit_as();
        let use_prlimit = !supports_rlimit && which::which("prlimit").is_ok();

        let mut cmd = if use_prlimit {
            // Use prlimit wrapper for older bwrap versions
            let mut prlimit_cmd = Command::new("prlimit");
            let memory_limit_bytes = self
                .config
                .max_memory_mb
                .saturating_mul(1024)
                .saturating_mul(1024);
            prlimit_cmd.arg(format!("--as={}", memory_limit_bytes));
            prlimit_cmd.arg("bwrap");
            prlimit_cmd
        } else {
            Command::new("bwrap")
        };

        cmd.args(["--unshare-all", "--die-with-parent", "--new-session"]);

        if self.config.allow_network {
            cmd.arg("--share-net");
        } else {
            cmd.arg("--unshare-net");
        }

        cmd.args(["--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp"]);

        // Bind essential system paths readonly for executables and libraries
        let system_paths = ["/usr", "/bin", "/lib", "/lib64", "/etc"];
        for sys_path in &system_paths {
            let path = std::path::Path::new(sys_path);
            if path.exists() {
                cmd.arg("--ro-bind").arg(sys_path).arg(sys_path);
            }
        }

        // Bind /run for DNS resolution (resolv.conf may be a symlink to /run/systemd/resolve/*)
        if std::path::Path::new("/run").exists() {
            cmd.arg("--ro-bind").arg("/run").arg("/run");
        }

        for path in &self.config.allow_paths {
            let path_host = path.to_string_lossy().into_owned();
            let path_guest = path_host.clone();
            cmd.arg("--bind").arg(&path_host).arg(&path_guest);
        }

        for path in &self.config.readonly_paths {
            let path_host = path.to_string_lossy().into_owned();
            let path_guest = path_host.clone();
            cmd.arg("--ro-bind").arg(&path_host).arg(&path_guest);
        }

        let work_dir = self.temp_dir.path().to_string_lossy().into_owned();
        cmd.arg("--bind").arg(&work_dir).arg("/work");
        cmd.arg("--chdir").arg("/work");

        // Add memory limits via bwrap's --rlimit-as if supported (version >= 0.12.0)
        // If not supported, we use prlimit wrapper (set earlier)
        if supports_rlimit && !use_prlimit {
            let memory_limit_bytes = self
                .config
                .max_memory_mb
                .saturating_mul(1024)
                .saturating_mul(1024);
            let memory_soft = memory_limit_bytes.to_string();
            let memory_hard = memory_limit_bytes.to_string();
            cmd.arg("--rlimit-as").arg(&memory_soft).arg(&memory_hard);
        }

        cmd.arg(command);
        cmd.args(args);

        let start = Instant::now();
        let timeout = Duration::from_secs(self.config.timeout_seconds);

        // Spawn the process instead of waiting immediately
        let mut child = cmd
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .context("Failed to spawn sandboxed command")?;

        let mut was_timeout = false;

        // Wait for the child with timeout
        let output = loop {
            match child.try_wait() {
                Ok(Some(_status)) => {
                    // Process exited
                    let output = child
                        .wait_with_output()
                        .context("Failed to collect process output")?;
                    break output;
                }
                Ok(None) => {
                    // Process still running, check timeout
                    if start.elapsed() >= timeout {
                        // Timeout exceeded, kill the process
                        was_timeout = true;
                        child.kill().context("Failed to kill timed-out process")?;
                        // Wait for the killed process to exit
                        let output = child
                            .wait_with_output()
                            .context("Failed to collect output from killed process")?;
                        break output;
                    }
                    // Sleep briefly before checking again
                    std::thread::sleep(Duration::from_millis(50));
                }
                Err(e) => {
                    bail!("Failed to check process status: {}", e);
                }
            }
        };

        let duration = start.elapsed();

        Ok(SandboxResult {
            stdout: String::from_utf8_lossy(&output.stdout).to_string(),
            stderr: String::from_utf8_lossy(&output.stderr).to_string(),
            exit_code: output.status.code().unwrap_or(-1),
            duration,
            was_timeout,
        })
    }

    /// Check if bubblewrap supports --rlimit-as option (version >= 0.12.0)
    fn supports_rlimit_as(&self) -> bool {
        // Try to get bwrap version
        let output = Command::new("bwrap").arg("--version").output();

        if let Ok(output) = output {
            let version_str = String::from_utf8_lossy(&output.stdout);
            // Parse version like "bubblewrap 0.11.0" or "0.11.0"
            return version_str
                .split_whitespace()
                .last()
                .and_then(|part| {
                    part.split_once('.').and_then(|(major, rest)| {
                        rest.split_once('.').and_then(|(minor, _)| {
                            let maj = major.parse::<u32>().ok()?;
                            let min = minor.parse::<u32>().ok()?;
                            Some((maj, min))
                        })
                    })
                })
                .map(|(maj, min)| maj > 0 || (maj == 0 && min >= 12))
                .unwrap_or(false);
        }

        // If we can't determine the version, assume it doesn't support it (safer default)
        false
    }
}

#[derive(Debug, Clone)]
pub struct SandboxResult {
    pub stdout: String,
    pub stderr: String,
    pub exit_code: i32,
    pub duration: Duration,
    pub was_timeout: bool,
}