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

mod cgroup;

use cgroup::{Cgroup, OpenFlags, openat2_simple, write_value};
use rustix::{
    fs::{FlockOperation, Mode, XattrFlags},
    io::Errno,
};
use std::{
    env::ArgsOs,
    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)?;
    Ok((path.parent().unwrap(), Path::new(path.file_name().unwrap())))
}

fn enable_subtree_control(fd: &dyn AsFd) -> Result<(), String> {
    let mut buf = Vec::new();
    File::from(
        openat2_simple(fd, c"cgroup.controllers", OpenFlags::Read)
            .map_err(|e| format!("Cannot open cgroup.controllers: {e}"))?,
    )
    .read_to_end(&mut buf)
    .map_err(|e| format!("Cannot read cgroup.controllers: {e}"))?;
    let mut subtree = vec![];
    for controller in buf.split(|&b| b == b' ') {
        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(args: ArgsOs) -> Result<(), String> {
    let mut leaf = false;
    let mut systemd_delegate = false;
    let mut wait = true;
    let mut args = args.peekable();
    while let Some(arg) = args.peek() {
        if !arg.as_bytes().starts_with(b"-") {
            break;
        }
        let arg = args.next().unwrap();
        let Some(arg_) = arg.as_bytes().strip_prefix(b"--") else {
            return Err("takes no short options".to_owned());
        };
        match arg_ {
            b"" => break,
            b"leaf" => leaf = true,
            b"no-wait" => wait = false,
            b"systemd-delegate" => systemd_delegate = true,
            _ => return Err(format!("unknown long option {arg:?}")),
        }
    }
    let Some(cgroup_path) = args.next().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 {child_cgroup_path:?}: {e}"
            ));
        }
    }
    let child = openat2_simple(&cgroup, child_cgroup_path, OpenFlags::Directory)
        .map_err(|e| format!("Cannot make child cgroup: {e}"))?;
    // While waiting, hold an exclusive lock on the child.
    // This avoids two processes both waiting for the same cgroup to become
    // empty, then spawning processes in the same cgroup.
    rustix::fs::flock(&child, FlockOperation::LockExclusive)
        .map_err(|e| format!("Cannot take an exclusive lock on child cgroup: {e}"))?;
    if wait {
        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_delegate {
            // 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: &Path, args: ArgsOs) -> Result<(), String> {
    match prog_name.file_name().map(|f| f.as_bytes()) {
        Some(b"cgroup-setup") => cgroup_setup(args),
        Some(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(Path::new(&prog_name), args) {
        Ok(()) => {}
        Err(e) => {
            eprintln!("{prog_name:?}: {}", e);
            std::process::exit(1);
        }
    }
}