//! The sandbox helper: the body of `forge-runner __sandbox`.
//!
//! The process is single-threaded and has no tokio runtime, because libcontainer forks and
//! clones. It reads one [`SandboxRequest`] from stdin, prepares the rootfs, creates and starts a
//! container, and waits for it. Its stdout and stderr are the container's, so the runner just
//! forwards them.
//!
//! Setup happens in this order:
//! 1. Enter a user namespace with the request's [`IdMap`](super::idmap::IdMap): root is the
//!    runner's uid, and in subuid mode container ids `1..=n` are the subordinate range. The
//!    overlay mount needs CAP_SYS_ADMIN in a user namespace that owns the mount namespace.
//! 2. An overlay mount of the image rootfs. If that fails, the rootfs is copied instead.
//! 3. The OCI spec is written to `bundle/config.json` (0600) and libcontainer builds the
//!    container from it. The file is deleted right after `build()`, since libcontainer has
//!    read it by then.
//! 4. The container starts and the helper waits, polling so that SIGTERM is never missed. On
//!    SIGTERM the container is killed and deleted. The exit code is the job's.
//!
//! Exit codes: the job's code, 128 + signal if the job was signalled, or 125 for a sandbox setup
//! failure. A setup failure prints `forge-sandbox: <reason>` on stderr.

use std::fs::{self, OpenOptions};
use std::io::{self, Write};
use std::os::fd::AsFd;
use std::os::unix::fs::{MetadataExt, OpenOptionsExt, PermissionsExt, lchown};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use libcontainer::container::Container;
use libcontainer::container::builder::ContainerBuilder;
use libcontainer::syscall::syscall::SyscallType;
use nix::errno::Errno;
use nix::mount::{MntFlags, MsFlags, mount, umount2};
use nix::sys::signal::{SaFlags, SigAction, SigHandler, SigSet, Signal, sigaction};
use nix::sys::wait::{WaitPidFlag, WaitStatus, waitpid};
use nix::unistd::Pid;
use oci_spec::runtime::Spec;

use super::exec::FilteredExecutor;
use super::idmap::enter_mapped_userns;
use super::seccomp::build_filter;
use super::spec::{HOSTNAME, build_spec};
use super::{SandboxError, SandboxRequest};
use crate::workspace::validate_job_id;

/// Exit code for a sandbox setup failure. The runner maps it to a sandbox error.
pub const EXIT_SETUP_FAILED: i32 = 125;
/// How often the wait loop checks for SIGTERM and polls the child.
const POLL: Duration = Duration::from_millis(50);
const HOSTS: &str = "127.0.0.1 localhost\n::1 localhost ip6-localhost ip6-loopback\n";

/// Set by the SIGTERM and SIGINT handler. The wait loop reads it.
static TERMINATE: AtomicBool = AtomicBool::new(false);

extern "C" fn on_terminate(_signal: i32) {
    // Only an atomic store here, which is async-signal-safe.
    TERMINATE.store(true, Ordering::SeqCst);
}

/// Entry point for `forge-runner __sandbox`. Returns the process exit code.
pub fn helper_main() -> i32 {
    // libcontainer logs through tracing. Its logs appear only when the operator sets RUST_LOG.
    if std::env::var_os("RUST_LOG").is_some() {
        let _ = tracing_subscriber::fmt()
            .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
            .with_writer(io::stderr)
            .try_init();
    }
    match run() {
        Ok(code) => code,
        Err(e) => {
            eprintln!("forge-sandbox: {e}");
            EXIT_SETUP_FAILED
        }
    }
}

/// The directories under `job_dir` that the helper uses.
struct Layout {
    job_dir: PathBuf,
    bundle: PathBuf,
    state: PathBuf,
    etc: PathBuf,
    upper: PathBuf,
    work: PathBuf,
    merged: PathBuf,
}

impl Layout {
    fn new(job_dir: &Path) -> Self {
        let overlay = job_dir.join("overlay");
        Self {
            job_dir: job_dir.to_owned(),
            bundle: job_dir.join("bundle"),
            state: job_dir.join("state"),
            etc: job_dir.join("etc"),
            upper: overlay.join("upper"),
            work: overlay.join("work"),
            merged: overlay.join("merged"),
        }
    }
}

