AWG2_harness/traffic_core/src/swap.rs

use std::{
    fs,
    io::Write,
    path::{Path, PathBuf},
    process::Command,
    time::Duration,
};

#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;

use anyhow::Context;

use crate::{config::Config, lock::FileLock};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ActiveSwap {
    pub name: PathBuf,
    pub kind: String,
    pub size_bytes: u64,
    pub priority: i32,
}

#[derive(Debug, Clone)]
struct DesiredSwap {
    path: PathBuf,
    size_bytes: u64,
    priority: i32,
    swappiness: u8,
}

#[derive(Debug, Clone)]
pub struct SwapStatus {
    pub enabled: bool,
    pub desired_path: PathBuf,
    pub active_entries: Vec<ActiveSwap>,
    pub desired_active: bool,
    pub exclusive: bool,
    pub fstab_swap_entries: Vec<String>,
    pub fstab_exclusive: bool,
    pub current_swappiness: Option<u8>,
    pub persisted_swappiness: Option<u8>,
    pub drift_reason: Option<String>,
}

pub fn ensure_swap(cfg: &Config) -> anyhow::Result<()> {
    if !cfg.swap.enabled {
        return Ok(());
    }

    let _lock = FileLock::acquire_timeout("/run/awg2-harness/swap.lock", Duration::from_secs(30))?;

    let desired = DesiredSwap::from_config(cfg)?;
    let active_before = read_active_swaps()?;
    disable_foreign_swaps(&desired, &active_before)?;
    ensure_desired_swapfile(&desired)?;
    rewrite_fstab_exclusive(&desired)?;
    enable_desired_swap(&desired)?;
    apply_swappiness(desired.swappiness)?;
    verify_exclusive_swap_state(&desired)?;
    Ok(())
}

pub fn inspect_swap(cfg: &Config) -> anyhow::Result<SwapStatus> {
    let desired_path = PathBuf::from(cfg.swap.path.clone());
    let active_entries = read_active_swaps()?;
    let desired_active = active_entries.iter().any(|s| s.name == desired_path);
    let exclusive = active_entries.len() == 1 && desired_active;
    let fstab_swap_entries = read_fstab_swap_entries("/etc/fstab")?;
    let expected_line = render_fstab_line(&desired_path, cfg.swap.priority);
    let fstab_exclusive = fstab_swap_entries.len() == 1
        && normalize_ws(&fstab_swap_entries[0])
            == normalize_ws(&expected_line.trim_end().to_string());
    let current_swappiness = read_current_swappiness().ok();
    let persisted_swappiness = read_persisted_swappiness()?;

    let drift_reason = if !cfg.swap.enabled {
        None
    } else if !exclusive {
        if let Some(foreign) = active_entries.iter().find(|s| s.name != desired_path) {
            Some(format!("foreign active swap: {}", foreign.name.display()))
        } else {
            Some("desired swap is not exclusively active".to_string())
        }
    } else if !fstab_exclusive {
        Some("/etc/fstab swap entries are not exclusive".to_string())
    } else if current_swappiness != Some(cfg.swap.swappiness) {
        Some(format!(
            "runtime swappiness mismatch: current={:?} desired={}",
            current_swappiness, cfg.swap.swappiness
        ))
    } else if persisted_swappiness != Some(cfg.swap.swappiness) {
        Some(format!(
            "persisted swappiness mismatch: persisted={:?} desired={}",
            persisted_swappiness, cfg.swap.swappiness
        ))
    } else {
        None
    };

    Ok(SwapStatus {
        enabled: cfg.swap.enabled,
        desired_path,
        active_entries,
        desired_active,
        exclusive,
        fstab_swap_entries,
        fstab_exclusive,
        current_swappiness,
        persisted_swappiness,
        drift_reason,
    })
}

impl DesiredSwap {
    fn from_config(cfg: &Config) -> anyhow::Result<Self> {
        let path = PathBuf::from(cfg.swap.path.clone());
        if !path.is_absolute() {
            anyhow::bail!("swap.path must be absolute: {}", path.display());
        }
        let size_bytes = cfg
            .swap
            .size_mb
            .checked_mul(1024 * 1024)
            .ok_or_else(|| anyhow::anyhow!("swap.size_mb overflow"))?;
        Ok(Self {
            path,
            size_bytes,
            priority: cfg.swap.priority,
            swappiness: cfg.swap.swappiness,
        })
    }
}

