rightkit-child-process 0.5.9

use crate::{JobAssignmentFailure, JobAssignmentMode, ProcessHandle, ProcessPriority};
use std::{
    io,
    os::windows::{
        io::{AsHandle, AsRawHandle, FromRawHandle, OwnedHandle},
        process::CommandExt,
    },
    process::Command,
};
use windows::{
    core::BOOL,
    Win32::{
        Foundation::{HANDLE, WAIT_OBJECT_0, WAIT_TIMEOUT},
        System::{
            Diagnostics::ToolHelp::{
                CreateToolhelp32Snapshot, Thread32First, Thread32Next, TH32CS_SNAPTHREAD,
                THREADENTRY32,
            },
            JobObjects::{
                AssignProcessToJobObject, CreateJobObjectW, IsProcessInJob,
                JobObjectExtendedLimitInformation, SetInformationJobObject, TerminateJobObject,
                JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_BREAKAWAY_OK,
                JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK,
            },
            Threading::{
                OpenProcess, OpenThread, ProcessPowerThrottling, ResumeThread,
                SetProcessInformation, TerminateProcess, WaitForSingleObject,
                BELOW_NORMAL_PRIORITY_CLASS, CREATE_BREAKAWAY_FROM_JOB, CREATE_NEW_CONSOLE,
                CREATE_NEW_PROCESS_GROUP, CREATE_NO_WINDOW, CREATE_SUSPENDED, DETACHED_PROCESS,
                IDLE_PRIORITY_CLASS, PROCESS_CREATION_FLAGS,
                PROCESS_POWER_THROTTLING_CURRENT_VERSION, PROCESS_POWER_THROTTLING_EXECUTION_SPEED,
                PROCESS_POWER_THROTTLING_STATE, PROCESS_QUERY_LIMITED_INFORMATION,
                PROCESS_SYNCHRONIZE, THREAD_SUSPEND_RESUME,
            },
        },
    },
};

fn raw(handle: &impl AsHandle) -> HANDLE {
    HANDLE(handle.as_handle().as_raw_handle() as _)
}

pub(crate) fn os_error(error: windows::core::Error) -> io::Error {
    // windows captures GetLastError before cleanup can overwrite it.
    // HRESULT_FROM_WIN32 retains the original error in its low 16 bits.
    let code = error.code().0 as u32;
    if code & 0xffff_0000 == 0x8007_0000 {
        io::Error::from_raw_os_error((code & 0xffff) as i32)
    } else {
        io::Error::other(error)
    }
}

fn own(handle: HANDLE) -> OwnedHandle {
    // Each successful Win32 creation returns one newly owned kernel handle.
    unsafe { OwnedHandle::from_raw_handle(handle.0 as _) }
}

/// Priority-class creation flag. `Normal` contributes no bits, so the default composition is
/// the one this crate always used.
pub(crate) fn priority_class(priority: ProcessPriority) -> PROCESS_CREATION_FLAGS {
    match priority {
        ProcessPriority::Normal => PROCESS_CREATION_FLAGS(0),
        ProcessPriority::BelowNormal | ProcessPriority::Efficient => BELOW_NORMAL_PRIORITY_CLASS,
        ProcessPriority::Idle => IDLE_PRIORITY_CLASS,
    }
}

/// EcoQoS: ask Windows to run the process on its efficiency class (power throttling of
/// execution speed). The process may still be suspended.
pub(crate) fn apply_efficiency(process: &impl AsHandle) -> io::Result<()> {
    let state = PROCESS_POWER_THROTTLING_STATE {
        Version: PROCESS_POWER_THROTTLING_CURRENT_VERSION,
        ControlMask: PROCESS_POWER_THROTTLING_EXECUTION_SPEED,
        StateMask: PROCESS_POWER_THROTTLING_EXECUTION_SPEED,
    };
    // SAFETY: `state` is a valid, live PROCESS_POWER_THROTTLING_STATE of the stated size.
    unsafe {
        SetProcessInformation(
            raw(process),
            ProcessPowerThrottling,
            &state as *const _ as *const std::ffi::c_void,
            std::mem::size_of::<PROCESS_POWER_THROTTLING_STATE>() as u32,
        )
    }
    .map_err(os_error)
}

/// Full creation flags for a spawn: the suspension or breakaway mode and the console flag this
/// crate already sets, with the priority class added. Nothing is removed.
pub(crate) fn creation_flags(
    detached: bool,
    hidden: bool,
    new_console: bool,
    priority: ProcessPriority,
) -> PROCESS_CREATION_FLAGS {
    let mode = if detached {
        CREATE_BREAKAWAY_FROM_JOB
    } else {
        CREATE_SUSPENDED
    };
    // A new console wins over hidden: `CREATE_NEW_CONSOLE` and `CREATE_NO_WINDOW` are not
    // combined (`windows_new_console` clears `hidden`, this keeps the rule in one place).
    let flags = if new_console {
        CREATE_NEW_CONSOLE | mode
    } else if hidden {
        CREATE_NO_WINDOW | mode
    } else {
        mode
    };
    flags | priority_class(priority)
}

