From 6f9b395b0b7dd57117a2cf20db67d3d136c0cc13 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Tue, 23 Jul 2024 13:54:46 +0800 Subject: [PATCH] Windows: Support wait all subprocess finished --- Cargo.toml | 2 +- src/main.rs | 9 ++++ src/windows.rs | 135 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 145 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 5317ee1..dd8fd8b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,4 +10,4 @@ subprocess = "0.2.9" yaml-rust = "0.4.5" [target.'cfg(windows)'.dependencies] -winapi = { version = "0.3", features = ["wincon", "winuser"] } +winapi = { version = "0.3", features = ["errhandlingapi", "impl-default", "ioapiset", "jobapi2", "wincon", "winuser"] } diff --git a/src/main.rs b/src/main.rs index 80b23d8..7a0726b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,7 +18,10 @@ pub fn print_usage(prog: &str, opts: &Options) { #[derive(Debug, derive_more::Display, derive_more::From)] enum Error { + #[cfg(not(windows))] Popen(subprocess::PopenError), + #[cfg(windows)] + Popen(windows::PopenError), Exited, } @@ -105,11 +108,17 @@ impl Main { } } + #[cfg(not(windows))] fn call(cml: Vec) -> Result { let mut p = subprocess::Popen::create(&cml, subprocess::PopenConfig::default())?; p.wait() } + #[cfg(windows)] + fn call(cml: Vec) -> Result { + windows::call(&cml).map(|c| ExitStatus::Exited(c)) + } + fn restore(&self) -> Result<(), Error> { let cml = match self._cfg.restore_command() { Some(cml) => cml, diff --git a/src/windows.rs b/src/windows.rs index 191f88c..33ba20d 100644 --- a/src/windows.rs +++ b/src/windows.rs @@ -1,4 +1,24 @@ +use std::ffi::{OsStr, OsString}; +use std::mem::size_of; +use std::os::windows::ffi::OsStrExt; +use std::ptr::{addr_of_mut, null, null_mut}; +use winapi::ctypes::c_void; +use winapi::shared::basetsd::ULONG_PTR; +use winapi::shared::minwindef::DWORD; +use winapi::um::errhandlingapi::GetLastError; +use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE}; +use winapi::um::ioapiset::{CreateIoCompletionPort, GetQueuedCompletionStatus}; +use winapi::um::jobapi2::{AssignProcessToJobObject, SetInformationJobObject}; +use winapi::um::minwinbase::LPOVERLAPPED; +use winapi::um::processthreadsapi::{ + CreateProcessW, GetExitCodeProcess, ResumeThread, PROCESS_INFORMATION, STARTUPINFOW, +}; +use winapi::um::winbase::{CreateJobObjectA, CREATE_SUSPENDED, INFINITE}; use winapi::um::wincon::GetConsoleWindow; +use winapi::um::winnt::{ + JobObjectAssociateCompletionPortInformation, JOBOBJECT_ASSOCIATE_COMPLETION_PORT, + JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO, +}; use winapi::um::winuser::{ShowWindow, SW_HIDE, SW_SHOW}; fn console_show_window(n_cmd_show: i32) -> bool { @@ -17,3 +37,118 @@ pub fn show_window() -> bool { pub fn hide_window() -> bool { console_show_window(SW_HIDE) } + +#[derive(Debug, derive_more::Display, derive_more::From)] +pub enum PopenError { + CreateJobFailed, + CreateProcessFailed, + AssignJobFailed, +} + +pub fn call>(argv: &[S]) -> Result { + let job = unsafe { CreateJobObjectA(null_mut(), null()) }; + if job.is_null() { + println!("Failed to create job: {}.", unsafe { GetLastError() }); + return Err(PopenError::CreateJobFailed); + } + let io_port = unsafe { CreateIoCompletionPort(INVALID_HANDLE_VALUE, null_mut(), 0, 1) }; + if io_port.is_null() { + unsafe { CloseHandle(job) }; + println!("CreateIoCompletionPort: {}", unsafe { GetLastError() }); + return Err(PopenError::CreateJobFailed); + } + let mut port = JOBOBJECT_ASSOCIATE_COMPLETION_PORT::default(); + port.CompletionKey = job; + port.CompletionPort = io_port; + let ok = unsafe { + SetInformationJobObject( + job, + JobObjectAssociateCompletionPortInformation, + addr_of_mut!(port) as *mut c_void, + size_of::() as u32, + ) != 0 + }; + if !ok { + unsafe { CloseHandle(job) }; + unsafe { CloseHandle(io_port) }; + println!("SetInformationJobObject: {}", unsafe { GetLastError() }); + return Err(PopenError::CreateJobFailed); + } + let mut si = STARTUPINFOW::default(); + let mut pi = PROCESS_INFORMATION::default(); + let mut cml = OsString::new(); + for i in argv.iter() { + let t = i.as_ref(); + if !cml.is_empty() { + cml.push(" "); + } + if t.to_string_lossy().find(' ').is_some() { + cml.push("\""); + cml.push( + t.to_string_lossy() + .replace("\\", "\\\\") + .replace("\"", "\\\""), + ); + cml.push("\""); + } else { + cml.push( + t.to_string_lossy() + .replace("\\", "\\\\") + .replace("\"", "\\\""), + ); + } + } + let mut cmlw: Vec<_> = cml.encode_wide().collect(); + cmlw.resize(cmlw.len() + 1000, 0); + let re = unsafe { + CreateProcessW( + null(), + cmlw.as_mut_ptr(), + null_mut(), + null_mut(), + 1, + CREATE_SUSPENDED, + null_mut(), + null_mut(), + addr_of_mut!(si), + addr_of_mut!(pi), + ) != 0 + }; + if !re { + println!("Failed to create process: {}.", unsafe { GetLastError() }); + unsafe { CloseHandle(job) }; + unsafe { CloseHandle(io_port) }; + return Err(PopenError::CreateProcessFailed); + } + let re = unsafe { AssignProcessToJobObject(job, pi.hProcess) != 0 }; + if !re { + println!("Failed to assign process to job."); + unsafe { CloseHandle(job) }; + unsafe { CloseHandle(pi.hProcess) }; + unsafe { CloseHandle(pi.hThread) }; + unsafe { CloseHandle(io_port) }; + return Err(PopenError::AssignJobFailed); + } + unsafe { ResumeThread(pi.hThread) }; + let mut code = DWORD::default(); + let mut key = ULONG_PTR::default(); + let mut overlapped: LPOVERLAPPED = null_mut(); + while unsafe { + GetQueuedCompletionStatus( + io_port, + addr_of_mut!(code), + addr_of_mut!(key), + addr_of_mut!(overlapped), + INFINITE, + ) + } != 0 + && !(job.wrapping_sub(key).is_null() && code == JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO) + {} + let mut c = DWORD::default(); + unsafe { GetExitCodeProcess(pi.hProcess, addr_of_mut!(c)) }; + unsafe { CloseHandle(job) }; + unsafe { CloseHandle(io_port) }; + unsafe { CloseHandle(pi.hThread) }; + unsafe { CloseHandle(pi.hProcess) }; + Ok(c) +}