injector/src/lib.rs

@@ -8,12 +8,14 @@ use log::{debug, error, info, trace};
 use path_clean::PathClean;
 use serde::{Deserialize, Serialize};
 use windows::core::{PSTR, s};
-use windows::Win32::Foundation::{CloseHandle, GetLastError, HANDLE, NTSTATUS};
+use windows::Win32::Foundation::{CloseHandle, GetLastError, HANDLE};
 use windows::Win32::System::Diagnostics::Debug::WriteProcessMemory;
 use windows::Win32::System::LibraryLoader::{GetModuleHandleA, GetProcAddress};
-use windows::Win32::System::Memory::{MEM_COMMIT, MEM_RELEASE, MEM_RESERVE, PAGE_READWRITE,
+use windows::Win32::System::Memory::{MEM_COMMIT, MEM_RELEASE, MEM_RESERVE, PAGE_READWRITE, 
                                      VirtualAllocEx, VirtualFreeEx};
-use windows::Win32::System::Threading::{CREATE_SUSPENDED, CreateProcessA, CreateRemoteThread, GetExitCodeProcess, PROCESS_INFORMATION, ResumeThread, STARTUPINFOA, TerminateProcess, TerminateThread, WaitForSingleObject};
+use windows::Win32::System::Threading::{CREATE_SUSPENDED, CreateProcessA, CreateRemoteThread, 
+                                        PROCESS_INFORMATION, ResumeThread, STARTUPINFOA, 
+                                        TerminateProcess, TerminateThread, WaitForSingleObject};
 
 type FarProcUnwrapped = unsafe extern "system" fn() -> isize;
 type LPThreadStartRoutine = unsafe extern "system" fn(param: *mut c_void) -> u32;
@@ -82,25 +84,20 @@ fn yes() -> bool {
 
 #[derive(Serialize, Deserialize, Debug, Default)]
 pub struct Environment {
-    pub vars: Option<Vec<String>>,
+    vars: Option<Vec<String>>,
     #[serde(default = "yes")]
-    pub use_system_env: bool,
+    use_system_env: bool,
     #[serde(default)]
-    pub environment_append: bool,
+    environment_append: bool,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Launcher {
-    pub executable_file: String,
-    pub cmd_line_args: Option<String>,
-    pub current_dir: Option<String>,
+    executable_file: String,
+    cmd_line_args: Option<String>,
+    current_dir: Option<String>,
     #[serde(default)]
-    pub dll_list: Vec<String>,
-}
-
-pub struct Context {
-    pub proc_info: PROCESS_INFORMATION,
-    pub environment: Vec<u8>,
+    dll_list: Vec<String>,
 }
 
 fn absolute_path(cwd: &Path, path: impl AsRef<Path>) -> PathBuf {
@@ -162,7 +159,7 @@ fn inject_standard(h_target: HANDLE, dll_path: &str) -> Result<(), Error> {
     Ok(())
 }
 
-pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<Context, Error<'a>> {
+pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<(), Error<'a>> {
     let working_dir = match &launcher.current_dir {
         None => std::env::current_dir()?,
         Some(dir) => PathBuf::from_str(dir.as_str())?
@@ -188,8 +185,7 @@ pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<Context
         |cmd_line_args| CString::new(cmd_line_args).unwrap() // TODO: avoid this panic!
     );
     trace!("cmd_line_args: {:?}", cmd_line_args);
-    let mut environment: Vec<u8> = Vec::with_capacity(4096);
-    let environment_ptr = match env.vars {
+    let environment = match env.vars {
         None => None,
         Some(variables) => {
             let mut system_variables = match env.use_system_env {
@@ -213,14 +209,15 @@ pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<Context
                     CString::new(format!("{key}={value}")).unwrap() // TODO: avoid this panic!
                 }).collect::<Vec<_>>();
 
+            let mut result = Vec::with_capacity(4096);
             for var in vars {
-                environment.extend_from_slice(var.as_bytes_with_nul());
+                result.extend_from_slice(var.as_bytes_with_nul());
             }
-            environment.extend_from_slice(b"\0");
-            Some(environment.as_ptr() as *const c_void)
+            result.extend_from_slice(b"\0");
+            Some(result.as_ptr() as *const c_void)
         }
     };
-    trace!("environment: {:?}", environment_ptr);
+    trace!("environment: {:?}", environment);
 
     let current_directory = CString::new(path_buf_to_str!(working_dir)?)?;
 
@@ -235,7 +232,7 @@ pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<Context
             None,
             false,
             CREATE_SUSPENDED,
-            environment_ptr,
+            environment,
             lp_current_directory,
             &startup_info,
             &mut proc_info,
@@ -256,30 +253,11 @@ pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<Context
         debug!("Injection of dll: {inject} succeeded!");
     }
 
-    unsafe { ResumeThread(proc_info.hThread); }
-    info!("Successful injection finished");
-    Ok(Context { proc_info, environment })
-}
-
-pub fn is_process_running(proc_info: PROCESS_INFORMATION) -> bool {
-    let mut exit_code = 0u32;
-    let result = unsafe { GetExitCodeProcess(proc_info.hProcess, &mut exit_code) };
-    match result {
-        Ok(_) => {
-            match NTSTATUS(exit_code as i32) {
-                windows::Win32::Foundation::STILL_ACTIVE => true,
-                _ => {
-                    unsafe {
-                        CloseHandle(proc_info.hThread).unwrap();
-                        CloseHandle(proc_info.hProcess).unwrap();
-                    }
-                    false
-                }
-            }
-        }
-        Err(err) => {
-            error!("Error while getting exit result: {err:?}");
-            true
-        }
+    unsafe {
+        ResumeThread(proc_info.hThread);
+        CloseHandle(proc_info.hThread)?;
+        CloseHandle(proc_info.hProcess)?;
     }
+    info!("Successful injection finished");
+    Ok(())
 }

launcher/src/main.rs

@@ -1,5 +1,5 @@
 use clap::Parser;
-use log::{error, info, trace};
+use log::{error, trace};
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -62,7 +62,7 @@ impl LogLevel {
 }
 
 #[inline]
-fn run(config_file: String) -> Result<injector::Context, Error> {
+fn run(config_file: String) -> Result<(), Error> {
     // Load launcher config
     let content = match std::fs::read_to_string(&config_file) {
         Ok(content) => Ok(content),
@@ -71,10 +71,8 @@ fn run(config_file: String) -> Result<injector::Context, Error> {
     let configuration: Configuration = toml::from_str(content.as_str())?;
 
     trace!("{:?}", configuration);
-    Ok(injector::spawn_process(
-        configuration.launcher,
-        configuration.environment.unwrap_or_default(),
-    )?)
+    injector::spawn_process(configuration.launcher, configuration.environment.unwrap_or_default())?;
+    Ok(())
 }
 
 fn main() {
@@ -82,16 +80,8 @@ fn main() {
     colog::default_builder()
         .filter_level(args.log_level.unwrap_or_default().into_level_filter())
         .init();
-
-    match run(args.config_file) {
-        Ok(context) => {
-            // Since environment pointer has to outlive the process, here we can do several things,
-            // either we join the thread of the process, or we just wait
-            while injector::is_process_running(context.proc_info) {
-                std::thread::sleep(std::time::Duration::from_secs(1))
-            }
-            info!("Application exited");
-        }
-        Err(err) => error!("{}", err.to_string())
+    
+    if let Err(err) = run(args.config_file) {
+        error!("{}", err.to_string())
     }
 }