fn run() -> Result<i32, SandboxError> {
    let req: SandboxRequest = serde_json::from_reader(io::stdin().lock())
        .map_err(|e| SandboxError::Invalid(format!("cannot read request: {e}")))?;
    validate_request(&req)?;
    install_signal_handlers()?;

    let layout = Layout::new(&req.job_dir);
    prepare_dirs(&layout)?;
    write_etc_files(&layout)?;

    enter_mapped_userns(&req.idmap).map_err(|e| SandboxError::Setup(e.to_string()))?;
    let rootfs = mount_rootfs(&req, &layout)?;

    // Clean up on failure too: the state dir holds the job env.
    let result = run_job(&req, &layout, &rootfs);
    // The container is gone, so nothing writes to the workspace now. In subuid mode the job may
    // have left files owned by subordinate ids, which the runner can't remove, so they go back
    // to the runner.
    if req.idmap.is_subuid() {
        reown_workspace(&req.workspace);
    }
    cleanup(&layout);
    result
}

/// Sets every entry under `workspace` to uid and gid 0 of this namespace, which is the runner's
/// own id on the host. Symlinks are not followed. Entries already owned by 0:0 are left alone.
/// Failures are warnings only, because the job has already run.
fn reown_workspace(workspace: &Path) {
    let mut stack = vec![workspace.to_owned()];
    while let Some(path) = stack.pop() {
        let meta = match fs::symlink_metadata(&path) {
            Ok(meta) => meta,
            Err(e) => {
                warn(format!("cannot stat {}: {e}", path.display()));
                continue;
            }
        };
        if (meta.uid() != 0 || meta.gid() != 0)
            && let Err(e) = lchown(&path, Some(0), Some(0))
        {
            warn(format!("cannot chown {}: {e}", path.display()));
        }
        if !meta.is_dir() {
            continue;
        }
        match fs::read_dir(&path) {
            Ok(entries) => {
                for entry in entries {
                    match entry {
                        Ok(entry) => stack.push(entry.path()),
                        Err(e) => warn(format!("cannot read {}: {e}", path.display())),
                    }
                }
            }
            Err(e) => warn(format!("cannot read {}: {e}", path.display())),
        }
    }
}

fn validate_request(req: &SandboxRequest) -> Result<(), SandboxError> {
    validate_job_id(&req.job_id).map_err(|e| SandboxError::Invalid(e.to_string()))?;
    for (what, path) in [
        ("rootfs", &req.rootfs),
        ("job_dir", &req.job_dir),
        ("workspace", &req.workspace),
    ] {
        if !path.is_absolute() {
            return Err(SandboxError::Invalid(format!("{what} must be absolute")));
        }
    }
    if !req.workspace.is_dir() {
        return Err(SandboxError::Invalid("workspace is not a directory".into()));
    }
    if !req.rootfs.is_dir() {
        return Err(SandboxError::Invalid("rootfs is not a directory".into()));
    }
    if req.run.contains('\0') {
        return Err(SandboxError::Invalid("run script contains NUL".into()));
    }
    Ok(())
}

/// Installs the SIGTERM and SIGINT handler. It does not set SA_RESTART. The wait loop polls,
/// so the handler only sets a flag.
fn install_signal_handlers() -> Result<(), SandboxError> {
    let action = SigAction::new(
        SigHandler::Handler(on_terminate),
        SaFlags::empty(),
        SigSet::empty(),
    );
    for sig in [Signal::SIGTERM, Signal::SIGINT] {
        // SAFETY: the handler only performs an atomic store, which is async-signal-safe.
        unsafe { sigaction(sig, &action) }
            .map_err(|e| SandboxError::Setup(format!("cannot install signal handler: {e}")))?;
    }
    Ok(())
}

fn prepare_dirs(layout: &Layout) -> Result<(), SandboxError> {
    for dir in [
        &layout.job_dir,
        &layout.bundle,
        &layout.state,
        &layout.etc,
        &layout.upper,
        &layout.work,
        &layout.merged,
    ] {
        fs::create_dir_all(dir)
            .map_err(SandboxError::io(format!("cannot create {}", dir.display())))?;
    }
    for dir in [&layout.job_dir, &layout.bundle, &layout.state] {
        fs::set_permissions(dir, fs::Permissions::from_mode(0o700))
            .map_err(SandboxError::io(format!("cannot chmod {}", dir.display())))?;
    }
    Ok(())
}

