// SPDX-License-Identifier: EUPL-1.2+
// SPDX-FileCopyrightText: 2026 Demi Marie Obenour <demiobenour@gmail.com>

mod cgroup;

use cgroup::{Cgroup, openat2_simple, write_value};
use rustix::{
    fs::{FlockOperation, Mode, OFlags, XattrFlags},
    io::Errno,
};
use std::{
    env::ArgsOs,
    ffi::OsStr,
    fs::File,
    io::Read as _,
    os::unix::prelude::*,
    path::{Path, PathBuf},
};

// Check that the path is canonical,
// then split it into basename and filename.
fn split_path(path: &Path) -> Result<(&Path, &Path), String> {
    cgroup::check_path(path)
        .map(|()| (path.parent().unwrap(), Path::new(path.file_name().unwrap())))
}

fn read_control_file(fd: &dyn AsFd, p: &Path) -> Result<Vec<u8>, String> {
    let mut buf = Vec::new();
    File::from(
        openat2_simple(&fd, Path::new(p), OFlags::RDONLY)
            .map_err(|e| format!("Cannot open {p:?}: {e}"))?,
    )
    .read_to_end(&mut buf)
    .map_err(|e| format!("Cannot read {p:?}: {e}"))?;
    Ok(buf)
}

fn enable_subtree_control(fd: &dyn AsFd) -> Result<(), String> {
    let p = Path::new("cgroup.controllers");
    let buf = read_control_file(fd, p)?;
    let mut subtree = vec![];
    for controller in buf.split(|&b| b == b' ').filter(|e| !e.is_empty()) {
        if !subtree.is_empty() {
            subtree.push(b' ');
        }
        subtree.push(b'+');
        subtree.extend_from_slice(controller);
    }
    if !subtree.is_empty() {
        write_value(&fd, Path::new("cgroup.subtree_control"), &subtree)?;
    }
    Ok(())
}

fn cgroup_setup(mut args: ArgsOs) -> Result<(), String> {
    let mut leaf = false;
    let mut cgroup_path;
    let mut systemd_compat = false;
    let mut wait = true;
    loop {
        cgroup_path = args.next();
        let Some(ref arg_) = cgroup_path else {
            break;
        };
        let arg_ = arg_.as_bytes();
        if arg_ == b"--" {
            cgroup_path = args.next();
            break;
        }
        if !arg_.starts_with(b"-") {
            break;
        }

        if !arg_.starts_with(b"--") {
            return Err("takes no short options".to_owned());
        }

        match &arg_[2..] {
            b"leaf" => leaf = true,
            b"wait" => wait = true,
            b"no-wait" => wait = false,
            b"systemd-compat" => systemd_compat = true,
            arg => return Err(format!("unknown long option {:?}", OsStr::from_bytes(arg))),
        }
    }
    let Some(cgroup_path) = cgroup_path.map(PathBuf::from) else {
        return Err("have no positional arguments, expected at least 1".to_owned());
    };

    let (parent_cgroup_path, child_cgroup_path) = split_path(&cgroup_path)?;
    let cgroup = Cgroup::new(parent_cgroup_path)?;
    match rustix::fs::mkdirat(&cgroup, child_cgroup_path, Mode::from_raw_mode(0o755)) {
        Ok(()) | Err(Errno::EXIST) => {}
        Err(e) => return Err(format!("Cannot make child cgroup: {e}")),
    }
    let child = cgroup
        .open_beneath(child_cgroup_path, OFlags::RDONLY | OFlags::DIRECTORY)
        .map_err(|e| format!("Cannot make child cgroup: {e}"))?;
    if wait {
        // While waiting, only hold an exclusive lock on the child, not the parent.
        rustix::fs::flock(&child, FlockOperation::LockExclusive)
            .map_err(|e| format!("Cannot take an exclusive lock on child cgroup: {e}"))?;
        rustix::fs::flock(&cgroup, FlockOperation::LockShared)
            .map_err(|e| format!("Cannot downgrade lock on cgroup to a shared lock: {e}"))?;
        Cgroup::wait_for_empty(&child)
            .map_err(|e| format!("Cannot wait for {parent_cgroup_path:?} to be empty: {e}"))?;
    }
    let pid = std::process::id().to_string();
    if leaf {
        if args.len() != 0 {
            // If we aren't delegating any cgroups, don't create a sub-cgroup.
            write_value(&child, Path::new("cgroup.procs"), pid.as_bytes())
                .map_err(|e| format!("Cannot move process to child cgroup: {e}"))?;
        }
    } else {
        // If the child process will need to manage cgroups itself, it will need
        // to set up a sub-cgroup due to the "no internal processes" rule.  It's
        // simplest to just do it automatically.  If the cgroup already exists,
        // that isn't an error.
        match rustix::fs::mkdirat(&child, cgroup::DEFAULT_LEAF, Mode::from_raw_mode(0o755)) {
            Ok(()) | Err(Errno::EXIST) => {}
            Err(e) => return Err(format!("Cannot make child cgroup: {e}")),
        }
        if args.len() != 0 {
            let child_proc_path = Path::new(cgroup::DEFAULT_LEAF).join(Path::new("cgroup.procs"));
            write_value(&child, &child_proc_path, pid.as_bytes())
                .map_err(|e| format!("Cannot move process to child cgroup: {e}"))?;
        }
        if systemd_compat {
            // systemd-aware programs expect to have user.delegate=1
            // and to set cgroup.subtree_control themselves
            rustix::fs::fsetxattr(&child, c"user.delegate", b"1", XattrFlags::empty()).map_err(
                |e| format!("Cannot enable cgroup delegation in {parent_cgroup_path:?}: {e}"),
            )?
        } else {
            // Spectrum's programs do not check for user.delegate=1
            // and expect the caller to set cgroup.subtree_control.
            enable_subtree_control(&child)?;
        }
    }
    let Some(program_name) = args.next() else {
        return Ok(());
    };
    let e = std::process::Command::new(&program_name).args(args).exec();
    Err(format!("Cannot spawn child {:?}: {}", program_name, e))
}

fn cgroup_purge(mut args: ArgsOs) -> Result<(), String> {
    if args.len() != 1 {
        return Err("usage: cgroup-purge CGROUP_TO_PURGE".to_owned());
    }
    let arg = args.next().unwrap();
    let (parent, child) = split_path(Path::new(&arg))?;
    Cgroup::new(parent)?.purge_child(child)
}

fn run(prog_name: &OsStr, args: ArgsOs) -> Result<(), String> {
    match prog_name
        .as_bytes()
        .split(|&b| b == b'/')
        .next_back()
        .unwrap()
    {
        b"cgroup-setup" => cgroup_setup(args),
        b"cgroup-purge" => cgroup_purge(args),
        _ => Err(format!(
            "must be invoked as \"cgroup-setup\" or \
                 \"cgroup-purge\", got {prog_name:?}",
        )),
    }
}

fn main() {
    let mut args = std::env::args_os();
    let Some(prog_name) = args.next() else {
        eprintln!("No command line arguments (argv[0] is NULL)");
        std::process::exit(1);
    };
    match run(&prog_name, args) {
        Ok(()) => {}
        Err(e) => {
            eprintln!("{prog_name:?}: {}", e);
            std::process::exit(1);
        }
    }
}