diff --git a/crates/vm/Cargo.toml b/crates/vm/Cargo.toml index 8a684c93dd9..4eac727c2fa 100644 --- a/crates/vm/Cargo.toml +++ b/crates/vm/Cargo.toml @@ -128,6 +128,7 @@ features = [ "Win32_System_Environment", "Win32_System_IO", "Win32_System_Ioctl", + "Win32_System_JobObjects", "Win32_System_Kernel", "Win32_System_LibraryLoader", "Win32_System_Memory", diff --git a/crates/vm/src/stdlib/_winapi.rs b/crates/vm/src/stdlib/_winapi.rs index fe2e752b99f..b210badd00a 100644 --- a/crates/vm/src/stdlib/_winapi.rs +++ b/crates/vm/src/stdlib/_winapi.rs @@ -615,6 +615,78 @@ mod _winapi { }) } + #[pyfunction] + fn CreateJobObject( + _security_attributes: PyObjectRef, + name: OptionalArg>, + vm: &VirtualMachine, + ) -> PyResult { + let handle = unsafe { + match name.flatten() { + Some(name) => { + let name_wide = name.as_wtf8().to_wide_with_nul(); + windows_sys::Win32::System::JobObjects::CreateJobObjectW( + null(), + name_wide.as_ptr(), + ) + } + None => windows_sys::Win32::System::JobObjects::CreateJobObjectW(null(), null()), + } + }; + if handle.is_null() { + return Err(vm.new_last_os_error()); + } + Ok(WinHandle(handle)) + } + + #[pyfunction] + fn AssignProcessToJobObject( + job: WinHandle, + process: WinHandle, + vm: &VirtualMachine, + ) -> PyResult<()> { + let ret = unsafe { + windows_sys::Win32::System::JobObjects::AssignProcessToJobObject(job.0, process.0) + }; + if ret == 0 { + return Err(vm.new_last_os_error()); + } + Ok(()) + } + + #[pyfunction] + fn TerminateJobObject(job: WinHandle, exit_code: u32, vm: &VirtualMachine) -> PyResult<()> { + let ret = + unsafe { windows_sys::Win32::System::JobObjects::TerminateJobObject(job.0, exit_code) }; + if ret == 0 { + return Err(vm.new_last_os_error()); + } + Ok(()) + } + + #[pyfunction] + fn SetJobObjectKillOnClose(job: WinHandle, vm: &VirtualMachine) -> PyResult<()> { + use windows_sys::Win32::System::JobObjects::{ + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, + SetInformationJobObject, + }; + let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { core::mem::zeroed() }; + info.BasicLimitInformation.LimitFlags = + windows_sys::Win32::System::JobObjects::JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + let ret = unsafe { + SetInformationJobObject( + job.0, + JobObjectExtendedLimitInformation, + &info as *const _ as *const core::ffi::c_void, + core::mem::size_of::() as u32, + ) + }; + if ret == 0 { + return Err(vm.new_last_os_error()); + } + Ok(()) + } + #[pyfunction] fn GetModuleFileName(handle: isize, vm: &VirtualMachine) -> PyResult { let mut path: Vec = vec![0; MAX_PATH as usize];