pub(crate) fn configure(
    command: &mut Command,
    detached: bool,
    hidden: bool,
    new_console: bool,
    priority: ProcessPriority,
) {
    command.creation_flags(creation_flags(detached, hidden, new_console, priority).0);
}

/// Creation flags of the detached fallback used when the caller's job denies breakaway:
/// `CREATE_NEW_PROCESS_GROUP` plus the priority class and one console flag. No
/// `CREATE_BREAKAWAY_FROM_JOB`. A hidden child gets `CREATE_NO_WINDOW` (its own windowless
/// console, so console programs such as powershell.exe keep working with piped stdio) instead of
/// `DETACHED_PROCESS`; the two must not be combined because `CREATE_NO_WINDOW` is ignored when
/// `DETACHED_PROCESS` is present (Microsoft "Process Creation Flags"). A visible child keeps
/// `DETACHED_PROCESS`; a child that asked for a new console gets `CREATE_NEW_CONSOLE` instead.
pub(crate) fn in_job_detached_flags(
    hidden: bool,
    new_console: bool,
    priority: ProcessPriority,
) -> PROCESS_CREATION_FLAGS {
    let console = if new_console {
        CREATE_NEW_CONSOLE
    } else if hidden {
        CREATE_NO_WINDOW
    } else {
        DETACHED_PROCESS
    };
    console | CREATE_NEW_PROCESS_GROUP | priority_class(priority)
}

/// Replaces the command's creation flags with the in-job detached fallback.
pub(crate) fn configure_in_job_detached(
    command: &mut Command,
    hidden: bool,
    new_console: bool,
    priority: ProcessPriority,
) {
    command.creation_flags(in_job_detached_flags(hidden, new_console, priority).0);
}

/// Whether a failed breakaway launch was refused by the caller's job object
/// (`ERROR_ACCESS_DENIED`).
pub(crate) fn is_breakaway_denied(error: &io::Error) -> bool {
    error.raw_os_error() == Some(5)
}

#[derive(Debug)]
pub(crate) struct KillOnCloseJob(OwnedHandle);

/// The processes that were in a job when [`KillOnCloseJob::hold_members`] ran, each held open
/// (process id and handle). The handle keeps the process object alive after the process exits,
/// so its exit can still be waited for once the job no longer lists it.
#[derive(Debug, Default)]
pub(crate) struct JobMembers(Vec<(u32, OwnedHandle)>);

impl JobMembers {
    /// Waits until `deadline` (the absolute instant shared with the caller's other steps; `None`
    /// waits as long as it takes) for every held process to have exited: its process object is
    /// signalled, so it holds no handle, lock or file any more. `Ok(false)` once the deadline has
    /// passed with one still running. A job's `ActiveProcesses` can reach zero earlier than this,
    /// when a process's last thread has exited (rightkit#1260 follow-up).
    ///
    /// Deadline-safe like [`KillOnCloseJob::wait_empty_until`]: exits confirmed only after
    /// `deadline` are `Ok(false)`.
    pub(crate) fn wait_exited_until(
        &self,
        deadline: Option<std::time::Instant>,
    ) -> io::Result<bool> {
        // INFINITE (u32::MAX) is reserved for an unbounded wait; this is just very long.
        const LONG_WAIT_MS: u32 = u32::MAX - 1;
        for (_, process) in &self.0 {
            loop {
                let millis = deadline.map_or(LONG_WAIT_MS, |d| {
                    let left = d.saturating_duration_since(std::time::Instant::now());
                    // Round up sub-millisecond waits so a remainder still waits.
                    let millis = left.as_millis()
                        + u128::from(!left.subsec_nanos().is_multiple_of(1_000_000));
                    millis.min(u128::from(LONG_WAIT_MS)) as u32
                });
                // SAFETY: the handle is open with PROCESS_SYNCHRONIZE.
                match unsafe { WaitForSingleObject(raw(process), millis) } {
                    WAIT_OBJECT_0 => break,
                    WAIT_TIMEOUT => {
                        if deadline.is_some_and(|d| std::time::Instant::now() >= d) {
                            return Ok(false);
                        }
                    }
                    _ => return Err(io::Error::last_os_error()),
                }
            }
        }
        Ok(!deadline.is_some_and(|d| std::time::Instant::now() > d))
    }
}

