bot-forge 1.0.2

#[cfg(windows)]
use std::os::windows::io::AsRawHandle;
#[cfg(windows)]
use std::process::Child;
#[cfg(windows)]
use windows_sys::Win32::System::JobObjects::{
    AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
    JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation,
    QueryInformationJobObject, SetInformationJobObject,
};

#[cfg(windows)]
use crate::error::ForgeError;

#[cfg(not(windows))]
pub(crate) struct ProcessJob;

#[cfg(windows)]
pub(crate) struct ProcessJob(windows_sys::Win32::Foundation::HANDLE);

#[cfg(not(windows))]
pub(crate) fn job_ref(_: &()) -> Option<&ProcessJob> {
    None
}

#[cfg(windows)]
pub(crate) fn job_ref(job: &ProcessJob) -> Option<&ProcessJob> {
    Some(job)
}

#[cfg(windows)]
impl ProcessJob {
    /// Assign a child process to a kill-on-close Windows Job Object.
    ///
    /// # Errors
    ///
    /// Returns [`ForgeError`] when the job cannot be created, configured, or associated with the
    /// live child process.
    pub(crate) fn assign(child: &Child) -> Result<Self, ForgeError> {
        // SAFETY: Null security attributes and name request an unnamed job with default security.
        let handle = unsafe { CreateJobObjectW(std::ptr::null(), std::ptr::null()) };
        if handle.is_null() {
            return Err(os_error("failed to create the Windows Job Object"));
        }
        let job = Self(handle);
        // SAFETY: The Windows structure is a plain C data structure for which an all-zero value is
        // the documented baseline before individual limit fields are populated.
        let mut information: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { std::mem::zeroed() };
        information.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
        // SAFETY: `handle` was created successfully above and `information` has the exact layout
        // and byte size required by `JobObjectExtendedLimitInformation`.
        let configured = unsafe {
            SetInformationJobObject(
                handle,
                JobObjectExtendedLimitInformation,
                std::ptr::from_ref(&information).cast(),
                std::mem::size_of_val(&information) as u32,
            )
        };
        let process = child.as_raw_handle() as windows_sys::Win32::Foundation::HANDLE;
        // SAFETY: `handle` is owned by `job`; `process` is borrowed from the live child process.
        if configured == 0 || unsafe { AssignProcessToJobObject(handle, process) } == 0 {
            return Err(os_error("failed to configure the Windows Job Object"));
        }
        Ok(job)
    }

    /// Terminate every process currently associated with this job.
    pub(crate) fn terminate(&self) {
        // SAFETY: `self.0` remains an owned, open job handle until `Drop` closes it.
        unsafe {
            windows_sys::Win32::System::JobObjects::TerminateJobObject(self.0, 1);
        }
    }

    /// Return peak memory charged to the job, or `None` when Windows cannot provide it.
    pub(crate) fn peak_memory_mib(&self) -> Option<u64> {
        // SAFETY: The output C structure accepts zero initialization before Windows fills it.
        let mut information: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { std::mem::zeroed() };
        // SAFETY: `self.0` is an open job handle and the mutable output buffer has the advertised
        // structure layout and size for the duration of the call.
        let success = unsafe {
            QueryInformationJobObject(
                self.0,
                JobObjectExtendedLimitInformation,
                std::ptr::from_mut(&mut information).cast(),
                std::mem::size_of_val(&information) as u32,
                std::ptr::null_mut(),
            )
        };
        (success != 0).then_some(information.PeakJobMemoryUsed as u64 / 1024 / 1024)
    }
}

#[cfg(windows)]
fn os_error(message: &str) -> ForgeError {
    ForgeError::Command(format!("{message}：{}", std::io::Error::last_os_error()))
}

#[cfg(windows)]
impl Drop for ProcessJob {
    fn drop(&mut self) {
        // SAFETY: `self.0` is owned by this value and is closed exactly once from `Drop`.
        unsafe {
            windows_sys::Win32::Foundation::CloseHandle(self.0);
        }
    }
}