Add failing test for overriding directly linked CUDA dll

2 files changed, 24 insertions, 20 deletions

diff --git a/zluda_inject/tests/inject.rs b/zluda_inject/tests/inject.rs
index f63186a..5a19d8a 100644
--- a/zluda_inject/tests/inject.rs
+++ b/zluda_inject/tests/inject.rs
@@ -1,13 +1,23 @@
-use std::{env, io, process::Command};

+use std::{env, io, path::PathBuf, process::Command};

 

 #[test]

 fn direct_cuinit() -> io::Result<()> {

+    let zluda_with_exe = PathBuf::from(env!("CARGO_BIN_EXE_zluda_with"));

+    let mut zluda_redirect_dll = zluda_with_exe.parent().unwrap().to_path_buf();

+    zluda_redirect_dll.push("zluda_redirect.dll");

     let helpers_dir = env!("HELPERS_OUT_DIR");

-    let mut main_exe = Command::new(format!(

+    let exe_under_test = format!(

         "{}{}direct_cuinit.exe",

         helpers_dir,

         std::path::MAIN_SEPARATOR

-    ));

-    assert!(main_exe.status()?.success());

+    );

+    let mut test_cmd = Command::new(&zluda_with_exe);

+    test_cmd

+        .arg(&zluda_redirect_dll)

+        .arg("--")

+        .arg(&exe_under_test);

+    let test_output = test_cmd.output()?;

+    let stderr_text = String::from_utf8(test_output.stderr).unwrap();

+    assert!(stderr_text.contains("ZLUDA_DUMP"));

     Ok(())

 }

diff --git a/zluda_redirect/src/lib.rs b/zluda_redirect/src/lib.rs
index dfe2385..d695ff7 100644
--- a/zluda_redirect/src/lib.rs
+++ b/zluda_redirect/src/lib.rs
@@ -361,7 +361,7 @@ struct DetourDetachGuard {
     state: DetourUndoState,
     suspended_threads: Vec<*mut c_void>,
     // First element is the original fn, second is the new fn
-    overriden_functions: Vec<(*mut *mut c_void, *mut c_void)>,
+    overriden_functions: Vec<(*mut c_void, *mut c_void)>,
 }
 
 impl DetourDetachGuard {
@@ -371,7 +371,7 @@ impl DetourDetachGuard {
     // also get overriden, so for example ZludaLoadLibraryExW ends calling
     // itself recursively until stack overflow exception occurs
     unsafe fn detour_functions<'a>(
-        override_fn_pairs: Vec<(*mut *mut c_void, *mut c_void)>,
+        override_fn_pairs: Vec<(*mut c_void, *mut c_void)>,
     ) -> Option<Self> {
         let mut result = DetourDetachGuard {
             state: DetourUndoState::DoNothing,
@@ -391,29 +391,23 @@ impl DetourDetachGuard {
             }
         }
         result.overriden_functions.extend_from_slice(&[
+            (CREATE_PROCESS_A as _, ZludaCreateProcessA as _),
+            (CREATE_PROCESS_W as _, ZludaCreateProcessW as _),
             (
-                &mut CREATE_PROCESS_A as *mut _ as _,
-                ZludaCreateProcessA as _,
-            ),
-            (
-                &mut CREATE_PROCESS_W as *mut _ as _,
-                ZludaCreateProcessW as _,
-            ),
-            (
-                &mut CREATE_PROCESS_AS_USER_W as *mut _ as _,
+                CREATE_PROCESS_AS_USER_W as _,
                 ZludaCreateProcessAsUserW as _,
             ),
             (
-                &mut CREATE_PROCESS_WITH_LOGON_W as *mut _ as _,
+                CREATE_PROCESS_WITH_LOGON_W as _,
                 ZludaCreateProcessWithLogonW as _,
             ),
             (
-                &mut CREATE_PROCESS_WITH_TOKEN_W as *mut _ as _,
+                CREATE_PROCESS_WITH_TOKEN_W as _,
                 ZludaCreateProcessWithTokenW as _,
             ),
         ]);
-        for (original_fn, new_fn) in result.overriden_functions.iter().copied() {
-            if DetourAttach(original_fn, new_fn) != NO_ERROR as i32 {
+        for (original_fn, new_fn) in result.overriden_functions.iter_mut() {
+            if DetourAttach(original_fn as *mut _, *new_fn) != NO_ERROR as i32 {
                 return None;
             }
         }
@@ -708,7 +702,7 @@ unsafe fn attach_cuinit(nvcuda_mod: HMODULE) -> Option<DetourDetachGuard> {
                     cuda_unsupported as _
                 }
             };
-        override_fn_pairs.push((&mut original_fn_address as *mut _ as _, override_fn_address));
+        override_fn_pairs.push((original_fn_address as _, override_fn_address));
     }
     DetourDetachGuard::detour_functions(override_fn_pairs)
 }