From c869a0d6110fe703057ce0a3979e4f10c34038a7 Mon Sep 17 00:00:00 2001
From: Andrzej Janik <vosen@vosen.pl>
Date: Thu, 3 Feb 2022 12:28:42 +0100
Subject: Add tests for injecting into CLR process

---
 zluda_inject/build.rs                             |   9 ++++++
 zluda_inject/tests/helpers/do_cuinit_main.rs      |   6 ++--
 zluda_inject/tests/helpers/do_cuinit_main_clr.cs  |  34 ++++++++++++++++++++++
 zluda_inject/tests/helpers/do_cuinit_main_clr.exe | Bin 0 -> 4608 bytes
 zluda_inject/tests/helpers/indirect_cuinit.rs     |   6 ++--
 zluda_inject/tests/inject.rs                      |   5 ++++
 6 files changed, 54 insertions(+), 6 deletions(-)
 create mode 100644 zluda_inject/tests/helpers/do_cuinit_main_clr.cs
 create mode 100644 zluda_inject/tests/helpers/do_cuinit_main_clr.exe

diff --git a/zluda_inject/build.rs b/zluda_inject/build.rs
index 743e708..1e425a7 100644
--- a/zluda_inject/build.rs
+++ b/zluda_inject/build.rs
@@ -50,6 +50,15 @@ fn main() -> Result<(), VarError> {
             .arg(full_file_path);
         assert!(rustc_cmd.status().unwrap().success());
     }
+    std::fs::copy(
+        format!(
+            "{}{}do_cuinit_main_clr.exe",
+            helpers_dir_as_string,
+            path::MAIN_SEPARATOR
+        ),
+        format!("{}{}do_cuinit_main_clr.exe", out_dir, path::MAIN_SEPARATOR),
+    )
+    .unwrap();
     println!("cargo:rustc-env=HELPERS_OUT_DIR={}", &out_dir);
     Ok(())
 }
diff --git a/zluda_inject/tests/helpers/do_cuinit_main.rs b/zluda_inject/tests/helpers/do_cuinit_main.rs
index 4a24752..ab3516d 100644
--- a/zluda_inject/tests/helpers/do_cuinit_main.rs
+++ b/zluda_inject/tests/helpers/do_cuinit_main.rs
@@ -17,7 +17,7 @@ fn main() {
     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) };
+    let cu_init = unsafe { GetProcAddress(nvcuda, b"do_cuinit\0".as_ptr()) };
+    let cu_init = unsafe { mem::transmute::<_, unsafe extern "system" fn(u32) -> u32>(cu_init) };
+    unsafe { cu_init(0) };
 }
diff --git a/zluda_inject/tests/helpers/do_cuinit_main_clr.cs b/zluda_inject/tests/helpers/do_cuinit_main_clr.cs
new file mode 100644
index 0000000..666c237
--- /dev/null
+++ b/zluda_inject/tests/helpers/do_cuinit_main_clr.cs
@@ -0,0 +1,34 @@
+using System;
+using System.IO;
+using System.Reflection;
+using System.Runtime.InteropServices;
+
+namespace Zluda
+{
+    class Program
+    {
+        [UnmanagedFunctionPointer(CallingConvention.Winapi)]
+        private delegate int CuInit(int flags);
+
+        static int Main(string[] args)
+        {
+            DirectoryInfo exeDirectory = Directory.GetParent(Assembly.GetEntryAssembly().Location);
+            string dllPath = Path.Combine(exeDirectory.ToString(), "do_cuinit.dll");
+            IntPtr nvcuda = NativeMethods.LoadLibrary(dllPath);
+            if (nvcuda == IntPtr.Zero)
+                return 1;
+            IntPtr doCuinitPtr = NativeMethods.GetProcAddress(nvcuda, "do_cuinit");
+            CuInit cuinit = (CuInit)Marshal.GetDelegateForFunctionPointer(doCuinitPtr, typeof(CuInit));
+            return cuinit(0);
+        }
+    }
+
+    static class NativeMethods
+    {
+        [DllImport("kernel32.dll")]
+        public static extern IntPtr LoadLibrary(string dllToLoad);
+
+        [DllImport("kernel32.dll")]
+        public static extern IntPtr GetProcAddress(IntPtr hModule, string procedureName);
+    }
+}
\ No newline at end of file
diff --git a/zluda_inject/tests/helpers/do_cuinit_main_clr.exe b/zluda_inject/tests/helpers/do_cuinit_main_clr.exe
new file mode 100644
index 0000000..b8e4975
Binary files /dev/null and b/zluda_inject/tests/helpers/do_cuinit_main_clr.exe differ
diff --git a/zluda_inject/tests/helpers/indirect_cuinit.rs b/zluda_inject/tests/helpers/indirect_cuinit.rs
index 3b4f1a7..f254dc1 100644
--- a/zluda_inject/tests/helpers/indirect_cuinit.rs
+++ b/zluda_inject/tests/helpers/indirect_cuinit.rs
@@ -10,7 +10,7 @@ extern "system" {
 
 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) };
+    let cu_init = unsafe { GetProcAddress(nvcuda, b"cuInit\0".as_ptr()) };
+    let cu_init = unsafe { mem::transmute::<_, unsafe extern "system" fn(u32) -> u32>(cu_init) };
+    unsafe { cu_init(0) };
 }
diff --git a/zluda_inject/tests/inject.rs b/zluda_inject/tests/inject.rs
index d6430e0..15e5e04 100644
--- a/zluda_inject/tests/inject.rs
+++ b/zluda_inject/tests/inject.rs
@@ -15,6 +15,11 @@ fn do_cuinit() -> io::Result<()> {
     run_process_and_check_for_zluda_dump("do_cuinit_main")
 }
 
+#[test]
+fn do_cuinit_clr() -> io::Result<()> {
+    run_process_and_check_for_zluda_dump("do_cuinit_main_clr")
+}
+
 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();
-- 
cgit v1.2.3