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

use std::cell::RefCell;
use std::ffi::OsStr;
use std::fs::File;
use std::io::{Read as _, Seek as _, Write as _};
use std::os::unix::prelude::*;

use std::path::{Component, Path, PathBuf};
use std::rc::Rc;

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

#[derive(Debug)]
pub(crate) struct Cgroup {
    fd: Vec<OwnedFd>,
}

impl AsFd for Cgroup {
    fn as_fd(&self) -> BorrowedFd<'_> {
        self.fd.last().unwrap().as_fd()
    }
}

fn assert_single_component(component: &Path) {
    match component.as_os_str().as_bytes() {
        b"" | b"." | b".." => panic!("bad component"),
        c if c.contains(&b'\0') => panic!("NUL in component"),
        c if c.contains(&b'/') => panic!("/ in component"),
        _ => {}
    }
}

// Wrapper around openat2() with better defaults.
pub fn openat2_simple(fd: &dyn AsFd, path: &Path, flags: OFlags) -> Result<OwnedFd, Errno> {
    rustix::fs::openat2(
        fd.as_fd(),
        path,
        OFlags::CLOEXEC | OFlags::NOCTTY | flags,
        Mode::empty(),
        ResolveFlags::NO_SYMLINKS | ResolveFlags::NO_MAGICLINKS | ResolveFlags::NO_XDEV,
    )
}

pub const DEFAULT_LEAF: &str = "$inner.service";

pub fn check_path(path: &Path) -> Result<(), String> {
    let bytes = path.as_os_str().as_bytes();
    // Path::components() skips ., so use string manipulation instead.
    for component in bytes[path.is_absolute() as usize..].split(|&b| b == b'/') {
        if matches!(component, b"" | b"." | b"..") {
            return Err(format!("cgroup path {path:?} isn't canonical"));
        }
    }
    Ok(())
}

fn push_child_fds(fds: &mut Vec<(Rc<RefCell<Dir>>, PathBuf)>, fd: OwnedFd) {
    // The rustix source code shows that Dir::new() never fails.
    let child_fd = Rc::new(RefCell::new(Dir::new(fd).unwrap()));
    while let Some(element) = child_fd.borrow_mut().next() {
        let element = element.expect("Iterating through a cgroup directory failed?");
        if element.file_type() != rustix::fs::FileType::Directory {
            continue;
        }
        match element.file_name().to_bytes() {
            b"." | b".." => {}
            other => {
                let other = Path::new(OsStr::from_bytes(other)).to_owned();
                assert_single_component(&other);
                fds.push((child_fd.clone(), other));
            }
        }
    }
}

// Remove all subdirectories of the given directory recursively,
// but not the directory itself.  The directory file descriptor
// is closed.
//
// This isn't the most efficient possible algorithm, but
// simplicity is more important than performance in this
// case.  Also, it keeps open more file descriptors than
// strictly necessary, but Spectrum runs with a very high
// limit for the number of open file descriptors, and it
// uses shallow control group hierarchies.
fn remove_child_directories(dirfd: OwnedFd) -> Result<(), Errno> {
    let mut fds = Vec::new();
    // Push the children of this directory onto the stack.
    push_child_fds(&mut fds, dirfd);
    while let Some((d, path)) = fds.pop() {
        assert_single_component(&path);
        // Try to delete the directory.  If that fails because there are child
        // directories, push the child directories onto the stack, then push
        // this directory again.
        match rustix::fs::unlinkat(
            // The rustix source code shows that Dir::fd() never fails.
            d.borrow().fd().unwrap(),
            &path,
            AtFlags::REMOVEDIR,
        ) {
            Err(Errno::NOTEMPTY) => {}
            Ok(()) => continue,
            Err(bad) => return Err(bad),
        }
        let fd = openat2_simple(
            // The rustix source code shows that Dir::fd() never fails.
            &d.borrow().fd().unwrap(),
            &path,
            OFlags::DIRECTORY | OFlags::RDONLY,
        )?;
        // Process child directories first, then attempt to delete the
        // directory again.
        fds.push((d, path));
        push_child_fds(&mut fds, fd);
    }
    Ok(())
}