/// The files bind mounted at `/etc/hosts` and `/etc/hostname`.
fn write_etc_files(layout: &Layout) -> Result<(), SandboxError> {
    let hosts = layout.etc.join("hosts");
    fs::write(&hosts, HOSTS).map_err(SandboxError::io("cannot write hosts"))?;
    let hostname = layout.etc.join("hostname");
    fs::write(&hostname, format!("{HOSTNAME}\n"))
        .map_err(SandboxError::io("cannot write hostname"))?;
    for file in [&hosts, &hostname] {
        fs::set_permissions(file, fs::Permissions::from_mode(0o644))
            .map_err(SandboxError::io(format!("cannot chmod {}", file.display())))?;
    }
    Ok(())
}

/// Mounts the overlay at `merged` and returns it as the rootfs. If the overlay fails, the image
/// rootfs is copied instead, and a warning is printed.
fn mount_rootfs(req: &SandboxRequest, layout: &Layout) -> Result<PathBuf, SandboxError> {
    match mount_overlay(&req.rootfs, layout) {
        Ok(()) => Ok(layout.merged.clone()),
        Err(e) => {
            warn(format!(
                "overlay mount failed ({e}), copying the rootfs instead"
            ));
            copy_rootfs(&req.rootfs, &layout.merged)?;
            Ok(layout.merged.clone())
        }
    }
}

fn mount_overlay(lower: &Path, layout: &Layout) -> Result<(), SandboxError> {
    // Overlay options are comma separated, and the colon separates lower dirs, so paths
    // containing either can't be expressed.
    for p in [lower, &layout.upper, &layout.work] {
        let s = p.to_string_lossy();
        if s.contains(',') || s.contains(':') {
            return Err(SandboxError::Setup(format!(
                "overlay path has ',' or ':': {s}"
            )));
        }
    }
    let options = format!(
        "lowerdir={},upperdir={},workdir={}",
        lower.display(),
        layout.upper.display(),
        layout.work.display()
    );
    mount(
        Some("overlay"),
        layout.merged.as_path(),
        Some("overlay"),
        MsFlags::empty(),
        Some(options.as_str()),
    )
    .map_err(|e| SandboxError::Setup(e.to_string()))
}

/// A full copy of the rootfs, reflinked where the filesystem supports it.
fn copy_rootfs(src: &Path, dst: &Path) -> Result<(), SandboxError> {
    let cp = ["/usr/bin/cp", "/bin/cp"]
        .into_iter()
        .find(|p| Path::new(p).exists())
        .unwrap_or("cp");
    let status = Command::new(cp)
        .arg("-a")
        .arg("--reflink=auto")
        .arg(format!("{}/.", src.display()))
        .arg(format!("{}/", dst.display()))
        .status()
        .map_err(SandboxError::io("cannot run cp for the rootfs copy"))?;
    if status.success() {
        Ok(())
    } else {
        Err(SandboxError::Setup(format!("rootfs copy failed: {status}")))
    }
}

/// Runs the job with `[limits]` applied through its cgroup. If the container can't be built that
/// way, it is rebuilt in the same cgroup without the limit values, and a warning goes to the job
/// log.
///
/// libcontainer 0.7 cannot run a container in a user namespace without a cgroup manager, and it
/// only avoids systemd when the cgroup path is absolute. So "run without cgroups" is not
/// available. If the job cgroup can't be created at all, the job fails with exit 125.
fn run_job(req: &SandboxRequest, layout: &Layout, rootfs: &Path) -> Result<i32, SandboxError> {
    let cgroup = job_cgroup_path(&req.job_id, std::process::id())?;
    let mut spec = build_spec(req, rootfs)?;
    set_cgroup_path(&mut spec, &cgroup);

    match launch_and_wait(layout, &req.job_id, &spec) {
        Ok(code) => Ok(code),
        Err(SandboxError::Setup(reason)) if reason.starts_with("container: ") => {
            warn(format!("resource limits unavailable: {reason}"));
            strip_limits(&mut spec);
            let retry_id = format!("{}-nolimits", req.job_id);
            launch_and_wait(layout, &retry_id, &spec)
        }
        Err(e) => Err(e),
    }
}

