// 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 main() {
    let mut args = std::env::args_os();
    let Some(prog_name) = args.next() else {
        fail!("No command line arguments (argv[0] is NULL)");
    };
    let mut cgroup_relative_path = args.next();
    if cgroup_relative_path.as_deref() == Some(OsStr::from_bytes(b"--")) {
        cgroup_relative_path = args.next();
    }

    let Some(cgroup_relative_path) = cgroup_relative_path else {
        fail!("{prog_name:?}: Have no positional arguments, expect at least 1")
    };
    let mut cgroup_relative_path = cgroup_relative_path.into_vec();

    // slow but we do not care
    if cgroup_relative_path.len() > 247 {
        fail!("{prog_name:?}: Cgroup name {cgroup_relative_path:?} too long");
    }
    if cgroup_relative_path.is_empty() {
        fail!("{prog_name:?}: Cgroup name {cgroup_relative_path:?} empty");
    }
    if cgroup_relative_path[0] == b'-' {
        fail!("{prog_name:?}: Cgroup name {cgroup_relative_path:?} starts with '-'");
    }
    for &i in &cgroup_relative_path {
        match i {
            b'/' => fail!(
                "{prog_name:?}: Cgroup name {cgroup_relative_path:?} \
contains /"
            ),
            b'$' => fail!(
                "{prog_name:?}: Cgroup name {cgroup_relative_path:?} \
contains $: did you forget an execline substitution?"
            ),
            b'!'..=b'~' => {}
            _ => fail!(
                "{prog_name:?}: Cgroup name {cgroup_relative_path:?} \
contains space, non-ASCII character, or control character (byte {i})"
            ),
        }
    }
    cgroup_relative_path.extend_from_slice(b".service");
    let cgroup_relative_path = Path::new(OsStr::from_bytes(&cgroup_relative_path));

    let local_cgroup = std::fs::read("/proc/thread-self/cgroup")
        .unwrap_or_else(|e| fail!("{prog_name:?}: cannot read /proc/thread-self/cgroup: {e}"));

    if !local_cgroup.starts_with(b"0::/")
        || !local_cgroup.ends_with(b"\n")
        || local_cgroup.contains(&b'\0')
    {
        fail!("{prog_name:?}: 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,
    )
    .unwrap_or_else(|e| fail!("{prog_name:?}: cannot open /sys/fs/cgroup: {e}"));

    let cgroup = match cgroup::LeafCgroup::open_cgroup(cgroup_root.as_fd(), &local_cgroup) {
        Ok(e) => e,
        Err(e) => fail!(
            "{prog_name:?}: Failed to open {}: {e:?}",
            local_cgroup.display()
        ),
    };

    match rustix::fs::flock(cgroup.as_fd(), FlockOperation::LockExclusive) {
        Ok(()) => {}
        Err(e) => fail!("{prog_name:?}: cannot lock cgroup: {e}"),
    }
    purge_cgroup(&prog_name, &cgroup, cgroup_relative_path);
    let Some(program_name) = args.next() else {
        return;
    };
    if let Err(e) = cgroup.make_child(cgroup_relative_path) {
        fail!("{prog_name:?}: cannot create child cgroup: {e}")
    }
    let child = cgroup
        .open_cgroup_at(cgroup_relative_path)
        .unwrap_or_else(|e| {
            fail!("{prog_name:?}: cannot create child cgroup {cgroup_relative_path:?}: {e}")
        });
    child
        .make_child(Path::new("@inner.service"))
        .unwrap_or_else(|e| {
            fail!("{prog_name:?}: cannot create child cgroup {cgroup_relative_path:?}/@inner.service: {e}")
        });
    let child = child
        .open_cgroup_at(Path::new("@inner.service"))
        .unwrap_or_else(|e| {
            fail!("{prog_name:?}: cannot open child cgroup {cgroup_relative_path:?}/@inner.service: {e}")
        });

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

    let e = std::process::Command::new(&program_name).args(args).exec();
    fail!(
        "{prog_name:?}: cannot spawn child {:?}: {}",
        program_name,
        e
    );
}

fn write_cgroup_value(prog_name: &OsStr, 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!("{prog_name:?}: cannot open {:?}: {}", path, e));
    let () = File::from(fd)
        .write_all(value.as_bytes())
        .unwrap_or_else(|e| {
            fail!(
                "{prog_name:?}: Cannot write {:?} to {:?}: {}",
                name,
                path,
                e
            )
        });
}

fn purge_cgroup(prog_name: &OsStr, cgroup: &cgroup::LeafCgroup, cgroup_relative_path: &Path) {
    let child = match cgroup.open_cgroup_at(cgroup_relative_path) {
        Ok(child_cgroup) => child_cgroup,
        Err(Errno::NOENT) => return,
        Err(other) => {
            fail!("{prog_name:?}: cannot open child cgroup {cgroup_relative_path:?}: {other}")
        }
    };
    child.kill_processes().unwrap_or_else(|e| {
        fail!("{prog_name:?}: cannot kill programs in {cgroup_relative_path:?}: {e}")
    });

    cgroup
        .delete_child(cgroup_relative_path)
        .unwrap_or_else(|e| {
            fail!("{prog_name:?}: delete child cgroup {cgroup_relative_path:?}: {e}")
        });
}