// If the path is absolute, make it relative.
// Otherwise, read the current cgroup from /proc/thread-self/cgroup
// and prepend it to the path.
fn prepend_current_cgroup_if_needed(path: &Path) -> PathBuf {
    if let Ok(suffix) = path.strip_prefix("/") {
        suffix.to_owned()
    } else {
        // /proc/thread-self is the same as /proc/self, except for the current
        // thread instead of the initial thread.  In this case, the two are
        // identical, but using /proc/thread-self is better practice as it is
        // correct in more cases.  Reading /proc/thread-self/cgroup should
        // never fail unless the system is seriously broken.
        let current_cgroup = std::fs::read("/proc/thread-self/cgroup")
            .expect("cannot read /proc/thread-self/cgroup");
        // Using this on a system without cgroups v2 mounted is user error
        // and not supported.
        let current_cgroup = current_cgroup
            .strip_prefix(b"0::/")
            .and_then(|e| e.strip_suffix(b"\n"))
            .expect("you don't have cgroups v2 mounted");
        let mut current_cgroup = PathBuf::from(OsStr::from_bytes(current_cgroup));
        // Strip the implied $inner.service suffix.
        // This is used to satisfy the "no internal processes" rule.
        if current_cgroup.ends_with(Path::new(DEFAULT_LEAF)) {
            assert!(current_cgroup.pop());
        }
        current_cgroup.push(path);
        current_cgroup
    }
}

pub(crate) fn write_value(fd: &dyn AsFd, name: &Path, value: &[u8]) -> Result<(), String> {
    let path = Path::new(name);
    let fd = openat2_simple(fd, path, OFlags::WRONLY)
        .map_err(|e| format!("Cannot open {name:?}: {e}"))?;
    File::from(fd).write_all(value).map_err(|e| {
        format!(
            "Cannot write {:?} to {name:?}: {e}",
            OsStr::from_bytes(value)
        )
    })
}

impl Cgroup {
    pub fn open_beneath(&self, path: &Path, flags: OFlags) -> Result<OwnedFd, Errno> {
        openat2_simple(self, path, flags)
    }

    pub fn new(path: &Path) -> Result<Self, String> {
        let cgroup_root = rustix::fs::openat2(
            CWD,
            Path::new("/sys/fs/cgroup"),
            OFlags::CLOEXEC | OFlags::DIRECTORY | OFlags::RDONLY,
            Mode::empty(),
            ResolveFlags::NO_SYMLINKS | ResolveFlags::NO_MAGICLINKS,
        )
        .map_err(|e| format!("Cannot open /sys/fs/cgroup: {e}"))?;
        let mut cgroup = Self {
            fd: vec![(cgroup_root)],
        };

        let path = prepend_current_cgroup_if_needed(path);
        for component in path.components() {
            let component = match component {
                Component::Normal(component) => component,
                _ => unreachable!(),
            };
            let sub_fd = cgroup
                .open_beneath(Path::new(component), OFlags::RDONLY | OFlags::DIRECTORY)
                .map_err(|e| format!("Cannot open sub-cgroup {component:?}: {e}"))?;
            // Take a shared lock on the *previous* file descriptor.
            rustix::fs::flock(&cgroup, FlockOperation::LockShared)
                .map_err(|e| format!("Cannot lock sub-cgroup {component:?}: {e}"))?;
            cgroup.fd.push(sub_fd);
        }
        // Take an exclusive lock on the final file descriptor.
        rustix::fs::flock(&cgroup, FlockOperation::LockExclusive)
            .map_err(|e| format!("Cannot lock {path:?}: {e}"))?;
        Ok(cgroup)
    }

