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

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

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

mod cgroup;

macro_rules! fail {
    ($($arg:tt)*) => {{
        eprintln!($($arg)*);
        std::process::exit(1)}
    };
}

fn parse_string(original: &str, msg: &str, num: usize) -> u64 {
    let (to_parse, scale, unit): (_, u64, &'static str) = match original.as_bytes() {
        [rest @ .., b'K', b'B'] => (rest, 1_000, "KB"),
        [rest @ .., b'M', b'B'] => (rest, 1_000_000, "MB"),
        [rest @ .., b'G', b'B'] => (rest, 1_000_000_000, "GB"),
        [rest @ .., b'T', b'B'] => (rest, 1_000_000_000_000, "TB"),
        [rest @ .., b'P', b'B'] => (rest, 1_000_000_000_000_000, "PB"),
        [rest @ .., b'E', b'B'] => (rest, 1_000_000_000_000_000_000, "EB"),
        [rest @ .., b'K', b'i', b'B'] => (rest, 1 << 10, "KiB"),
        [rest @ .., b'M', b'i', b'B'] => (rest, 1 << 20, "MiB"),
        [rest @ .., b'G', b'i', b'B'] => (rest, 1 << 30, "GiB"),
        [rest @ .., b'T', b'i', b'B'] => (rest, 1 << 40, "TiB"),
        [rest @ .., b'P', b'i', b'B'] => (rest, 1 << 50, "PiB"),
        [rest @ .., b'E', b'i', b'B'] => (rest, 1 << 60, "EiB"),
        [rest @ .., b'B'] if rest.last().map(u8::is_ascii_digit).unwrap_or(false) => (rest, 1, "B"),
        _ => fail!("{msg}: {original:?} is missing a unit suffix"),
    };
    let (to_parse, radix) = match to_parse {
        [b'0', b'x' | b'X', s @ ..] => (s, 16),
        _ => (to_parse, 10),
    };
    // SAFETY: lopping ASCII bytes off start
    // and end of UTF-8 string leaves valid UTF_8
    let result: u64 =
        match u64::from_str_radix(unsafe { str::from_utf8_unchecked(to_parse) }, radix) {
            Ok(e) => e,
            Err(e) => {
                fail!("{msg}: argument {num}: Cannot parse as u64: {e}")
            }
        };
    if to_parse[0] == b'0' || to_parse[0] == b'+' {
        fail!(
            "{msg}: argument {num}: Leading {} in {original} not allowed",
            to_parse[0] as char
        )
    }
    result.checked_mul(scale).unwrap_or_else(|| {
        fail!(
            "{msg}: {original} is too large for unit {unit} ({original} × {scale} = {} > {})",
            u128::from(result) * u128::from(scale),
            u64::MAX,
        )
    })
}