fn read_active_swaps() -> anyhow::Result<Vec<ActiveSwap>> {
    let out = Command::new("swapon")
        .args([
            "--show",
            "--bytes",
            "--raw",
            "--noheadings",
            "--output",
            "NAME,TYPE,SIZE,PRIO",
        ])
        .output()?;
    if !out.status.success() {
        anyhow::bail!(
            "swapon --show failed: {}",
            String::from_utf8_lossy(&out.stderr)
        );
    }
    parse_swapon_output(&String::from_utf8_lossy(&out.stdout))
}

fn parse_swapon_output(raw: &str) -> anyhow::Result<Vec<ActiveSwap>> {
    let mut out = Vec::new();
    for line in raw.lines() {
        let t = line.trim();
        if t.is_empty() {
            continue;
        }
        let parts: Vec<&str> = t.split_whitespace().collect();
        if parts.len() != 4 {
            anyhow::bail!("malformed swapon row: {t}");
        }
        let size_bytes = parts[2]
            .parse::<u64>()
            .map_err(|_| anyhow::anyhow!("malformed swapon size: {}", parts[2]))?;
        let priority = parts[3]
            .parse::<i32>()
            .map_err(|_| anyhow::anyhow!("malformed swapon priority: {}", parts[3]))?;
        out.push(ActiveSwap {
            name: PathBuf::from(parts[0]),
            kind: parts[1].to_string(),
            size_bytes,
            priority,
        });
    }
    Ok(out)
}

fn disable_foreign_swaps(desired: &DesiredSwap, active: &[ActiveSwap]) -> anyhow::Result<()> {
    for entry in active {
        if entry.name != desired.path {
            run("swapoff", &[entry.name.to_string_lossy().as_ref()])?;
            maybe_delete_foreign_swapfile(entry, &desired.path)?;
        } else if entry.size_bytes != desired.size_bytes || entry.priority != desired.priority {
            run("swapoff", &[desired.path.to_string_lossy().as_ref()])?;
        }
    }
    Ok(())
}

fn maybe_delete_foreign_swapfile(entry: &ActiveSwap, desired_path: &Path) -> anyhow::Result<()> {
    if entry.name == desired_path {
        return Ok(());
    }
    let md = match fs::symlink_metadata(&entry.name) {
        Ok(v) => v,
        Err(_) => return Ok(()),
    };
    if md.file_type().is_symlink() {
        return Ok(());
    }
    if md.file_type().is_file() {
        let _ = fs::remove_file(&entry.name);
    }
    Ok(())
}

fn ensure_desired_swapfile(desired: &DesiredSwap) -> anyhow::Result<()> {
    let p = &desired.path;
    if !p.is_absolute() {
        anyhow::bail!("swap.path must be absolute: {}", p.display());
    }

    let mut need_create = true;
    if p.exists() {
        let md = fs::symlink_metadata(p)?;
        if md.file_type().is_symlink() {
            anyhow::bail!("desired swap path is symlink: {}", p.display());
        }
        if !md.file_type().is_file() {
            anyhow::bail!("desired swap path is not regular file: {}", p.display());
        }

        if md.len() == desired.size_bytes {
            need_create = false;
        } else {
            fs::remove_file(p)?;
        }
    }

    if need_create {
        create_swapfile(p, desired.size_bytes)?;
        run("mkswap", &[p.to_string_lossy().as_ref()])?;
    }

    #[cfg(unix)]
    {
        fs::set_permissions(p, fs::Permissions::from_mode(0o600))?;
    }
    Ok(())
}

fn create_swapfile(path: &Path, size_bytes: u64) -> anyhow::Result<()> {
    let size_mb = size_bytes / (1024 * 1024);
    let p = path.to_string_lossy().to_string();

    let out = Command::new("fallocate")
        .args(["-l", &format!("{}M", size_mb), &p])
        .output()?;
    if out.status.success() {
        return Ok(());
    }

    let out_dd = Command::new("dd")
        .args([
            "if=/dev/zero",
            &format!("of={p}"),
            "bs=1M",
            &format!("count={size_mb}"),
            "status=none",
        ])
        .output()?;
    if !out_dd.status.success() {
        anyhow::bail!(
            "swapfile create failed (fallocate+dd): fallocate={}, dd={}",
            String::from_utf8_lossy(&out.stderr),
            String::from_utf8_lossy(&out_dd.stderr)
        );
    }
    Ok(())
}

fn rewrite_fstab_exclusive(desired: &DesiredSwap) -> anyhow::Result<()> {
    let fstab_path = "/etc/fstab";
    let bak_path = "/etc/fstab.awg2-harness.bak";
    let tmp_path = "/etc/fstab.awg2-harness.tmp";

    let content = fs::read_to_string(fstab_path).unwrap_or_default();
    let rewritten = rewrite_fstab_content(&content, &desired.path, desired.priority);
    if rewritten == content {
        return Ok(());
    }

    let _ = fs::write(bak_path, &content);
    let mut f = fs::File::create(tmp_path)?;
    f.write_all(rewritten.as_bytes())?;
    f.sync_all()?;
    drop(f);
    fs::rename(tmp_path, fstab_path)?;
    Ok(())
}