/// The job's cgroup path, absolute and relative to the cgroup2 mount.
///
/// It must not be the runner's own cgroup. libcontainer kills every process in the job cgroup on
/// delete, so that would kill the runner. A cgroup that holds processes also can't enable
/// controllers for its children. The job cgroup goes in the parent slice instead, when the parent
/// is a plain domain cgroup. The runner's user owns both, because systemd delegates the user
/// slice.
///
/// The name is `forge-job-<job id>-<helper pid>`, so the runner can compute it after it has
/// spawned the helper. It uses the runner's own cgroup, which is also the helper's, to find the
/// parent slice. The path is relative to the cgroup2 mount, so the directory is
/// `/sys/fs/cgroup` joined with it.
pub fn job_cgroup_path(job_id: &str, helper_pid: u32) -> Result<PathBuf, SandboxError> {
    let text = fs::read_to_string("/proc/self/cgroup")
        .map_err(SandboxError::io("cannot read /proc/self/cgroup"))?;
    let own = text
        .lines()
        .find_map(|line| line.strip_prefix("0::"))
        .ok_or_else(|| SandboxError::Setup("cgroup v2 is not available".into()))?;
    let own = PathBuf::from(own.trim());
    let base = match own.parent() {
        Some(parent) if is_plain_domain(parent) => parent.to_path_buf(),
        _ => own,
    };
    // The helper's pid keeps the name unique even if two helpers share a job id.
    Ok(base.join(format!("forge-job-{job_id}-{helper_pid}")))
}

/// Whether `cgroup` (absolute, inside the cgroup2 mount) is a normal domain cgroup. A threaded
/// one can't hold domain children.
fn is_plain_domain(cgroup: &Path) -> bool {
    let file = Path::new("/sys/fs/cgroup")
        .join(cgroup.strip_prefix("/").unwrap_or(cgroup))
        .join("cgroup.type");
    fs::read_to_string(file).is_ok_and(|t| t.trim() == "domain")
}

fn set_cgroup_path(spec: &mut Spec, cgroup: &Path) {
    if let Some(linux) = spec.linux_mut() {
        linux.set_cgroups_path(Some(cgroup.to_owned()));
    }
}

/// Removes the resource values. The cgroup path stays, so libcontainer still uses the fs driver.
fn strip_limits(spec: &mut Spec) {
    if let Some(linux) = spec.linux_mut() {
        linux.set_resources(None);
    }
}

/// Writes the bundle, builds, starts, and waits. Build and start errors are returned as
/// `Setup("container: …")` so the caller can tell them apart from setup errors.
fn launch_and_wait(layout: &Layout, id: &str, spec: &Spec) -> Result<i32, SandboxError> {
    // Start each attempt with an empty state dir, so a failed attempt leaves no state behind.
    // The dir may not exist yet on the first attempt, so a failure here is not an error.
    let _ = fs::remove_dir_all(&layout.state);
    fs::create_dir_all(&layout.state).map_err(SandboxError::io("cannot create state dir"))?;
    fs::set_permissions(&layout.state, fs::Permissions::from_mode(0o700))
        .map_err(SandboxError::io("cannot chmod state dir"))?;

    write_bundle(layout, spec)?;
    let built = build_container(layout, id);
    // libcontainer has read the bundle config during build, and the env it holds is now in
    // state. Delete the bundle copy now.
    let _ = fs::remove_file(layout.bundle.join("config.json"));
    let mut container = match built {
        Ok(c) => c,
        Err(e) => {
            return Err(SandboxError::Setup(format!(
                "container: cannot create: {e}"
            )));
        }
    };
    let pid = match start_container(&mut container) {
        Ok(pid) => pid,
        Err(e) => {
            let _ = container.delete(true);
            return Err(e);
        }
    };
    wait_for_init(&mut container, pid)
}

/// Writes `bundle/config.json` (0600). It holds the job env and secrets, so it is removed right
/// after `build()`.
fn write_bundle(layout: &Layout, spec: &Spec) -> Result<(), SandboxError> {
    let json = serde_json::to_vec(spec)
        .map_err(|e| SandboxError::Setup(format!("cannot serialize spec: {e}")))?;
    let path = layout.bundle.join("config.json");
    let mut file = OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .mode(0o600)
        .open(&path)
        .map_err(SandboxError::io(format!(
            "cannot create {}",
            path.display()
        )))?;
    file.set_permissions(fs::Permissions::from_mode(0o600))
        .map_err(SandboxError::io("cannot chmod bundle config"))?;
    file.write_all(&json)
        .map_err(SandboxError::io("cannot write bundle config"))
}