impl KillOnCloseJob {
    /// Waits until `deadline` (an absolute instant shared with the caller's other steps; `None`
    /// waits as long as it takes) for the job to hold no active process
    /// (`JobObjectBasicAccountingInformation.ActiveProcesses == 0`). `Ok(false)` once the
    /// deadline has passed; a query error is returned, never read as empty. Job termination is
    /// asynchronous, so a leader's exit alone does not show the tree is gone (rightkit#1260).
    ///
    /// Deadline-safe: an empty job counts only when the query that saw it returned by
    /// `deadline`; one seen empty later (the last poll's sleep overran it, or the caller was
    /// descheduled before calling) is `Ok(false)`.
    pub(crate) fn wait_empty_until(
        &self,
        deadline: Option<std::time::Instant>,
    ) -> io::Result<bool> {
        const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(10);
        loop {
            let count = crate::windows_job::WindowsJob::active_count_of(raw(&self.0))?;
            let now = std::time::Instant::now();
            let left = deadline.map(|d| d.saturating_duration_since(now));
            let past = deadline.is_some_and(|d| now > d);
            if count == 0 {
                return Ok(!past);
            }
            if past || left.is_some_and(|l| l.is_zero()) {
                return Ok(false);
            }
            std::thread::sleep(left.map_or(POLL_INTERVAL, |l| POLL_INTERVAL.min(l)));
        }
    }

    /// The process ids in the job now (leader included while it is listed), grandchildren
    /// included. A snapshot; see [`crate::WindowsJob::process_ids`].
    pub(crate) fn process_ids(&self) -> io::Result<Vec<u32>> {
        crate::windows_job::WindowsJob::ids_of(raw(&self.0), 32)
    }

    /// Opens, for waiting, every process in the job now. Call it before the kill: a process is
    /// out of the job's accounting (and unlistable) once its last thread has exited, which can
    /// be before the kernel has closed its handles and signalled its process object, so only a
    /// handle taken beforehand can still wait for that last step (rightkit#1260 follow-up).
    ///
    /// Best effort by design: the job-emptiness wait stays the primary proof, and this only adds
    /// to it. A listing that fails holds nothing, and a member that cannot be opened (or that
    /// left before it was opened) is skipped.
    pub(crate) fn hold_members(&self) -> JobMembers {
        let Ok(pids) = self.process_ids() else {
            return JobMembers::default();
        };
        let mut held = Vec::with_capacity(pids.len());
        for pid in pids {
            if let Some(process) = self.open_member(pid) {
                held.push((pid, process));
            }
        }
        JobMembers(held)
    }

    /// The process object of `pid`, opened with the rights to wait on it, only if that process is
    /// in this job. The check is what keeps a pid that was listed, exited and was handed to an
    /// unrelated process before this open from being held and waited for. The handle pins the
    /// process object, so once it is held the pid cannot name anything else.
    pub(crate) fn open_member(&self, pid: u32) -> Option<OwnedHandle> {
        let access = PROCESS_SYNCHRONIZE | PROCESS_QUERY_LIMITED_INFORMATION;
        // SAFETY: OpenProcess takes plain values and returns a new handle, owned from here on.
        let process = own(unsafe { OpenProcess(access, false, pid) }.ok()?);
        let mut inside = BOOL(0);
        // SAFETY: both handles are open and `inside` is a valid out parameter.
        unsafe { IsProcessInJob(raw(&process), Some(raw(&self.0)), &mut inside) }.ok()?;
        inside.as_bool().then_some(process)
    }

    fn assign(process: &ProcessHandle, breakaway: crate::owned::BreakawayMode) -> io::Result<Self> {
        let job = Self(own(
            unsafe { CreateJobObjectW(None, None) }.map_err(os_error)?
        ));
        let mut limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default();
        limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
        match breakaway {
            crate::owned::BreakawayMode::None => {}
            // Only a process that asks (CREATE_BREAKAWAY_FROM_JOB) leaves the job.
            crate::owned::BreakawayMode::Explicit => {
                limits.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_BREAKAWAY_OK;
            }
            // Every created process leaves the job without asking.
            crate::owned::BreakawayMode::Silent => {
                limits.BasicLimitInformation.LimitFlags |=
                    JOB_OBJECT_LIMIT_BREAKAWAY_OK | JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK;
            }
        }
        unsafe {
            SetInformationJobObject(
                raw(&job.0),
                JobObjectExtendedLimitInformation,
                &limits as *const _ as *const _,
                std::mem::size_of_val(&limits) as u32,
            )
        }
        .map_err(os_error)?;
        unsafe { AssignProcessToJobObject(raw(&job.0), raw(process)) }.map_err(os_error)?;
        Ok(job)
    }
}

