vetto 0.3.7

//! cgroup v2 transient lifecycle and resource quota management.
//!
//! Handles cgroup creation in delegated hierarchies, writes resource ceilings
//! (memory.max, memory.swap.max, pids.max, cpu.max), migrates the child process,
//! and ensures cleanup on teardown or SIGKILL.

use std::fs;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use crate::error::{VettoError, VettoResult};
use crate::policy::CgroupConfig;

#[derive(Debug)]
pub struct CgroupHandle {
    path: PathBuf,
    cleaned: Arc<AtomicBool>,
}

impl CgroupHandle {
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Add a process PID to this cgroup.
    pub fn add_process(&self, pid: u32) -> VettoResult<()> {
        let procs_file = self.path.join("cgroup.procs");
        fs::write(&procs_file, pid.to_string()).map_err(|e| {
            VettoError::Sandbox(format!(
                "failed to move pid {pid} into cgroup {}: {e}",
                self.path.display()
            ))
        })
    }

    /// Clean up the cgroup directory.
    pub fn cleanup(&self) {
        // Kill remaining procs once if cgroup.kill is available (Linux 5.14+)
        if !self.cleaned.swap(true, Ordering::SeqCst) {
            let kill_file = self.path.join("cgroup.kill");
            if kill_file.exists() {
                let _ = fs::write(&kill_file, "1");
            }
        }
        // Attempt removing the directory on every cleanup invocation until success
        if self.path.exists() {
            let _ = fs::remove_dir(&self.path);
        }
    }
}

impl Drop for CgroupHandle {
    fn drop(&mut self) {
        self.cleanup();
    }
}

/// Parse human-readable memory limit into bytes or string representation.
pub fn parse_memory_bytes(input: &str) -> Option<String> {
    let s = input.trim();
    if s.is_empty() || s.eq_ignore_ascii_case("max") {
        return Some("max".to_string());
    }
    let (num_part, unit_part) = match s.find(|c: char| !c.is_ascii_digit() && c != '.') {
        Some(idx) => (&s[..idx], s[idx..].trim().to_uppercase()),
        None => (s, String::new()),
    };
    let num: f64 = num_part.parse().ok()?;
    let multiplier: f64 = match unit_part.as_str() {
        "" | "B" => 1.0,
        "K" | "KB" | "KIB" => 1024.0,
        "M" | "MB" | "MIB" => 1024.0 * 1024.0,
        "G" | "GB" | "GIB" => 1024.0 * 1024.0 * 1024.0,
        "T" | "TB" | "TIB" => 1024.0 * 1024.0 * 1024.0 * 1024.0,
        _ => return None,
    };
    let bytes = (num * multiplier) as u64;
    Some(bytes.to_string())
}

/// Parse CPU limit (e.g. "50%", "100%", "200%", or raw quota/period "50000 100000").
pub fn parse_cpu_max(input: &str) -> Option<String> {
    let s = input.trim();
    if s.is_empty() || s.eq_ignore_ascii_case("max") {
        return Some("max 100000".to_string());
    }
    if s.ends_with('%') {
        let pct_str = s.trim_end_matches('%').trim();
        let pct: f64 = pct_str.parse().ok()?;
        let period = 100_000u64;
        let quota = ((pct / 100.0) * period as f64) as u64;
        return Some(format!("{quota} {period}"));
    }
    if s.contains(' ') {
        return Some(s.to_string());
    }
    if let Ok(quota) = s.parse::<u64>() {
        return Some(format!("{quota} 100000"));
    }
    None
}

/// Read available cgroup v2 controllers from the cgroup root if mounted.
pub fn available_controllers() -> Vec<String> {
    if let Ok(content) = fs::read_to_string("/sys/fs/cgroup/cgroup.controllers") {
        content.split_whitespace().map(|s| s.to_string()).collect()
    } else {
        Vec::new()
    }
}

