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

use std::{
    ffi::{OsStr, OsString},
    fs::File,
    io::Read as _,
    os::unix::prelude::*,
    path::{Component, Path, PathBuf},
};

use rustix::{
    fs::{FlockOperation, Mode, OFlags, ResolveFlags},
    io::Errno,
};

use crate::cgroup::Access;

mod cgroup;

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 main_(&prog_name, args) {
        Ok(()) => {}
        Err(e) => {
            eprintln!("{prog_name:?}: {}", e);
            std::process::exit(1);
        }
    }
}

fn main_(prog_name: &OsStr, mut args: std::env::ArgsOs) -> Result<(), String> {
    let mut purge = false;
    match prog_name
        .as_bytes()
        .split(|&b| b == b'/')
        .next_back()
        .unwrap()
    {
        b"finish" => {
            return s6_finish(&mut args);
        }
        b"cgroup-setup" => {}
        e => {
            return Err(format!(
                "must be invoked as \"cgroup-setup\" or \"finish\", got {:?}",
                e
            ));
        }
    };
    let mut leaf = false;
    let mut cgroup_relative_path;
    let mut delegate = false;
    let mut init_subtree = false;
    let mut child_name: Option<&'static OsStr> = None;
    let mut wait = None;
    loop {
        cgroup_relative_path = args.next();
        let Some(ref arg_) = cgroup_relative_path else {
            break;
        };
        let arg_ = arg_.as_bytes();
        if arg_ == b"--" {
            cgroup_relative_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"purge" => purge = true,
            b"leaf" => leaf = true,
            b"delegate" => delegate = true,
            b"init-subtree" => init_subtree = true,
            b"wait" => wait = Some(true),
            b"no-wait" => wait = Some(false),
            b"child-name" if child_name.is_none() => match args.next() {
                Some(arg) => child_name = Some(arg.leak()),
                None => return Err("--child-name: missing argument".to_owned()),
            },
            b"child-name" => return Err("--child-name: cannot be used twice".to_owned()),
            arg => match str::from_utf8(arg) {
                Ok(e) => return Err(format!("unknown long option {e:?}")),
                Err(_) => return Err("long option isn't UTF-8".to_owned()),
            },
        }
    }

    let default_child_name = OsStr::from_bytes(b"$inner.service");

    let child_name = Path::new(child_name.unwrap_or(default_child_name));

    let Some(cgroup_relative_path) = cgroup_relative_path else {
        return Err("have no positional arguments, expected at least 1".to_owned());
    };

    // The kernel doesn't care, but displaying messages does.
    let cgroup_path = String::try_from(cgroup_relative_path.into_vec())
        .map_err(|e| format!("non-UTF-8 cgroup path not supported (error is {e})"))?;

    if cgroup_path.is_empty() {
        return Err("cgroup name is empty".to_owned());
    }
    if cgroup_path == ".." {
        return Err("cgroup name is ..".to_owned());
    }
    // If we aren't asked to create a child process, don't wait for existing
    // processes to die unless explicitly asked to.  Waiting for a cgroup
    // we are in to be empty is a guaranteed deadlock.
    let may_wait = cgroup_path != "." && cgroup_path != "/";
    let wait = match wait {
        Some(false) => false,
        None => may_wait && args.len() > 0,
        Some(true) if !may_wait => {
            let msg = "Cannot wait for the program's own cgroup or \
                       root cgroup to be empty";
            return Err(msg.to_owned());
        }
        Some(true) => true,
    };

    let params = CgroupParams {
        purge,
        leaf,
        delegate,
        init_subtree,
        wait,
    };

    let (full_path, cgroup_target, cgroup) = cgroup_parse(cgroup_path)?;
    if params.purge {
        purge_cgroup(&cgroup, &cgroup_target)?;
    }
    let child = match cgroup.open_cgroup_at(&cgroup_target) {
        Ok(child_cgroup) => child_cgroup,
        Err(Errno::NOENT) => {
            if let Err(e) = cgroup.make_child(&cgroup_target) {
                return Err(format!("Cannot create child cgroup {full_path:?}: {e}"));
            }
            cgroup
                .open_cgroup_at(&cgroup_target)
                .map_err(|e| format!("Cannot open child cgroup {full_path:?}: {e}"))?
        }
        Err(other) => {
            return Err(format!("Cannot open child cgroup {full_path:?}: {other}"));
        }
    };
    if params.wait {
        child
            .wait_for_empty(false)
            .map_err(|e| format!("Cannot wait for {full_path:?} to be empty: {e}"))?;
    }
    let pid = std::process::id().to_string();
    if params.leaf {
        // If we aren't delegating any cgroups, don't create a sub-cgroup.
        child
            .write_cgroup_value("cgroup.procs", &pid)
            .map_err(|e| format!("Cannot write to {full_path:?}/cgroup.procs: {e}"))?;
    } else {
        child.make_child(child_name).map_err(|e| {
            format!(
                "Cannot create child cgroup {}/{}: {e}",
                full_path.display(),
                child_name.display()
            )
        })?;

        // 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.
        let grandchild = child.open_cgroup_at(Path::new(child_name)).map_err(|e| {
            format!(
                "Cannot open child cgroup {}/{}: {e}",
                full_path.display(),
                child_name.display()
            )
        })?;
        grandchild
            .write_cgroup_value("cgroup.procs", &pid)
            .map_err(|e| {
                format!(
                    "Cannot write to {}/{}/cgroup.procs: {e}",
                    full_path.display(),
                    child_name.display()
                )
            })?;
    }
    if params.init_subtree {
        enable_subtree_control(&cgroup)?;
    }
    if !params.leaf {
        enable_subtree_control(&child)?;
    }
    if params.delegate {
        child
            .enable_delegation()
            .map_err(|e| format!("Cannot enable cgroup delegation in {full_path:?}: {e}"))?;
    }
    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 s6_finish(args: &mut std::env::ArgsOs) -> Result<(), String> {
    if args.len() < 3 {
        return Err(format!(
            "s6 finish scripts take 3 arguments, got {}",
            args.len()
        ));
    }
    let status = parse_digit_string(&args.next().unwrap(), "exit status")?;
    let signal = args.next().unwrap();
    let signal = if status == 256 {
        Some(parse_digit_string(&signal, "signal number")?)
    } else {
        None
    };
    let service = args
        .next()
        .unwrap()
        .into_string()
        .map_err(|e| format!("Service name {e:?} is not UTF-8"))?;
    let (_full_path, cgroup_target, cgroup) = cgroup_parse(service)?;
    let r = purge_cgroup(&cgroup, &cgroup_target);
    if let Some(signal) = signal {
        match signal as libc::c_int {
            libc::SIGBUS
            | libc::SIGFPE
            | libc::SIGABRT
            | libc::SIGTRAP
            | libc::SIGSEGV
            | libc::SIGILL => {
                // Process *crashed*, indicating a *possible exploit attempt*.
                // s6 should *not* restart it.  This is distinct from a Rust panic,
                // which is much less likely to indicate memory corruption.
                if let Err(e) = r {
                    // do not panic on stderr write failure
                    eprintln!("Could not purge cgroup: {e}");
                }
                std::process::exit(125)
            }
            _ => return r,
        }
    }
    r
}

fn parse_digit_string(digits: &OsStr, msg: &str) -> Result<u16, String> {
    let checked = match str::from_utf8(digits.as_bytes()) {
        Ok(s) => s,
        Err(e) => return Err(format!("{msg} is not UTF-8: {e}")),
    };
    let r = checked
        .parse::<u16>()
        .map_err(|e| format!("{msg} {digits:?} is a bad 16-bit number: {e}"))?;
    match checked.as_bytes() {
        b"0" | [b'1'..=b'9', ..] => Ok(r),
        [b'0', ..] => Err(format!("{msg} {} has a leading 0", digits.display())),
        _ => Err(format!("{msg} {} starts with +", digits.display())),
    }
}

fn cgroup_parse(mut arg: String) -> Result<(PathBuf, PathBuf, cgroup::LeafCgroup), String> {
    let cgroup_root = rustix::fs::openat2(
        rustix::fs::CWD,
        Path::new("/sys/fs/cgroup"),
        OFlags::DIRECTORY | OFlags::RDONLY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
        Mode::empty(),
        ResolveFlags::NO_MAGICLINKS | ResolveFlags::NO_SYMLINKS,
    )
    .map_err(|e| format!("cannot open /sys/fs/cgroup: {e}"))?;
    let cgroup_absolute = arg.starts_with("/");
    let (full_path, mut cgroup_parent, cgroup_target) = if cgroup_absolute {
        if arg.contains("//") {
            return Err(format!("Cgroup path {arg:?} contains //"));
        }
        arg.remove(0);
        let to_create = PathBuf::from(arg);
        for component in to_create.components() {
            match component {
                Component::Prefix(_) => unreachable!("not present on Unix"),
                Component::CurDir | Component::ParentDir => {
                    return Err("Cgroup path has . or .. components".to_owned());
                }
                Component::RootDir | Component::Normal(_) => {}
            }
        }
        let Some(last_component) = to_create.file_name() else {
            return Err(format!("Cgroup path {to_create:?} has no file name"));
        };
        if to_create.parent().is_none() {
            return Err(format!("Cgroup path {to_create:?} has no parent name"));
        }
        let r = last_component.as_bytes().to_owned();
        let mut prefix = to_create.clone();
        prefix.pop();
        (to_create, prefix, r)
    } else {
        if arg.is_empty() {
            return Err("Cgroup path is empty".to_string());
        }

        let mut local_cgroup: Vec<u8> = std::fs::read("/proc/thread-self/cgroup")
            .map_err(|e| format!("cannot read /proc/thread-self/cgroup: {e}"))?;
        let local_cgroup_len = local_cgroup.len();
        if local_cgroup_len < 5
            || local_cgroup[..4] != *b"0::/"
            || local_cgroup[local_cgroup_len - 1] != b'\n'
        {
            return Err(format!(
                "Invalid contents {local_cgroup:?} of /proc/thread-self/cgroup - \
                 do you have cgroups v1 mounted instead of cgroups v2?"
            ));
        }

        local_cgroup.copy_within(4..local_cgroup_len - 1, 0);
        local_cgroup.truncate(local_cgroup_len - 5);
        let total_path = PathBuf::from(<OsString as OsStringExt>::from_vec(local_cgroup));
        let mut r = total_path.clone();
        r.push(&arg);
        (r, total_path, arg.into())
    };
    if cgroup_parent.file_name() == Some(OsStr::from_bytes(b"$inner.service")) {
        cgroup_parent.pop();
    }
    if cgroup_parent.as_os_str().is_empty() {
        cgroup_parent = ".".into();
    }
    let cgroup_target = OsString::from_vec(cgroup_target).into();
    let cgroup = cgroup::LeafCgroup::open_cgroup(cgroup_root.as_fd(), &cgroup_parent)
        .map_err(|e| format!("Failed to open {}: {e:?}", cgroup_parent.display()))?;
    match rustix::fs::flock(cgroup.as_fd(), FlockOperation::LockExclusive) {
        Ok(()) => {}
        Err(e) => return Err(format!("Cannot lock cgroup: {e}")),
    }
    Ok((full_path, cgroup_target, cgroup))
}

struct CgroupParams {
    purge: bool,
    leaf: bool,
    delegate: bool,
    init_subtree: bool,
    wait: bool,
}

fn enable_subtree_control(cgroup: &cgroup::LeafCgroup) -> Result<(), String> {
    let mut buf = Vec::new();
    File::from(
        cgroup
            .open_subtree(Path::new("cgroup.controllers"), Access::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![];
    if buf.ends_with(b"\n") {
        buf.pop();
    }
    for controller in buf.split(|&b| b == b' ').filter(|e| !e.is_empty()) {
        for &c in controller {
            if c <= b' ' || c >= 0x7F {
                return Err(format!("Bad byte {c} in cgroup.controllers"));
            }
        }
        if !subtree.is_empty() {
            subtree.push(b' ');
        }
        subtree.push(b'+');
        subtree.extend_from_slice(controller);
    }
    if !subtree.is_empty() {
        cgroup.write_cgroup_value("cgroup.subtree_control", str::from_utf8(&subtree).unwrap())?;
    }
    Ok(())
}

fn purge_cgroup(cgroup: &cgroup::LeafCgroup, cgroup_target: &Path) -> Result<(), String> {
    let child = match cgroup.open_cgroup_at(cgroup_target) {
        Ok(child_cgroup) => child_cgroup,
        Err(Errno::NOENT) => return Ok(()),
        Err(other) => {
            return Err(format!(
                "Cannot open child cgroup {cgroup_target:?}: {other}"
            ));
        }
    };
    child
        .wait_for_empty(true)
        .map_err(|e| format!("Cannot kill programs in {cgroup_target:?}: {e}"))?;

    cgroup
        .delete_child(cgroup_target)
        .map_err(|e| format!("Delete child cgroup {cgroup_target:?}: {e}"))
}