    pub fn wait_for_empty(fd: &dyn AsFd) -> std::io::Result<()> {
        let wait_file = openat2_simple(fd, Path::new("cgroup.events"), OFlags::RDONLY)?;
        let poll_fd = wait_file.as_raw_fd();
        let mut wait_fd = File::from(wait_file);
        let mut fds = libc::pollfd {
            fd: poll_fd,
            events: libc::POLLPRI | libc::POLLERR,
            revents: 0,
        };
        let mut v = vec![];
        loop {
            v.clear();
            wait_fd
                .seek(std::io::SeekFrom::Start(0))
                .expect("Seek on control group file should succeed");
            wait_fd
                .read_to_end(&mut v)
                .expect("reading from control group should work");
            // Check that the cgroup isn't already empty.  If it was,
            // the kernel would not send an event and poll() would wait
            // forever.
            if v.split(|&c| c == b'\n').any(|line| line == b"populated 0") {
                break;
            }
            // SAFETY: FFI call, valid arguments, fds contains 1 element
            if unsafe { libc::poll(&raw mut fds, 1, -1) } != 1 {
                panic!("poll failed");
            }
        }
        Ok(())
    }

    pub fn purge_child(&mut self, path: &Path) -> Result<(), String> {
        assert_single_component(path);
        // See if we can just delete the child directly.
        match rustix::fs::unlinkat(&self, Path::new(path), AtFlags::REMOVEDIR) {
            // If the cgroup was successfully deleted, or if it
            // has already been deleted, we are done.
            Ok(()) | Err(Errno::NOENT) => return Ok(()),
            // If this cgroup is in use, keep going.
            Err(Errno::BUSY) => {}
            Err(e) => return Err(format!("Cannot purge {path:?}: {e}")),
        }

        let sub_fd = match self.open_beneath(path, OFlags::RDONLY | OFlags::DIRECTORY) {
            Ok(sub_fd) => sub_fd,
            Err(Errno::NOENT) => return Ok(()),
            Err(e) => {
                return Err(format!("Cannot open sub-cgroup {path:?}: {e}",));
            }
        };

        // Take an exclusive lock on the cgroup that is about to be
        // removed.  This avoids concurrent executions of this program
        // operating on deleted sub-cgroups.
        rustix::fs::flock(&sub_fd, FlockOperation::LockExclusive)
            .map_err(|e| format!("Cannot lock sub-cgroup: {e}"))?;

        // Drop the exclusive lock on the original cgroup,
        // This avoids blocking concurrent operations on other
        // child cgroups while the cgroup is being purged,
        // or while waiting for programs to exit.
        rustix::fs::flock(&self, FlockOperation::LockShared)
            .map_err(|e| format!("Cannot relock: {e}"))?;

        // Kill all processes in the child cgroup.
        write_value(&sub_fd, Path::new("cgroup.kill"), b"1")?;

        // Wait for the child cgroup to become empty.
        Self::wait_for_empty(&sub_fd)
            .map_err(|e| format!("Cannot wait for cgroup to become empty: {e}"))
            .inspect_err(|_| {
                self.fd.pop().unwrap();
            })?;

        // Remove the child cgroup and its contents recursively.
        remove_child_directories(sub_fd).map_err(|e| format!("Cannot remove: {e}"))?;
        // Re-take an exclusive lock on the parent of the cgroup being purged.
        // Otherwise, a concurrent instance of cgroup-setup might create a cgroup
        // only for this one to delete it.  The other instance could then try to
        // create a sub-cgroup of a deleted cgroup, which would fail.  Waiting
        // until nobody is using the parent cgroup ensures these problems can't
        // happen.
        //
        // This must happen *after* the lock on the cgroup being purged is released.
        // Another instance of the program might have a shared lock on the parent
        // and be waiting for an exclusive lock on the child.  Trying to take an
        // exclusive lock on the parent while a lock is held on the child would
        // result in an ABBA deadlock.
        rustix::fs::flock(&self, FlockOperation::LockExclusive)
            .map_err(|e| format!("Cannot re-lock exclusively: {e}"))?;
        // Delete the cgroup.  If it's been re-created in the meantime
        // and is currently in use, this is not an error.  Another
        // process deleting the cgroup is also not an error.  Both of
        // these can happen because of the time period between
        // remove_child_directories() closing the file descriptor
        // (releasing its lock) and the above call to flock().
        match rustix::fs::unlinkat(&self, path, AtFlags::REMOVEDIR) {
            Ok(()) | Err(Errno::BUSY) | Err(Errno::NOENT) => Ok(()),
            Err(e) => Err(format!("Cannot delete: {e}")),
        }
    }
}