From 8ac19ef59514d60a359bb9cebc8542e6467dfba2 Mon Sep 17 00:00:00 2001
From: beerpiss <lacvtg.a1.2023@gmail.com>
Date: Sun, 26 Nov 2023 16:22:56 +0700
Subject: [PATCH] feat: all the logs are output now

---
 Cargo.lock |  2 +-
 Cargo.toml |  4 ++--
 src/lib.rs | 57 ++++++++++++++++++++++++++++++++++++++++++++++++------
 3 files changed, 54 insertions(+), 9 deletions(-)

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 <lacvtg.a1.2023@gmail.com>"]
 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;
             }
         }
         _ => {}