fn rewrite_fstab_content(content: &str, desired_path: &Path, priority: i32) -> String {
    let mut out = Vec::<String>::new();
    for line in content.lines() {
        let t = line.trim();
        if t.is_empty() || t.starts_with('#') {
            out.push(line.to_string());
            continue;
        }
        let cols: Vec<&str> = t.split_whitespace().collect();
        if cols.len() >= 3 && cols[2] == "swap" {
            continue;
        }
        out.push(line.to_string());
    }
    out.push(
        render_fstab_line(desired_path, priority)
            .trim_end()
            .to_string(),
    );
    let mut s = out.join("\n");
    if !s.ends_with('\n') {
        s.push('\n');
    }
    s
}

fn render_fstab_line(path: &Path, priority: i32) -> String {
    format!("{} none swap sw,pri={} 0 0\n", path.display(), priority)
}

fn enable_desired_swap(desired: &DesiredSwap) -> anyhow::Result<()> {
    let active = read_active_swaps()?;
    if let Some(found) = active.iter().find(|v| v.name == desired.path)
        && found.size_bytes == desired.size_bytes
        && found.priority == desired.priority
    {
        return Ok(());
    }

    run(
        "swapon",
        &[
            "-p",
            &desired.priority.to_string(),
            desired.path.to_string_lossy().as_ref(),
        ],
    )?;
    Ok(())
}

fn apply_swappiness(v: u8) -> anyhow::Result<()> {
    run("sysctl", &["-w", &format!("vm.swappiness={v}")])?;
    let current = read_current_swappiness()?;
    if current != v {
        anyhow::bail!("failed to apply vm.swappiness: got {current}, expected {v}");
    }

    let conf_path = "/etc/sysctl.d/99-awg2-harness.conf";
    let expected = format!("# Managed by awg2-harness\nvm.swappiness = {v}\n");
    let current_file = fs::read_to_string(conf_path).unwrap_or_default();
    if current_file != expected {
        fs::write(conf_path, expected)?;
    }

    let persisted = read_persisted_swappiness()?;
    if persisted != Some(v) {
        anyhow::bail!(
            "failed to persist vm.swappiness: got {:?}, expected {}",
            persisted,
            v
        );
    }
    Ok(())
}

pub fn read_current_swappiness() -> anyhow::Result<u8> {
    let out = Command::new("sysctl")
        .args(["-n", "vm.swappiness"])
        .output()?;
    if !out.status.success() {
        anyhow::bail!("sysctl read vm.swappiness failed");
    }
    let s = String::from_utf8_lossy(&out.stdout).trim().to_string();
    s.parse::<u8>()
        .with_context(|| format!("invalid vm.swappiness value: {s}"))
}

pub fn read_persisted_swappiness() -> anyhow::Result<Option<u8>> {
    let conf_path = "/etc/sysctl.d/99-awg2-harness.conf";
    let content = match fs::read_to_string(conf_path) {
        Ok(v) => v,
        Err(_) => return Ok(None),
    };
    for line in content.lines() {
        let t = line.trim();
        if t.is_empty() || t.starts_with('#') {
            continue;
        }
        let Some((k, v)) = t.split_once('=') else {
            continue;
        };
        if k.trim() == "vm.swappiness" {
            return Ok(Some(v.trim().parse::<u8>()?));
        }
    }
    Ok(None)
}

fn verify_exclusive_swap_state(desired: &DesiredSwap) -> anyhow::Result<()> {
    let active = read_active_swaps()?;
    if active.len() != 1 {
        anyhow::bail!(
            "swap verify failed: expected exactly one active swap, got {}",
            active.len()
        );
    }
    let only = &active[0];
    if only.name != desired.path {
        anyhow::bail!(
            "swap verify failed: active swap {} != desired {}",
            only.name.display(),
            desired.path.display()
        );
    }
    if only.size_bytes != desired.size_bytes {
        anyhow::bail!(
            "swap verify failed: size mismatch {} != {}",
            only.size_bytes,
            desired.size_bytes
        );
    }
    if only.priority != desired.priority {
        anyhow::bail!(
            "swap verify failed: priority mismatch {} != {}",
            only.priority,
            desired.priority
        );
    }

    let entries = read_fstab_swap_entries("/etc/fstab")?;
    if entries.len() != 1 {
        anyhow::bail!(
            "swap verify failed: expected one fstab swap entry, got {}",
            entries.len()
        );
    }
    let expected = render_fstab_line(&desired.path, desired.priority);
    if normalize_ws(&entries[0]) != normalize_ws(expected.trim_end()) {
        anyhow::bail!(
            "swap verify failed: fstab entry mismatch: got='{}' expected='{}'",
            entries[0],
            expected.trim_end()
        );
    }

    let current = read_current_swappiness()?;
    if current != desired.swappiness {
        anyhow::bail!(
            "swap verify failed: vm.swappiness current={} expected={}",
            current,
            desired.swappiness
        );
    }
    Ok(())
}