/// Locate a writable cgroup v2 hierarchy.
pub fn find_cgroup_root() -> Option<PathBuf> {
    if std::env::var_os("VETTO_TEST_NO_CGROUP").is_some() {
        return None;
    }
    let cgroup2_mount = Path::new("/sys/fs/cgroup");
    if !cgroup2_mount.join("cgroup.controllers").exists() {
        return None;
    }

    // 1. Check /proc/self/cgroup to attach to caller's current delegated cgroup subtree
    if let Ok(cgroup_content) = fs::read_to_string("/proc/self/cgroup") {
        for line in cgroup_content.lines() {
            // format: 0::<path>
            if let Some(path_part) = line.strip_prefix("0::") {
                let rel = path_part.trim().trim_start_matches('/');
                let mut cur = cgroup2_mount.join(rel);
                while cur.starts_with(cgroup2_mount) {
                    if is_dir_writable(&cur) {
                        return Some(cur);
                    }
                    if let Some(parent) = cur.parent() {
                        cur = parent.to_path_buf();
                    } else {
                        break;
                    }
                }
            }
        }
    }

    // 2. Check user slice under systemd: /sys/fs/cgroup/user.slice/user-<uid>.slice/user@<uid>.service/
    let uid = unsafe { libc::getuid() };
    let user_slice = cgroup2_mount.join(format!("user.slice/user-{uid}.slice/user@{uid}.service"));
    if is_dir_writable(&user_slice) {
        return Some(user_slice);
    }

    // 3. Check general user.slice
    let user_slice_general = cgroup2_mount.join(format!("user.slice/user-{uid}.slice"));
    if is_dir_writable(&user_slice_general) {
        return Some(user_slice_general);
    }

    // 4. Direct cgroup root if running privileged/root
    if is_dir_writable(cgroup2_mount) {
        return Some(cgroup2_mount.to_path_buf());
    }

    // 5. Check user home cgroup (~/.cgroup or $XDG_RUNTIME_DIR/cgroup)
    if let Some(runtime_dir) = std::env::var_os("XDG_RUNTIME_DIR") {
        let runtime_cgroup = PathBuf::from(runtime_dir).join("cgroup");
        if is_dir_writable(&runtime_cgroup) {
            return Some(runtime_cgroup);
        }
    }

    None
}

fn is_dir_writable(path: &Path) -> bool {
    if !path.is_dir() {
        return false;
    }
    let test_file = path.join(format!(".vetto-write-test-{}", std::process::id()));
    if fs::write(&test_file, b"test").is_ok() {
        let _ = fs::remove_file(&test_file);
        true
    } else {
        false
    }
}

