use std::{
hash::{Hash, Hasher},
io,
os::windows::{
prelude::{
AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle, OwnedHandle,
RawHandle,
},
raw::HANDLE,
},
path::Path,
process::Child,
time::Duration,
};
use sysinfo::{PidExt, ProcessExt, SystemExt};
use winapi::{shared::minwindef::FALSE, um::processthreadsapi::OpenProcess};
use crate::process::{BorrowedProcess, OwnedProcessModule, Process, PROCESS_INJECTION_ACCESS};
#[repr(transparent)]
#[derive(Debug)]
pub struct OwnedProcess(OwnedHandle);
unsafe impl Send for OwnedProcess {}
unsafe impl Sync for OwnedProcess {}
impl AsRawHandle for OwnedProcess {
fn as_raw_handle(&self) -> HANDLE {
self.0.as_raw_handle()
}
}
impl AsHandle for OwnedProcess {
fn as_handle(&self) -> BorrowedHandle<'_> {
self.0.as_handle()
}
}
impl IntoRawHandle for OwnedProcess {
fn into_raw_handle(self) -> RawHandle {
self.0.into_raw_handle()
}
}
impl FromRawHandle for OwnedProcess {
unsafe fn from_raw_handle(handle: HANDLE) -> Self {
Self(unsafe { OwnedHandle::from_raw_handle(handle) })
}
}
impl From<Child> for OwnedProcess {
fn from(child: Child) -> Self {
Self::from_child(child)
}
}
impl TryFrom<BorrowedProcess<'_>> for OwnedProcess {
type Error = io::Error;
fn try_from(process: BorrowedProcess<'_>) -> Result<Self, Self::Error> {
process.try_to_owned()
}
}
impl PartialEq for OwnedProcess {
fn eq(&self, other: &Self) -> bool {
self.borrowed() == other.borrowed()
}
}
impl Eq for OwnedProcess {}
impl Hash for OwnedProcess {
fn hash<H: Hasher>(&self, state: &mut H) {
self.borrowed().hash(state);
}
}
impl Process for OwnedProcess {
type Handle = OwnedHandle;
fn borrowed(&self) -> BorrowedProcess<'_> {
unsafe { BorrowedProcess::from_handle_unchecked(self.as_handle()) }
}
fn try_clone(&self) -> Result<Self, io::Error> {
self.borrowed().try_to_owned()
}
fn into_handle(self) -> Self::Handle {
self.0
}
unsafe fn from_handle_unchecked(handle: Self::Handle) -> Self {
Self(handle)
}
fn current_handle() -> Self::Handle {
unsafe { OwnedHandle::from_raw_handle(Self::raw_current_handle()) }
}
fn find_module_by_name(
&self,
module_name: impl AsRef<Path>,
) -> Result<Option<OwnedProcessModule>, io::Error> {
if let Some(module) = self.borrowed().find_module_by_name(module_name)? {
Ok(Some(module.try_to_owned()?))
} else {
Ok(None)
}
}
fn find_module_by_path(
&self,
module_path: impl AsRef<Path>,
) -> Result<Option<OwnedProcessModule>, io::Error> {
if let Some(module) = self.borrowed().find_module_by_path(module_path)? {
Ok(Some(module.try_to_owned()?))
} else {
Ok(None)
}
}
fn wait_for_module_by_name(
&self,
module_name: impl AsRef<Path>,
timeout: Duration,
) -> Result<Option<OwnedProcessModule>, io::Error> {
if let Some(module) = self
.borrowed()
.wait_for_module_by_name(module_name, timeout)?
{
Ok(Some(module.try_to_owned()?))
} else {
Ok(None)
}
}
fn wait_for_module_by_path(
&self,
module_path: impl AsRef<Path>,
timeout: Duration,
) -> Result<Option<OwnedProcessModule>, io::Error> {
if let Some(module) = self
.borrowed()
.wait_for_module_by_path(module_path, timeout)?
{
Ok(Some(module.try_to_owned()?))
} else {
Ok(None)
}
}
}
impl OwnedProcess {
pub fn from_pid(pid: u32) -> Result<OwnedProcess, io::Error> {
let handle = unsafe {
OpenProcess(
PROCESS_INJECTION_ACCESS,
FALSE,
pid,
)
};
if handle.is_null() {
return Err(io::Error::last_os_error());
}
Ok(unsafe { OwnedProcess::from_raw_handle(handle) })
}
#[must_use]
pub fn all() -> Vec<OwnedProcess> {
let mut system = sysinfo::System::new();
system.refresh_processes();
system
.processes()
.values()
.map(|process| process.pid())
.filter_map(|pid| OwnedProcess::from_pid(pid.as_u32()).ok())
.collect()
}
#[must_use]
pub fn find_all_by_name(name: impl AsRef<str>) -> Vec<OwnedProcess> {
let mut system = sysinfo::System::new();
system.refresh_processes();
system
.processes()
.values()
.filter(move |process| process.name().contains(name.as_ref()))
.map(|process| process.pid())
.filter_map(|pid| OwnedProcess::from_pid(pid.as_u32()).ok())
.collect()
}
#[must_use]
pub fn find_first_by_name(name: impl AsRef<str>) -> Option<OwnedProcess> {
let mut system = sysinfo::System::new();
system.refresh_processes();
system
.processes()
.values()
.filter(move |process| process.name().contains(name.as_ref()))
.map(|process| process.pid())
.find_map(|pid| OwnedProcess::from_pid(pid.as_u32()).ok())
}
#[must_use]
pub fn from_child(child: Child) -> OwnedProcess {
unsafe { OwnedProcess::from_raw_handle(child.into_raw_handle()) }
}
#[must_use]
pub unsafe fn borrowed_static(&self) -> BorrowedProcess<'static> {
unsafe {
BorrowedProcess::from_handle_unchecked(BorrowedHandle::borrow_raw(self.as_raw_handle()))
}
}
pub fn try_clone(&self) -> Result<Self, io::Error> {
self.borrowed().try_to_owned()
}
#[allow(clippy::must_use_candidate)]
pub fn leak(self) -> BorrowedProcess<'static> {
unsafe { self.borrowed_static() }
}
#[must_use]
pub const fn kill_on_drop(self) -> ProcessKillGuard {
ProcessKillGuard(self)
}
}
#[derive(Debug, shrinkwraprs::Shrinkwrap)]
#[shrinkwrap(mutable)]
pub struct ProcessKillGuard(pub OwnedProcess);
impl Drop for ProcessKillGuard {
fn drop(&mut self) {
let _ = self.0.kill();
}
}