From 26bf0eeaf278be11ac7ddf25bcddc752fe49f7aa Mon Sep 17 00:00:00 2001
From: Andrzej Janik <vosen@vosen.pl>
Date: Thu, 2 Dec 2021 23:47:37 +0100
Subject: Add more tests for CUDA redirection

---
 zluda_inject/build.rs                         |  2 +-
 zluda_inject/tests/helpers/do_cuinit.rs       | 10 ++++++++++
 zluda_inject/tests/helpers/do_cuinit_main.rs  | 23 +++++++++++++++++++++++
 zluda_inject/tests/helpers/indirect_cuinit.rs | 16 ++++++++++++++++
 zluda_inject/tests/inject.rs                  | 20 +++++++++++++++-----
 5 files changed, 65 insertions(+), 6 deletions(-)
 create mode 100644 zluda_inject/tests/helpers/do_cuinit.rs
 create mode 100644 zluda_inject/tests/helpers/do_cuinit_main.rs
 create mode 100644 zluda_inject/tests/helpers/indirect_cuinit.rs

diff --git a/zluda_inject/build.rs b/zluda_inject/build.rs
index c9146e0..743e708 100644
--- a/zluda_inject/build.rs
+++ b/zluda_inject/build.rs
@@ -21,6 +21,7 @@ fn main() -> Result<(), VarError> {
     helpers_dir.push("tests");
     helpers_dir.push("helpers");
     let helpers_dir_as_string = helpers_dir.to_string_lossy();
+    println!("cargo:rerun-if-changed={}", helpers_dir_as_string);
     for rust_file in fs::read_dir(&helpers_dir).unwrap().filter_map(rust_file) {
         let full_file_path = format!(
             "{}{}{}",
@@ -28,7 +29,6 @@ fn main() -> Result<(), VarError> {
             path::MAIN_SEPARATOR,
             rust_file
         );
-        println!("cargo:rerun-if-changed={}", full_file_path);
         let mut rustc_cmd = Command::new(&*rustc_exe);
         if debug {
             rustc_cmd.arg("-g");
diff --git a/zluda_inject/tests/helpers/do_cuinit.rs b/zluda_inject/tests/helpers/do_cuinit.rs
new file mode 100644
index 0000000..468d56c
--- /dev/null
+++ b/zluda_inject/tests/helpers/do_cuinit.rs
@@ -0,0 +1,10 @@
+#![crate_type = "cdylib"]
+
+extern "system" {
+    fn cuInit(flags: u32) -> u32;
+}
+
+#[no_mangle]
+unsafe extern "system" fn do_cuinit(flags: u32) -> u32 {
+    cuInit(flags)
+}
diff --git a/zluda_inject/tests/helpers/do_cuinit_main.rs b/zluda_inject/tests/helpers/do_cuinit_main.rs
new file mode 100644
index 0000000..4a24752
--- /dev/null
+++ b/zluda_inject/tests/helpers/do_cuinit_main.rs
@@ -0,0 +1,23 @@
+#![crate_type = "bin"]
+
+use std::ffi::c_void;
+use std::mem;
+use std::env;
+use std::path::PathBuf;
+use std::ffi::CString;
+
+extern "system" {
+    fn LoadLibraryA(lpFileName: *const i8) -> *mut c_void;
+    fn GetProcAddress(hModule: *mut c_void, lpProcName: *const u8) -> *mut c_void;
+}
+
+fn main() {
+    let current_exe = env::current_exe().unwrap();
+    let mut dll = PathBuf::from(current_exe.parent().unwrap());
+    dll.push("do_cuinit.dll");
+    let dll_cstring = CString::new(dll.to_str().unwrap()).unwrap();
+    let nvcuda = unsafe { LoadLibraryA(dll_cstring.as_ptr()) };
+    let cuInit = unsafe { GetProcAddress(nvcuda, b"do_cuinit\0".as_ptr()) };
+    let cuInit = unsafe { mem::transmute::<_, unsafe extern "system" fn(u32) -> u32>(cuInit) };
+    unsafe { cuInit(0) };
+}
diff --git a/zluda_inject/tests/helpers/indirect_cuinit.rs b/zluda_inject/tests/helpers/indirect_cuinit.rs
new file mode 100644
index 0000000..3b4f1a7
--- /dev/null
+++ b/zluda_inject/tests/helpers/indirect_cuinit.rs
@@ -0,0 +1,16 @@
+#![crate_type = "bin"]
+
+use std::ffi::c_void;
+use std::mem;
+
+extern "system" {
+    fn LoadLibraryA(lpFileName: *const u8) -> *mut c_void;
+    fn GetProcAddress(hModule: *mut c_void, lpProcName: *const u8) -> *mut c_void;
+}
+
+fn main() {
+    let nvcuda = unsafe { LoadLibraryA(b"C:\\Windows\\System32\\nvcuda.dll\0".as_ptr()) };
+    let cuInit = unsafe { GetProcAddress(nvcuda, b"cuInit\0".as_ptr()) };
+    let cuInit = unsafe { mem::transmute::<_, unsafe extern "system" fn(u32) -> u32>(cuInit) };
+    unsafe { cuInit(0) };
+}
diff --git a/zluda_inject/tests/inject.rs b/zluda_inject/tests/inject.rs
index de5fef8..d6430e0 100644
--- a/zluda_inject/tests/inject.rs
+++ b/zluda_inject/tests/inject.rs
@@ -2,15 +2,25 @@ use std::{env, io, path::PathBuf, process::Command};
 
 #[test]
 fn direct_cuinit() -> io::Result<()> {
+    run_process_and_check_for_zluda_dump("direct_cuinit")
+}
+
+#[test]
+fn indirect_cuinit() -> io::Result<()> {
+    run_process_and_check_for_zluda_dump("indirect_cuinit")
+}
+
+#[test]
+fn do_cuinit() -> io::Result<()> {
+    run_process_and_check_for_zluda_dump("do_cuinit_main")
+}
+
+fn run_process_and_check_for_zluda_dump(name: &'static str) -> io::Result<()> {
     let zluda_with_exe = PathBuf::from(env!("CARGO_BIN_EXE_zluda_with"));
     let mut zluda_dump_dll = zluda_with_exe.parent().unwrap().to_path_buf();
     zluda_dump_dll.push("zluda_dump.dll");
     let helpers_dir = env!("HELPERS_OUT_DIR");
-    let exe_under_test = format!(
-        "{}{}direct_cuinit.exe",
-        helpers_dir,
-        std::path::MAIN_SEPARATOR
-    );
+    let exe_under_test = format!("{}{}{}.exe", helpers_dir, std::path::MAIN_SEPARATOR, name);
     let mut test_cmd = Command::new(&zluda_with_exe);
     let test_cmd = test_cmd.arg(&zluda_dump_dll).arg("--").arg(&exe_under_test);
     let test_output = test_cmd.output()?;
-- 
cgit v1.2.3