From 50aee7dd64f58f8eb87746bbef6fe2debee0ea08 Mon Sep 17 00:00:00 2001
From: xavo95 <xavo95@xeondev.com>
Date: Mon, 7 Oct 2024 00:36:34 +0200
Subject: [PATCH] Improve environment handling

---
 injector/src/lib.rs  | 23 ++++++++++++++---------
 launcher/src/main.rs | 18 +++++++++++++-----
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/injector/src/lib.rs b/injector/src/lib.rs
index b917625..6675f0f 100644
--- a/injector/src/lib.rs
+++ b/injector/src/lib.rs
@@ -100,6 +100,11 @@ pub struct Launcher {
     pub dll_list: Vec<String>,
 }
 
+pub struct Context {
+    pub proc_info: PROCESS_INFORMATION,
+    pub environment: Vec<u8>,
+}
+
 fn absolute_path(cwd: &Path, path: impl AsRef<Path>) -> PathBuf {
     let path = path.as_ref();
     if path.is_absolute() {
@@ -159,7 +164,7 @@ fn inject_standard(h_target: HANDLE, dll_path: &str) -> Result<(), Error> {
     Ok(())
 }
 
-pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<(), Error<'a>> {
+pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<Context, Error<'a>> {
     let working_dir = match &launcher.current_dir {
         None => std::env::current_dir()?,
         Some(dir) => PathBuf::from_str(dir.as_str())?
@@ -185,7 +190,8 @@ pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<(), Err
         |cmd_line_args| CString::new(cmd_line_args).unwrap() // TODO: avoid this panic!
     );
     trace!("cmd_line_args: {:?}", cmd_line_args);
-    let environment = match env.vars {
+    let mut environment: Vec<u8> = Vec::with_capacity(4096);
+    let environment_ptr = match env.vars {
         None => None,
         Some(variables) => {
             let mut system_variables = match env.use_system_env {
@@ -209,15 +215,14 @@ pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<(), Err
                     CString::new(format!("{key}={value}")).unwrap() // TODO: avoid this panic!
                 }).collect::<Vec<_>>();
 
-            let mut result = Vec::with_capacity(4096);
             for var in vars {
-                result.extend_from_slice(var.as_bytes_with_nul());
+                environment.extend_from_slice(var.as_bytes_with_nul());
             }
-            result.extend_from_slice(b"\0");
-            Some(result.as_ptr() as *const c_void)
+            environment.extend_from_slice(b"\0");
+            Some(environment.as_ptr() as *const c_void)
         }
     };
-    trace!("environment: {:?}", environment);
+    trace!("environment: {:?}", environment_ptr);
 
     let current_directory = CString::new(path_buf_to_str!(working_dir)?)?;
 
@@ -232,7 +237,7 @@ pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<(), Err
             None,
             false,
             CREATE_SUSPENDED,
-            environment,
+            environment_ptr,
             lp_current_directory,
             &startup_info,
             &mut proc_info,
@@ -259,5 +264,5 @@ pub fn spawn_process<'a>(launcher: Launcher, env: Environment) -> Result<(), Err
         CloseHandle(proc_info.hProcess)?;
     }
     info!("Successful injection finished");
-    Ok(())
+    Ok(Context { proc_info, environment })
 }
\ No newline at end of file
diff --git a/launcher/src/main.rs b/launcher/src/main.rs
index 0719043..f4ec487 100644
--- a/launcher/src/main.rs
+++ b/launcher/src/main.rs
@@ -1,6 +1,7 @@
 use clap::Parser;
 use log::{error, trace};
 use serde::{Deserialize, Serialize};
+use injector::Context;
 
 #[derive(Serialize, Deserialize, Debug)]
 struct Configuration {
@@ -62,7 +63,7 @@ impl LogLevel {
 }
 
 #[inline]
-fn run(config_file: String) -> Result<(), Error> {
+fn run(config_file: String) -> Result<injector::Context, Error> {
     // Load launcher config
     let content = match std::fs::read_to_string(&config_file) {
         Ok(content) => Ok(content),
@@ -71,8 +72,10 @@ fn run(config_file: String) -> Result<(), Error> {
     let configuration: Configuration = toml::from_str(content.as_str())?;
 
     trace!("{:?}", configuration);
-    injector::spawn_process(configuration.launcher, configuration.environment.unwrap_or_default())?;
-    Ok(())
+    Ok(injector::spawn_process(
+        configuration.launcher, 
+        configuration.environment.unwrap_or_default()
+    )?)
 }
 
 fn main() {
@@ -81,7 +84,12 @@ fn main() {
         .filter_level(args.log_level.unwrap_or_default().into_level_filter())
         .init();
     
-    if let Err(err) = run(args.config_file) {
-        error!("{}", err.to_string())
+    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
+            std::thread::sleep(std::time::Duration::from_secs(u64::MAX))
+        },
+        Err(err) => error!("{}", err.to_string())
     }
 }
\ No newline at end of file