/// Create a transient cgroup v2 scope for the session.
pub fn setup_cgroup(
    cgroup_config: Option<&CgroupConfig>,
    cpu_max_override: Option<&str>,
) -> VettoResult<Option<CgroupHandle>> {
    let quotas_mandated = cgroup_config.is_some()
        || cpu_max_override.is_some()
        || std::env::var_os("VETTO_REQUIRE_CGROUP").is_some();

    let effective_cgroup = match (cgroup_config, cpu_max_override) {
        (None, None) => CgroupConfig::default(),
        (Some(c), None) => c.clone(),
        (None, Some(cpu)) => CgroupConfig {
            cpu_max: Some(cpu.to_string()),
            ..CgroupConfig::default()
        },
        (Some(c), Some(cpu)) => {
            let mut merged = c.clone();
            merged.cpu_max = Some(cpu.to_string());
            merged
        }
    };

    let has_quotas = effective_cgroup.memory_max.is_some()
        || effective_cgroup.swap_max.is_some()
        || effective_cgroup.pids_max.is_some()
        || effective_cgroup.cpu_max.is_some();

    let is_required = quotas_mandated || has_quotas;

    // Validate quota specifications early: fail-closed if invalid
    if let Some(mem) = &effective_cgroup.memory_max {
        if parse_memory_bytes(mem).is_none() && is_required {
            return Err(VettoError::Sandbox(format!(
                "invalid memory_max spec '{mem}' (fail-closed exit 125)"
            )));
        }
    }
    if let Some(swap) = &effective_cgroup.swap_max {
        if parse_memory_bytes(swap).is_none() && is_required {
            return Err(VettoError::Sandbox(format!(
                "invalid swap_max spec '{swap}' (fail-closed exit 125)"
            )));
        }
    }
    if let Some(cpu) = &effective_cgroup.cpu_max {
        if parse_cpu_max(cpu).is_none() && is_required {
            return Err(VettoError::Sandbox(format!(
                "invalid cpu_max spec '{cpu}' (fail-closed exit 125)"
            )));
        }
    }

    let Some(root) = find_cgroup_root() else {
        if is_required {
            return Err(VettoError::Sandbox(
                "cgroup v2 is unavailable or not writable on this system; \
                 cannot enforce mandated cgroup resource quotas (fail-closed exit 125)"
                    .into(),
            ));
        }
        tracing::debug!(
            "cgroup v2 is unavailable or not writable on this system; \
             continuing without cgroup resource quotas"
        );
        return Ok(None);
    };

    // Enable subtree controllers in parent if possible
    let subtree_file = root.join("cgroup.subtree_control");
    if subtree_file.exists() {
        let _ = fs::write(&subtree_file, "+memory +pids +cpu");
    }

    let nonce = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0);
    let cgroup_dir = root.join(format!("vetto-session-{}-{}", std::process::id(), nonce));

    if let Err(e) = fs::create_dir(&cgroup_dir) {
        if is_required {
            return Err(VettoError::Sandbox(format!(
                "failed to create cgroup directory {}: {e} (fail-closed exit 125)",
                cgroup_dir.display()
            )));
        }
        tracing::debug!(
            "failed to create cgroup directory {}: {e}; continuing without cgroup",
            cgroup_dir.display()
        );
        return Ok(None);
    }

    // Write limits
    if let Some(mem) = &effective_cgroup.memory_max {
        let bytes = parse_memory_bytes(mem).ok_or_else(|| {
            if is_required {
                let _ = fs::remove_dir(&cgroup_dir);
            }
            VettoError::Sandbox(format!(
                "invalid memory_max spec '{mem}' (fail-closed exit 125)"
            ))
        })?;
        if let Err(e) = fs::write(cgroup_dir.join("memory.max"), &bytes) {
            if is_required {
                let _ = fs::remove_dir(&cgroup_dir);
                return Err(VettoError::Sandbox(format!(
                    "failed to write memory.max ({bytes}): {e} (fail-closed exit 125)"
                )));
            }
        }
    }
    if let Some(swap) = &effective_cgroup.swap_max {
        let bytes = parse_memory_bytes(swap).ok_or_else(|| {
            if is_required {
                let _ = fs::remove_dir(&cgroup_dir);
            }
            VettoError::Sandbox(format!(
                "invalid swap_max spec '{swap}' (fail-closed exit 125)"
            ))
        })?;
        if let Err(e) = fs::write(cgroup_dir.join("memory.swap.max"), &bytes) {
            if is_required {
                let _ = fs::remove_dir(&cgroup_dir);
                return Err(VettoError::Sandbox(format!(
                    "failed to write memory.swap.max ({bytes}): {e} (fail-closed exit 125)"
                )));
            }
        }
    }
    if let Some(pids) = &effective_cgroup.pids_max {
        if let Err(e) = fs::write(cgroup_dir.join("pids.max"), pids) {
            if is_required {
                let _ = fs::remove_dir(&cgroup_dir);
                return Err(VettoError::Sandbox(format!(
                    "failed to write pids.max ({pids}): {e} (fail-closed exit 125)"
                )));
            }
        }
    }
    if let Some(cpu) = &effective_cgroup.cpu_max {
        let val = parse_cpu_max(cpu).ok_or_else(|| {
            if is_required {
                let _ = fs::remove_dir(&cgroup_dir);
            }
            VettoError::Sandbox(format!(
                "invalid cpu_max spec '{cpu}' (fail-closed exit 125)"
            ))
        })?;
        if let Err(e) = fs::write(cgroup_dir.join("cpu.max"), &val) {
            if is_required {
                let _ = fs::remove_dir(&cgroup_dir);
                return Err(VettoError::Sandbox(format!(
                    "failed to write cpu.max ({val}): {e} (fail-closed exit 125)"
                )));
            }
        }
    }

    Ok(Some(CgroupHandle {
        path: cgroup_dir,
        cleaned: Arc::new(AtomicBool::new(false)),
    }))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_memory_units() {
        assert_eq!(parse_memory_bytes("2g"), Some("2147483648".into()));
        assert_eq!(parse_memory_bytes("512M"), Some("536870912".into()));
        assert_eq!(parse_memory_bytes("0"), Some("0".into()));
        assert_eq!(parse_memory_bytes("max"), Some("max".into()));
        assert_eq!(parse_memory_bytes("1024"), Some("1024".into()));
    }

    #[test]
    fn parse_cpu_percent_and_raw() {
        assert_eq!(parse_cpu_max("50%"), Some("50000 100000".into()));
        assert_eq!(parse_cpu_max("100%"), Some("100000 100000".into()));
        assert_eq!(parse_cpu_max("200%"), Some("200000 100000".into()));
        assert_eq!(parse_cpu_max("50000 100000"), Some("50000 100000".into()));
        assert_eq!(parse_cpu_max("max"), Some("max 100000".into()));
    }

    #[test]
    fn test_cgroup_root_or_graceful_none() {
        let _root = find_cgroup_root();
        let scope = setup_cgroup(None, None);
        assert!(scope.is_ok());
    }

    #[test]
    fn test_mandated_cgroup_fails_closed_when_unavailable() {
        unsafe {
            let pid = libc::fork();
            assert!(pid >= 0, "fork failed");
            if pid == 0 {
                std::env::set_var("VETTO_TEST_NO_CGROUP", "1");
                let cfg = CgroupConfig {
                    memory_max: Some("2g".into()),
                    pids_max: Some("128".into()),
                    ..CgroupConfig::default()
                };
                let res = setup_cgroup(Some(&cfg), None);
                let ok = match res {
                    Err(e)
                        if e.exit_code() == crate::exit_codes::EXIT_FAIL_CLOSED
                            && e.to_string().contains("fail-closed exit 125") =>
                    {
                        0
                    }
                    _ => 1,
                };
                libc::_exit(ok);
            }
            let mut status = 0;
            libc::waitpid(pid, &mut status, 0);
            assert!(
                libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
                "mandated cgroup test failed in child"
            );
        }
    }

    #[test]
    fn test_invalid_quota_fails_closed() {
        let cfg_mem = CgroupConfig {
            memory_max: Some("invalid_bytes_spec".into()),
            ..CgroupConfig::default()
        };
        let res_mem = setup_cgroup(Some(&cfg_mem), None);
        assert!(res_mem.is_err());
        let err_mem = res_mem.err().unwrap();
        assert_eq!(err_mem.exit_code(), crate::exit_codes::EXIT_FAIL_CLOSED);
        assert!(err_mem.to_string().contains("fail-closed exit 125"));

        let cfg_cpu = CgroupConfig {
            cpu_max: Some("invalid_cpu_percent".into()),
            ..CgroupConfig::default()
        };
        let res_cpu = setup_cgroup(Some(&cfg_cpu), None);
        assert!(res_cpu.is_err());
        let err_cpu = res_cpu.err().unwrap();
        assert_eq!(err_cpu.exit_code(), crate::exit_codes::EXIT_FAIL_CLOSED);
        assert!(err_cpu.to_string().contains("fail-closed exit 125"));
    }

    #[test]
    fn test_cgroup_handle_idempotent_cleanup() {
        let temp_dir =
            std::env::temp_dir().join(format!("vetto-cgroup-test-{}", std::process::id()));
        let _ = fs::create_dir(&temp_dir);
        assert!(temp_dir.exists());

        let handle = CgroupHandle {
            path: temp_dir.clone(),
            cleaned: Arc::new(AtomicBool::new(false)),
        };

        // First cleanup removes directory
        handle.cleanup();
        assert!(!temp_dir.exists());

        // Second cleanup is an idempotent no-op without errors
        handle.cleanup();
        assert!(!temp_dir.exists());
    }
}