From fd1c13560f29e9f6e43d19b5cbe48dcd1351bcd6 Mon Sep 17 00:00:00 2001
From: Andrzej Janik <vosen@vosen.pl>
Date: Tue, 30 Nov 2021 01:25:51 +0100
Subject: Add failing test for overriding directly linked CUDA dll

---
 zluda_inject/tests/inject.rs | 18 ++++++++++++++----
 zluda_redirect/src/lib.rs    | 26 ++++++++++----------------
 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)
 }
-- 
cgit v1.2.3