diff --git a/Cargo.lock b/Cargo.lock index 4061290..0d3ab14 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -719,7 +719,7 @@ checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" [[package]] name = "saekawa" -version = "0.2.0" +version = "0.3.1" dependencies = [ "aes", "anyhow", diff --git a/Cargo.toml b/Cargo.toml index bef34ae..525969d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "saekawa" -version = "0.2.0" +version = "0.3.1" authors = ["beerpsi "] edition = "2021" license = "0BSD" @@ -39,4 +39,4 @@ sha1 = "0.10.6" ureq = { version = "2.8.0", features = ["json"] } url = "2.4.1" widestring = "1.0.2" -winapi = { version = "0.3.9", features = ["winhttp", "minwindef", "debugapi"] } +winapi = { version = "0.3.9", features = ["winhttp", "minwindef", "debugapi", "synchapi"] } diff --git a/src/lib.rs b/src/lib.rs index 4299a55..e5e0759 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,11 +6,18 @@ mod log; mod saekawa; mod types; -use ::log::error; +use std::ffi::c_void; +use std::{ptr, thread}; + +use ::log::{error, warn}; use lazy_static::lazy_static; use url::Url; -use winapi::shared::minwindef::{BOOL, DWORD, HINSTANCE, LPVOID, TRUE}; -use winapi::um::winnt::{DLL_PROCESS_ATTACH, DLL_PROCESS_DETACH}; +use winapi::shared::minwindef::{BOOL, DWORD, HINSTANCE, LPVOID, TRUE, FALSE}; +use winapi::um::errhandlingapi::GetLastError; +use winapi::um::handleapi::{DuplicateHandle, CloseHandle}; +use winapi::um::processthreadsapi::{GetCurrentProcess, GetCurrentThread}; +use winapi::um::synchapi::WaitForSingleObject; +use winapi::um::winnt::{DLL_PROCESS_ATTACH, DLL_PROCESS_DETACH, SYNCHRONIZE}; use crate::configuration::Configuration; use crate::helpers::hash_endpoint; @@ -88,6 +95,18 @@ fn init_logger() { .init(); } +struct ThreadHandle(*mut c_void); + +impl ThreadHandle { + pub unsafe fn wait_and_close(self, ms: u32) { + WaitForSingleObject(self.0, ms); + CloseHandle(self.0); + } +} + +unsafe impl Send for ThreadHandle {} +unsafe impl Sync for ThreadHandle {} + #[no_mangle] #[allow(non_snake_case, unused_variables)] extern "system" fn DllMain(dll_module: HINSTANCE, call_reason: DWORD, reserved: LPVOID) -> BOOL { @@ -95,13 +114,39 @@ extern "system" fn DllMain(dll_module: HINSTANCE, call_reason: DWORD, reserved: DLL_PROCESS_ATTACH => { init_logger(); - if let Err(err) = hook_init() { - error!("{:#}", err); - } + let (cur_thread, result) = unsafe { + let mut cur_thread = ptr::null_mut(); + let result = DuplicateHandle( + GetCurrentProcess(), + GetCurrentThread(), + GetCurrentProcess(), + &mut cur_thread, + SYNCHRONIZE, + FALSE, + 0 + ); + + if result == 0 { + warn!("Failed to get current thread handle, error code: {}", GetLastError()); + } + + (ThreadHandle(cur_thread), result) + }; + + thread::spawn(move || { + if result != 0 { + unsafe { cur_thread.wait_and_close(100) }; + } + + if let Err(err) = hook_init() { + error!("Failed to initialize hook: {:#}", err); + } + }); } DLL_PROCESS_DETACH => { if let Err(err) = hook_release() { error!("{:#}", err); + return FALSE; } } _ => {}