fn main() {
    let mut args = std::env::args_os();
    if args.next().is_none() {
        fail!("No command line arguments (argv[0] is NULL)");
    }
    let mut args = args.enumerate().map(|(num, arg)| {
        match str::from_utf8(<OsStr as OsStrExt>::as_bytes(&arg)) {
            Ok(v) => (num, v.to_owned()),
            Err(e) => {
                fail!(
                    "Argument {num} bytes {} through {} are invalid UTF-8",
                    e.valid_up_to(),
                    e.error_len().unwrap_or_else(|| arg.len())
                )
            }
        }
    });
    let mut finished = false;
    let mut positional_arguments;
    let mut values = [
        ("pids.max", None),
        ("memory.high", None),
        ("memory.max", None),
    ];

    let last = loop {
        let Some((num, arg_str)) = args.next() else {
            finished = true;
            break None;
        };
        if arg_str.as_bytes().get(0) != Some(&b'-') {
            break Some(arg_str.to_owned());
        }
        let Some(suffix) = arg_str.strip_prefix("--") else {
            fail!("Short options (with a single '-', as in {arg_str:?}) are not supported")
        };

        if suffix.is_empty() {
            break None;
        }

        for value in &mut values {
            if let Some(suffix) = suffix.strip_prefix(value.0)
                && let Some(suffix) = suffix.strip_prefix("=")
            {
                value.1 = Some(parse_string(suffix, &value.0, num));
            }
        }

        fail!("Unknown option {arg_str:?}")
    };

    if finished {
        positional_arguments = vec![]
    } else {
        positional_arguments = last.into_iter().chain(args.map(|a| a.1)).collect()
    }

    if positional_arguments.is_empty() {
        fail!(
            "Have {} arguments, expect at least 1",
            positional_arguments.len()
        )
    }

    // slow but we do not care
    let mut child_path = positional_arguments.remove(0);
    if child_path.len() > 247 {
        fail!("Cgroup name {child_path:?} too long");
    }
    if child_path.as_bytes().contains(&b'/') {
        fail!("Cgroup name {child_path:?} contains /");
    }
    child_path += ".service";
    let child_path = Path::new(&child_path);

    let local_cgroup =
        std::fs::read("/proc/thread-self/cgroup").expect("Should be able to read from procfs");
    if !local_cgroup.starts_with(b"0::/")
        || !local_cgroup.ends_with(b"\n")
        || local_cgroup.contains(&b'\0')
    {
        fail!("Kernel bug: Bad contents of /proc/thread-self/cgroup");
    }

    let mut local_cgroup = PathBuf::from(<OsString as OsStringExt>::from_vec(
        local_cgroup[4..local_cgroup.len() - 1].to_owned(),
    ));

    match local_cgroup.components().next_back() {
        Some(Component::Prefix(_)) => unreachable!("does not occur on Linux"),
        Some(Component::RootDir) | None => {}
        Some(Component::CurDir) => unreachable!("not produced by kernel"),
        Some(Component::ParentDir) => unreachable!("not produced by kernel"),
        Some(Component::Normal(os_str)) => {
            if os_str.as_bytes() == b"@inner.service" {
                let _ = local_cgroup.pop();
            }
        }
    }

    if local_cgroup.as_os_str().is_empty() {
        local_cgroup = PathBuf::from(".");
    }

    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,
    )
    .expect("Cannot get cgroup root");

    let cgroup = match cgroup::LeafCgroup::open_cgroup(cgroup_root.as_fd(), &local_cgroup) {
        Ok(e) => e,
        Err(e) => fail!(
            "Bad cgroup operation {e:?} opening {}",
            local_cgroup.display()
        ),
    };

    rustix::fs::flock(cgroup.as_fd(), FlockOperation::LockExclusive).expect("oops");
    purge_cgroup(&cgroup, child_path);
    if positional_arguments.is_empty() {
        return;
    }
    if let Err(e) = cgroup.make_child(child_path) {
        fail!("Cannot create child cgroup: {e}")
    }
    let child = cgroup
        .open_cgroup_at(Path::new(&child_path))
        .expect("cannot open child cgroup");
    child
        .make_child(Path::new("@inner.service"))
        .expect("cannot make child cgroup");
    let child = child
        .open_cgroup_at(Path::new("@inner.service"))
        .expect("cannot open child cgroup");

    for (name, value) in values {
        if let Some(value) = value {
            write_cgroup_value(&child, name, &value.to_string());
        }
    }

    // SAFETY: safe FFI call
    let pid = unsafe { libc::getpid() };
    write_cgroup_value(&child, "cgroup.procs", &pid.to_string());

    let program = positional_arguments.remove(0);
    let e = std::process::Command::new(program.clone())
        .args(positional_arguments)
        .exec();
    fail!("Cannot spawn child {:?}: {}", program, e);
}

fn write_cgroup_value(child: &cgroup::LeafCgroup, name: &str, value: &str) {
    let path = Path::new(name);
    let fd = rustix::fs::openat2(
        child.as_fd(),
        path,
        OFlags::NOATIME | OFlags::CLOEXEC | OFlags::NOFOLLOW | OFlags::WRONLY,
        Mode::empty(),
        ResolveFlags::NO_SYMLINKS | ResolveFlags::BENEATH | ResolveFlags::NO_XDEV,
    )
    .unwrap_or_else(|e| fail!("Cannot open {:?}: {}", path, e));
    let () = File::from(fd)
        .write_all(value.as_bytes())
        .unwrap_or_else(|e| fail!("Cannot write {:?} to {:?}: {}", name, path, e));
}

fn purge_cgroup(cgroup: &cgroup::LeafCgroup, child_path: &Path) {
    let child = match cgroup.open_cgroup_at(child_path) {
        Ok(child_cgroup) => child_cgroup,
        Err(Errno::NOENT) => return,
        Err(other) => fail!("Cannot open child cgroup {child_path:?}: {other}"),
    };
    child.kill_processes().expect("cannot kill children");
    cgroup
        .delete_child(child_path)
        .expect("cannot delete child");
}