Add a library for dumping kernels arguments before and after launch (#18)

1 files changed, 39 insertions, 26 deletions

diff --git a/zluda_inject/src/bin.rs b/zluda_inject/src/bin.rs
index af44d74..ce83fe9 100644
--- a/zluda_inject/src/bin.rs
+++ b/zluda_inject/src/bin.rs
@@ -1,8 +1,7 @@
-use std::env;

-use std::env::Args;

 use std::mem;

 use std::path::Path;

 use std::ptr;

+use std::{env, ops::Deref};

 use std::{error::Error, process};

 

 use mem::size_of_val;

@@ -25,15 +24,15 @@ static ZLUDA_DLL: &'static str = "nvcuda.dll";
 include!("../../zluda_redirect/src/payload_guid.rs");

 

 pub fn main_impl() -> Result<(), Box<dyn Error>> {

-    let args = env::args();

-    if args.len() == 0 {

-        print_help();

-        process::exit(1);

+    let args = env::args().collect::<Vec<_>>();

+    if args.len() <= 1 {

+        print_help_and_exit();

     }

-    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 inject_path, cmd) = create_inject_path(&args[1..], injector_dir);

+    let mut cmd_line = construct_command_line(cmd);

     let mut startup_info = unsafe { mem::zeroed::<detours_sys::_STARTUPINFOW>() };

     let mut proc_info = unsafe { mem::zeroed::<detours_sys::_PROCESS_INFORMATION>() };

     os_call!(

@@ -54,13 +53,12 @@ pub fn main_impl() -> Result<(), Box<dyn Error>> {
         |x| x != 0

     );

     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,

             &PAYLOAD_GUID,

-            zluda_path.as_mut_ptr() as *mut _,

-            (zluda_path.len() * mem::size_of::<u16>()) as u32

+            inject_path.as_mut_ptr() as *mut _,

+            (inject_path.len() * mem::size_of::<u16>()) as u32

         ),

         |x| x != 0

     );

@@ -93,22 +91,29 @@ fn kill_child_on_process_exit(child: HANDLE) -> Result<(), Box<dyn Error>> {
     Ok(())

 }

 

-fn print_help() {

+fn print_help_and_exit() -> ! {

+    let current_exe = env::current_exe().unwrap();

+    let exe_name = current_exe.file_name().unwrap().to_string_lossy();

     println!(

         "USAGE:

-    zluda <EXE> [ARGS]...

+    {0} -- <EXE> [ARGS]...

+    {0} <DLL> -- <EXE> [ARGS]...

 ARGS:

-    <EXE>        Path to the executable to be injected with ZLUDA

+    <DLL>        DLL to ne injected instead of system nvcuda.dll, if not provided

+                 will use nvcuda.dll from the directory where {0} is located

+    <EXE>        Path to the executable to be injected with <DLL>

     <ARGS>...    Arguments that will be passed to <EXE>

-"

+",

+        exe_name

     );

+    process::exit(1)

 }

 

 // 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<u16> {

+fn construct_command_line(args: &[String]) -> Vec<u16> {

     let mut cmd_line = Vec::new();

     let args_len = args.len();

-    for (idx, arg) in args.enumerate().skip(1) {

+    for (idx, arg) in args.iter().enumerate() {

         if !arg.contains(&[' ', '\t', '\n', '\u{2B7F}', '\"'][..]) {

             cmd_line.extend(arg.encode_utf16());

         } else {

@@ -168,14 +173,22 @@ fn create_redirect_path(injector_dir: &Path) -> Vec<u8> {
     result

 }

 

-fn create_zluda_path(injector_dir: &Path) -> Vec<u16> {

-    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::<Vec<_>>();

-    result.push(0);

-    result

+fn create_inject_path<'a>(args: &'a [String], injector_dir: &Path) -> (Vec<u16>, &'a [String]) {

+    if args.get(0).map(Deref::deref) == Some("--") {

+        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::<Vec<_>>();

+        result.push(0);

+        (result, &args[1..])

+    } else if args.get(1).map(Deref::deref) == Some("--") {

+        let mut dll_path = args[0].encode_utf16().collect::<Vec<_>>();

+        dll_path.push(0);

+        (dll_path, &args[2..])

+    } else {

+        print_help_and_exit()

+    }

 }