#![cfg(windows)]
#![warn(unsafe_op_in_unsafe_fn)]
use std::{mem, os::windows::io::AsRawHandle, process, ptr};
use winapi::{
shared::{ntdef::NTSTATUS, ntstatus::STATUS_SUCCESS},
um::{
handleapi::CloseHandle,
jobapi2::{AssignProcessToJobObject, CreateJobObjectW, QueryInformationJobObject},
libloaderapi::{GetModuleHandleA, GetProcAddress},
winnt::{
JobObjectBasicAccountingInformation, HANDLE, JOBOBJECT_BASIC_ACCOUNTING_INFORMATION,
},
},
};
#[cfg(windows_process_extensions_main_thread_handle)]
use winapi::shared::minwindef::DWORD;
#[cfg(not(windows_process_extensions_main_thread_handle))]
use once_cell::sync::Lazy;
use crate::util::units::Second;
const HUNDRED_NS_PER_MS: i64 = 10;
#[cfg(not(windows_process_extensions_main_thread_handle))]
#[allow(non_upper_case_globals)]
static NtResumeProcess: Lazy<unsafe extern "system" fn(ProcessHandle: HANDLE) -> NTSTATUS> =
Lazy::new(|| {
let ntdll = unsafe { GetModuleHandleA(b"ntdll.dll\0".as_ptr().cast()) };
assert!(!ntdll.is_null(), "GetModuleHandleA failed");
let nt_resume_process =
unsafe { GetProcAddress(ntdll, b"NtResumeProcess\0".as_ptr().cast()) };
assert!(!nt_resume_process.is_null(), "GetProcAddress failed");
unsafe { mem::transmute(nt_resume_process) }
});
pub struct CPUTimer {
job_object: HANDLE,
}
impl CPUTimer {
pub unsafe fn start_suspended_process(child: &process::Child) -> Self {
let job_object = unsafe { CreateJobObjectW(ptr::null_mut(), ptr::null_mut()) };
assert!(!job_object.is_null(), "CreateJobObjectW failed");
let ret = unsafe { AssignProcessToJobObject(job_object, child.as_raw_handle()) };
assert!(ret != 0, "AssignProcessToJobObject failed");
#[cfg(windows_process_extensions_main_thread_handle)]
{
let ret = unsafe { ResumeThread(child.main_thread_handle().as_raw_handle()) };
assert!(ret != -1 as DWORD, "ResumeThread failed");
}
#[cfg(not(windows_process_extensions_main_thread_handle))]
{
let ret = unsafe { NtResumeProcess(child.as_raw_handle()) };
assert!(ret == STATUS_SUCCESS, "NtResumeProcess failed");
}
Self { job_object }
}
pub fn stop(&self) -> (Second, Second) {
let mut job_object_info =
mem::MaybeUninit::<JOBOBJECT_BASIC_ACCOUNTING_INFORMATION>::uninit();
let res = unsafe {
QueryInformationJobObject(
self.job_object,
JobObjectBasicAccountingInformation,
job_object_info.as_mut_ptr().cast(),
mem::size_of::<JOBOBJECT_BASIC_ACCOUNTING_INFORMATION>() as u32,
ptr::null_mut(),
)
};
if res != 0 {
let job_object_info = unsafe { job_object_info.assume_init() };
let user: i64 = unsafe { job_object_info.TotalUserTime.QuadPart() } / HUNDRED_NS_PER_MS;
let kernel: i64 =
unsafe { job_object_info.TotalKernelTime.QuadPart() } / HUNDRED_NS_PER_MS;
(user as f64 * 1e-6, kernel as f64 * 1e-6)
} else {
(0.0, 0.0)
}
}
}
impl Drop for CPUTimer {
fn drop(self: &mut Self) {
unsafe { CloseHandle(self.job_object) };
}
}