use std::{io, mem, os::windows::io::AsRawHandle, process::Child, ptr};
use windows_sys::Win32::{
Foundation::{CloseHandle, HANDLE, INVALID_HANDLE_VALUE},
System::{
Diagnostics::ToolHelp::{
CreateToolhelp32Snapshot, TH32CS_SNAPTHREAD, THREADENTRY32, Thread32First, Thread32Next,
},
JobObjects::{AssignProcessToJobObject, CreateJobObjectW, TerminateJobObject},
Threading::{OpenThread, ResumeThread, THREAD_SUSPEND_RESUME},
},
};
#[derive(Debug)]
pub(crate) struct JobObject(HANDLE);
unsafe impl Send for JobObject {}
unsafe impl Sync for JobObject {}
impl JobObject {
pub(crate) fn assign(child: &Child) -> io::Result<Self> {
let job = unsafe { CreateJobObjectW(ptr::null(), ptr::null()) };
if job.is_null() {
return Err(io::Error::last_os_error());
}
let job = Self(job);
if unsafe { AssignProcessToJobObject(job.0, child.as_raw_handle()) } == 0 {
return Err(io::Error::last_os_error());
}
resume_threads(child.id())?;
Ok(job)
}
pub(crate) fn terminate(&self) -> io::Result<()> {
if unsafe { TerminateJobObject(self.0, 1) } == 0 {
return Err(io::Error::last_os_error());
}
Ok(())
}
}
impl Drop for JobObject {
fn drop(&mut self) {
unsafe { CloseHandle(self.0) };
}
}
fn resume_threads(pid: u32) -> io::Result<()> {
let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0) };
if snapshot == INVALID_HANDLE_VALUE {
return Err(io::Error::last_os_error());
}
let result = resume_snapshot_threads(snapshot, pid);
unsafe { CloseHandle(snapshot) };
result
}
fn resume_snapshot_threads(snapshot: HANDLE, pid: u32) -> io::Result<()> {
let mut entry: THREADENTRY32 = unsafe { mem::zeroed() };
entry.dwSize = mem::size_of::<THREADENTRY32>() as u32;
let mut found = unsafe { Thread32First(snapshot, &mut entry) } != 0;
while found {
if entry.th32OwnerProcessID == pid {
let thread = unsafe { OpenThread(THREAD_SUSPEND_RESUME, 0, entry.th32ThreadID) };
if thread.is_null() {
return Err(io::Error::last_os_error());
}
if unsafe { ResumeThread(thread) } == u32::MAX {
let err = io::Error::last_os_error();
unsafe { CloseHandle(thread) };
return Err(err);
}
unsafe { CloseHandle(thread) };
}
found = unsafe { Thread32Next(snapshot, &mut entry) } != 0;
}
Ok(())
}