fn read_fstab_swap_entries(path: &str) -> anyhow::Result<Vec<String>> {
    let content = fs::read_to_string(path).unwrap_or_default();
    let mut out = Vec::new();
    for line in content.lines() {
        let t = line.trim();
        if t.is_empty() || t.starts_with('#') {
            continue;
        }
        let cols: Vec<&str> = t.split_whitespace().collect();
        if cols.len() >= 3 && cols[2] == "swap" {
            out.push(t.to_string());
        }
    }
    Ok(out)
}

fn normalize_ws(s: &str) -> String {
    s.split_whitespace().collect::<Vec<_>>().join(" ")
}

fn run(program: &str, args: &[&str]) -> anyhow::Result<()> {
    let out = Command::new(program).args(args).output()?;
    if !out.status.success() {
        anyhow::bail!(
            "{} {:?} failed: {}",
            program,
            args,
            String::from_utf8_lossy(&out.stderr)
        );
    }
    Ok(())
}

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

    #[test]
    fn parse_swapon_empty() {
        let got = parse_swapon_output("").expect("parse");
        assert!(got.is_empty());
    }

    #[test]
    fn parse_swapon_single_file() {
        let got = parse_swapon_output("/swapfile file 2147479552 1\n").expect("parse");
        assert_eq!(got.len(), 1);
        assert_eq!(got[0].name, PathBuf::from("/swapfile"));
        assert_eq!(got[0].kind, "file");
        assert_eq!(got[0].size_bytes, 2147479552);
        assert_eq!(got[0].priority, 1);
    }

    #[test]
    fn parse_swapon_multiple_entries() {
        let raw = "/swapfile file 2147479552 1\n/dev/zram0 partition 1073741824 100\n";
        let got = parse_swapon_output(raw).expect("parse");
        assert_eq!(got.len(), 2);
    }

    #[test]
    fn parse_swapon_rejects_malformed_size() {
        let err = parse_swapon_output("/swapfile file abc 1\n");
        assert!(err.is_err());
    }

    #[test]
    fn fstab_rewrite_removes_all_swap_entries() {
        let raw =
            "UUID=x / ext4 defaults 0 1\n/swap-old none swap sw 0 0\n/dev/zram0 none swap sw 0 0\n";
        let out = rewrite_fstab_content(raw, Path::new("/swapfile"), 1);
        assert!(!out.contains("/swap-old"));
        assert!(!out.contains("/dev/zram0"));
        assert!(out.contains("/swapfile none swap sw,pri=1 0 0"));
    }

    #[test]
    fn fstab_rewrite_preserves_comments_and_mounts() {
        let raw = "# head\nUUID=x / ext4 defaults 0 1\n\n# tail\n";
        let out = rewrite_fstab_content(raw, Path::new("/swapfile"), 1);
        assert!(out.contains("# head"));
        assert!(out.contains("UUID=x / ext4 defaults 0 1"));
        assert!(out.contains("# tail"));
    }

    #[test]
    fn fstab_rewrite_is_idempotent() {
        let raw = "UUID=x / ext4 defaults 0 1\n/swapfile none swap sw,pri=1 0 0\n";
        let one = rewrite_fstab_content(raw, Path::new("/swapfile"), 1);
        let two = rewrite_fstab_content(&one, Path::new("/swapfile"), 1);
        assert_eq!(one, two);
    }

    #[test]
    fn fstab_rewrite_adds_trailing_newline() {
        let raw = "UUID=x / ext4 defaults 0 1";
        let out = rewrite_fstab_content(raw, Path::new("/swapfile"), 1);
        assert!(out.ends_with('\n'));
    }

    #[test]
    fn swappiness_conf_render_is_stable() {
        let expected = "# Managed by awg2-harness\nvm.swappiness = 1\n";
        assert_eq!(
            expected,
            format!("# Managed by awg2-harness\nvm.swappiness = {}\n", 1)
        );
    }
}