fn assignment_result(
    result: io::Result<()>,
    mode: JobAssignmentMode,
    stage: &'static str,
    failures: &mut Vec<JobAssignmentFailure>,
) -> io::Result<()> {
    match result {
        Ok(()) => Ok(()),
        Err(error) if mode == JobAssignmentMode::BestEffort => {
            failures.push(JobAssignmentFailure { stage, error });
            Ok(())
        }
        Err(error) => Err(error),
    }
}

pub(crate) fn assign_jobs(
    process: &ProcessHandle,
    caller_job: Option<OwnedHandle>,
    mode: JobAssignmentMode,
    breakaway: crate::owned::BreakawayMode,
    owned_job: &mut Option<std::sync::Arc<KillOnCloseJob>>,
    failures: &mut Vec<JobAssignmentFailure>,
) -> io::Result<()> {
    // Caller job precedes the empty nested RightKit job. No user code runs yet.
    if let Some(caller) = caller_job {
        assignment_result(
            unsafe { AssignProcessToJobObject(raw(&caller), raw(process)) }.map_err(os_error),
            mode,
            "caller job assignment",
            failures,
        )?;
    }
    match KillOnCloseJob::assign(process, breakaway) {
        Ok(job) => *owned_job = Some(std::sync::Arc::new(job)),
        Err(error) => assignment_result(
            Err(error),
            mode,
            "kill-on-close job setup/assignment",
            failures,
        )?,
    }
    Ok(())
}

/// std closes the primary thread handle after CreateProcessW. The child is
/// still suspended, so its sole initial thread can be reopened safely by ID.
pub(crate) fn resume(pid: u32) -> io::Result<()> {
    let snapshot =
        own(unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0) }.map_err(os_error)?);
    let mut entry = THREADENTRY32 {
        dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
        ..Default::default()
    };
    unsafe { Thread32First(raw(&snapshot), &mut entry) }.map_err(os_error)?;
    loop {
        if entry.th32OwnerProcessID == pid {
            let thread =
                own(
                    unsafe { OpenThread(THREAD_SUSPEND_RESUME, false, entry.th32ThreadID) }
                        .map_err(os_error)?,
                );
            if unsafe { ResumeThread(raw(&thread)) } == u32::MAX {
                return Err(io::Error::last_os_error());
            }
            return Ok(());
        }
        if unsafe { Thread32Next(raw(&snapshot), &mut entry) }.is_err() {
            break;
        }
    }
    Err(io::Error::new(
        io::ErrorKind::NotFound,
        "suspended child's primary thread was not found",
    ))
}

/// How long an unbounded termination waits for a leader that `TerminateProcess` found already
/// exiting.
pub(crate) const EXITING_LEADER_WAIT: std::time::Duration = std::time::Duration::from_secs(2);

/// Terminates the job (when there is one) and the leader. A leader that `TerminateProcess`
/// refuses because it is already exiting is waited for at most [`EXITING_LEADER_WAIT`], and never
/// past `deadline`: a bounded termination passes its absolute deadline, read when the wait
/// starts, so this step never outlasts it (rightkit#1260).
pub(crate) fn terminate(
    process: &ProcessHandle,
    job: Option<&KillOnCloseJob>,
    code: u32,
    deadline: Option<std::time::Instant>,
) -> io::Result<()> {
    let refused = false;
    // Do not skip the job when its original parent already exited.
    let job_result = match job {
        Some(_) if refused => Ok(()),
        Some(job) => unsafe { TerminateJobObject(raw(&job.0), code) }.map_err(os_error),
        None => Ok(()),
    };
    // Even on failed assignment, preserve direct-child termination/reap.
    let direct_result = if process.try_wait()?.is_none() {
        let killed = if refused {
            Err(windows::core::Error::from_hresult(
                windows::Win32::Foundation::ERROR_ACCESS_DENIED.to_hresult(),
            ))
        } else {
            unsafe { TerminateProcess(raw(process), code) }
        };
        match killed {
            Ok(()) => Ok(()),
            // TerminateProcess fails with ACCESS_DENIED while a process is already exiting
            // (job termination is asynchronous): give that exit a bounded moment to land.
            Err(_)
                if process
                    .wait_timeout(deadline.map_or(EXITING_LEADER_WAIT, |d| {
                        d.saturating_duration_since(std::time::Instant::now())
                            .min(EXITING_LEADER_WAIT)
                    }))?
                    .is_some() =>
            {
                Ok(())
            }
            Err(error) => Err(os_error(error)),
        }
    } else {
        Ok(())
    };
    job_result.and(direct_result)
}