/// Creates the container. Its stdin is `/dev/null`. Its stdout and stderr are dups of the
/// helper's, so the runner receives them.
fn build_container(layout: &Layout, id: &str) -> Result<Container, String> {
    let io_err = |e: io::Error| e.to_string();
    let devnull = fs::File::open("/dev/null").map_err(io_err)?;
    let stdout = io::stdout().as_fd().try_clone_to_owned().map_err(io_err)?;
    let stderr = io::stderr().as_fd().try_clone_to_owned().map_err(io_err)?;
    // The seccomp BPF is compiled here, before libcontainer forks. The executor only installs it.
    let filter = build_filter().map_err(|e| e.to_string())?;
    let builder = ContainerBuilder::new(id.to_owned(), SyscallType::default())
        .validate_id()
        .and_then(|b| b.with_root_path(&layout.state))
        .map_err(|e| e.to_string())?;
    builder
        .with_stdin(devnull)
        .with_stdout(stdout)
        .with_stderr(stderr)
        .with_executor(FilteredExecutor::new(filter))
        .as_init(&layout.bundle)
        .with_systemd(false)
        .build()
        .map_err(|e| e.to_string())
}

/// Starts the container and returns the init pid. The init is a child of this process.
fn start_container(container: &mut Container) -> Result<Pid, SandboxError> {
    container
        .start()
        .map_err(|e| SandboxError::Setup(format!("container: cannot start: {e}")))?;
    // libcontainer has its own nix, so its Pid is a different type. The raw value is the same.
    container
        .pid()
        .map(|pid| Pid::from_raw(pid.as_raw()))
        .ok_or_else(|| SandboxError::Setup("container: no init pid".into()))
}

/// SIGKILL, as libcontainer's own signal type. Its `Signal` is not nix's, so it is built from
/// the number.
const SIGKILL_NUMBER: i32 = 9;

fn kill_container(container: &mut Container) -> Result<(), String> {
    let sigkill = libcontainer::signal::Signal::try_from(SIGKILL_NUMBER)
        .map_err(|_| "cannot build SIGKILL".to_owned())?;
    container.kill(sigkill, true).map_err(|e| e.to_string())
}

/// Waits for the init process. On SIGTERM the container is killed. The child is then reaped,
/// and `delete` removes the container's state. Returns the job's exit code.
fn wait_for_init(container: &mut Container, pid: Pid) -> Result<i32, SandboxError> {
    let mut killed = false;
    loop {
        if !killed && TERMINATE.load(Ordering::SeqCst) {
            killed = true;
            if let Err(e) = kill_container(container) {
                warn(format!("cannot kill container: {e}"));
            }
        }
        match waitpid(pid, Some(WaitPidFlag::WNOHANG)) {
            Ok(WaitStatus::StillAlive) => std::thread::sleep(POLL),
            Ok(WaitStatus::Exited(_, code)) => {
                delete_container(container);
                return Ok(code);
            }
            Ok(WaitStatus::Signaled(_, sig, _)) => {
                delete_container(container);
                return Ok(128 + sig as i32);
            }
            Ok(_) => {}
            Err(Errno::EINTR) => {}
            Err(e) => {
                let _ = kill_container(container);
                delete_container(container);
                return Err(SandboxError::Setup(format!("waiting for the job: {e}")));
            }
        }
    }
}

fn delete_container(container: &mut Container) {
    if let Err(e) = container.delete(true) {
        warn(format!("cannot delete container: {e}"));
    }
}

/// Removes the whole job dir: state and bundle (which hold the env), and the overlay upper dir
/// (which may hold files the job wrote). The container is gone by now. The overlay is still
/// mounted in this process's mount namespace, so detach it first. The helper is still root in its
/// user namespace, so it can remove the overlay work files that the runner's uid could not.
fn cleanup(layout: &Layout) {
    let _ = umount2(layout.merged.as_path(), MntFlags::MNT_DETACH);
    if let Err(e) = fs::remove_dir_all(&layout.job_dir) {
        warn(format!("cannot remove {}: {e}", layout.job_dir.display()));
    }
}

fn warn(message: impl std::fmt::Display) {
    eprintln!("forge-sandbox: warning: {message}");
}