From 2c0e9b912fe341bd1e513614014fa43b666d257d Mon Sep 17 00:00:00 2001 From: Andrzej Janik Date: Sun, 3 Jan 2021 18:45:48 +0100 Subject: Fix Windows ZLUDA injector (#26) Fix various bugs in injector and redirector, make them more robust and enable building them by default --- zluda_inject/Cargo.toml | 7 +- zluda_inject/src/bin.rs | 199 ++++++++++++++++++++++++++++++++--------------- zluda_inject/src/main.rs | 20 +++-- zluda_inject/src/win.rs | 21 ++--- 4 files changed, 166 insertions(+), 81 deletions(-) (limited to 'zluda_inject') diff --git a/zluda_inject/Cargo.toml b/zluda_inject/Cargo.toml index 5dbb14b..8f1cde5 100644 --- a/zluda_inject/Cargo.toml +++ b/zluda_inject/Cargo.toml @@ -9,8 +9,5 @@ name = "zluda" path = "src/main.rs" [target.'cfg(windows)'.dependencies] -zluda_redirect = { path = "../zluda_redirect" } -winapi = { version = "0.3", features = ["processthreadsapi", "std", "synchapi"] } -detours-sys = "0.1" -clap = "2.33" -guid = "0.1" \ No newline at end of file +winapi = { version = "0.3", features = ["jobapi2", "processthreadsapi", "std", "synchapi"] } +detours-sys = { path = "../detours-sys" } diff --git a/zluda_inject/src/bin.rs b/zluda_inject/src/bin.rs index e021a41..af44d74 100644 --- a/zluda_inject/src/bin.rs +++ b/zluda_inject/src/bin.rs @@ -1,58 +1,39 @@ -extern crate clap; -#[macro_use] -extern crate guid; -extern crate detours_sys; -extern crate winapi; - -use std::error::Error; -use std::ffi::OsStr; +use std::env; +use std::env::Args; use std::mem; -use std::os::windows::ffi::OsStrExt; +use std::path::Path; use std::ptr; +use std::{error::Error, process}; + +use mem::size_of_val; +use winapi::um::{ + jobapi2::{AssignProcessToJobObject, SetInformationJobObject}, + processthreadsapi::{GetExitCodeProcess, ResumeThread}, + synchapi::WaitForSingleObject, + winbase::CreateJobObjectA, + winnt::{ + JobObjectExtendedLimitInformation, HANDLE, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, + JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + }, +}; -use winapi::um::processthreadsapi::{GetExitCodeProcess, ResumeThread}; -use winapi::um::synchapi::WaitForSingleObject; use winapi::um::winbase::{INFINITE, WAIT_FAILED}; -use clap::{App, AppSettings, Arg}; +static REDIRECT_DLL: &'static str = "zluda_redirect.dll"; +static ZLUDA_DLL: &'static str = "nvcuda.dll"; -#[macro_use] -mod win; +include!("../../zluda_redirect/src/payload_guid.rs"); -fn main() -> Result<(), Box> { - let matches = App::new("ZLUDA injector") - .setting(AppSettings::TrailingVarArg) - .arg( - Arg::with_name("EXE") - .help("Path to the executable to be injected with ZLUDA") - .required(true), - ) - .arg( - Arg::with_name("ARGS") - .multiple(true) - .help("Arguments that will be passed to "), - ) - .get_matches(); - let exe = matches.value_of_os("EXE").unwrap(); - let args: Vec<&OsStr> = matches - .values_of_os("ARGS") - .map(|x| x.collect()) - .unwrap_or_else(|| Vec::new()); - let mut cmd_line = Vec::::with_capacity(exe.len() + 2); - cmd_line.push('\"' as u16); - copy_to(exe, &mut cmd_line); - cmd_line.push('\"' as u16); - cmd_line.push(' ' as u16); - args.split_last().map(|(last_arg, args)| { - for arg in args { - cmd_line.reserve(arg.len()); - copy_to(arg, &mut cmd_line); - cmd_line.push(' ' as u16); - } - copy_to(last_arg, &mut cmd_line); - }); - - cmd_line.push(0); +pub fn main_impl() -> Result<(), Box> { + let args = env::args(); + if args.len() == 0 { + print_help(); + process::exit(1); + } + let mut cmd_line = construct_command_line(args); + let injector_path = env::current_exe()?; + let injector_dir = injector_path.parent().unwrap(); + let redirect_path = create_redirect_path(injector_dir); let mut startup_info = unsafe { mem::zeroed::() }; let mut proc_info = unsafe { mem::zeroed::() }; os_call!( @@ -67,22 +48,19 @@ fn main() -> Result<(), Box> { ptr::null(), &mut startup_info as *mut _, &mut proc_info as *mut _, - "zluda_redirect.dll\0".as_ptr() as *const i8, + redirect_path.as_ptr() as *const i8, Option::None ), |x| x != 0 ); - let mut exe_path = std::env::current_dir()? - .as_os_str() - .encode_wide() - .collect::>(); - let guid = guid! {"C225FC0C-00D7-40B8-935A-7E342A9344C1"}; + kill_child_on_process_exit(proc_info.hProcess)?; + let mut zluda_path = create_zluda_path(injector_dir); os_call!( detours_sys::DetourCopyPayloadToProcess( proc_info.hProcess, - mem::transmute(&guid), - exe_path.as_mut_ptr() as *mut _, - (exe_path.len() * mem::size_of::()) as u32 + &PAYLOAD_GUID, + zluda_path.as_mut_ptr() as *mut _, + (zluda_path.len() * mem::size_of::()) as u32 ), |x| x != 0 ); @@ -94,13 +72,110 @@ fn main() -> Result<(), Box> { GetExitCodeProcess(proc_info.hProcess, &mut child_exit_code as *mut _), |x| x != 0 ); - std::process::exit(child_exit_code as i32) + process::exit(child_exit_code as i32) } -fn copy_to(from: &OsStr, to: &mut Vec) { - for x in from.encode_wide() { - to.push(x); +fn kill_child_on_process_exit(child: HANDLE) -> Result<(), Box> { + let job_handle = os_call!(CreateJobObjectA(ptr::null_mut(), ptr::null()), |x| x + != ptr::null_mut()); + let mut info = unsafe { mem::zeroed::() }; + info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + os_call!( + SetInformationJobObject( + job_handle, + JobObjectExtendedLimitInformation, + &mut info as *mut _ as *mut _, + size_of_val(&info) as u32 + ), + |x| x != 0 + ); + os_call!(AssignProcessToJobObject(job_handle, child), |x| x != 0); + Ok(()) +} + +fn print_help() { + println!( + "USAGE: + zluda [ARGS]... +ARGS: + Path to the executable to be injected with ZLUDA + ... Arguments that will be passed to +" + ); +} + +// Adapted from https://docs.microsoft.com/en-us/archive/blogs/twistylittlepassagesallalike/everyone-quotes-command-line-arguments-the-wrong-way +fn construct_command_line(args: Args) -> Vec { + let mut cmd_line = Vec::new(); + let args_len = args.len(); + for (idx, arg) in args.enumerate().skip(1) { + if !arg.contains(&[' ', '\t', '\n', '\u{2B7F}', '\"'][..]) { + cmd_line.extend(arg.encode_utf16()); + } else { + cmd_line.push('"' as u16); // " + let mut char_iter = arg.chars().peekable(); + loop { + let mut current = char_iter.next(); + let mut backslashes = 0; + match current { + Some('\\') => { + backslashes = 1; + while let Some('\\') = char_iter.peek() { + backslashes += 1; + char_iter.next(); + } + current = char_iter.next(); + } + _ => {} + } + match current { + None => { + for _ in 0..(backslashes * 2) { + cmd_line.push('\\' as u16); + } + break; + } + Some('"') => { + for _ in 0..(backslashes * 2 + 1) { + cmd_line.push('\\' as u16); + } + cmd_line.push('"' as u16); + } + Some(c) => { + for _ in 0..backslashes { + cmd_line.push('\\' as u16); + } + let mut temp = [0u16; 2]; + cmd_line.extend(&*c.encode_utf16(&mut temp)); + } + } + } + cmd_line.push('"' as u16); + } + if idx < args_len - 1 { + cmd_line.push(' ' as u16); + } } + cmd_line.push(0); + cmd_line +} + +fn create_redirect_path(injector_dir: &Path) -> Vec { + let mut injector_dir = injector_dir.to_path_buf(); + injector_dir.push(REDIRECT_DLL); + let mut result = injector_dir.to_string_lossy().into_owned().into_bytes(); + result.push(0); + result } -// +fn create_zluda_path(injector_dir: &Path) -> Vec { + let mut injector_dir = injector_dir.to_path_buf(); + injector_dir.push(ZLUDA_DLL); + let mut result = injector_dir + .to_string_lossy() + .as_ref() + .encode_utf16() + .collect::>(); + result.push(0); + result +} diff --git a/zluda_inject/src/main.rs b/zluda_inject/src/main.rs index 6e292f2..f8c1921 100644 --- a/zluda_inject/src/main.rs +++ b/zluda_inject/src/main.rs @@ -1,5 +1,15 @@ -#[cfg(target_os = "windows")] -mod bin; - -#[cfg(not(target_os = "windows"))] -fn main() {} \ No newline at end of file +#[macro_use] +#[cfg(target_os = "windows")] +mod win; +#[cfg(target_os = "windows")] +mod bin; + +use std::error::Error; + +#[cfg(target_os = "windows")] +fn main() -> Result<(), Box> { + bin::main_impl() +} + +#[cfg(not(target_os = "windows"))] +fn main() {} diff --git a/zluda_inject/src/win.rs b/zluda_inject/src/win.rs index ec57ffb..4d7fcdd 100644 --- a/zluda_inject/src/win.rs +++ b/zluda_inject/src/win.rs @@ -48,15 +48,18 @@ macro_rules! last_ident { macro_rules! os_call { ($($path:ident)::+ ($($args:expr),*), $success:expr) => { - let result = unsafe{ $($path)::+ ($($args),+) }; - if !($success)(result) { - let name = last_ident!($($path),+); - let err_code = $crate::win::errno(); - Err($crate::win::OsError{ - function: name, - error_code: err_code as u32, - message: $crate::win::error_string(err_code) - })?; + { + let result = unsafe{ $($path)::+ ($($args),*) }; + if !($success)(result) { + let name = last_ident!($($path),+); + let err_code = $crate::win::errno(); + Err($crate::win::OsError{ + function: name, + error_code: err_code as u32, + message: $crate::win::error_string(err_code) + })?; + } + result } }; } -- cgit v1.2.3