From 65187f401e01b1a3b4f1636be5fbd05c1f21ac2b Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sat, 17 Sep 2022 12:19:59 +0200
Subject: [PATCH] Fix destruction behaviour

---
 src/client/component/arxan.cpp | 19 +++---------------
 src/client/main.cpp            |  4 ++--
 src/common/utils/nt.cpp        | 36 +++++++++++++++++++++++-----------
 src/common/utils/nt.hpp        |  4 +++-
 4 files changed, 33 insertions(+), 30 deletions(-)

diff --git a/src/client/component/arxan.cpp b/src/client/component/arxan.cpp
index 155233b9..48db4b45 100644
--- a/src/client/component/arxan.cpp
+++ b/src/client/component/arxan.cpp
@@ -211,7 +211,7 @@ namespace arxan
 
 			if (NT_SUCCESS(status))
 			{
-				if (system_information_class == SystemProcessInformation)
+				if (system_information_class == SystemProcessInformation && !utils::nt::is_shutdown_in_progress())
 				{
 					bool injected_steam = false;
 					auto addr = static_cast<uint8_t*>(system_information);
@@ -429,7 +429,8 @@ namespace arxan
 
 			if (!context)
 			{
-				MessageBoxA(nullptr, utils::string::va("No frame offset for: %llX", handler_address), "Error", MB_ICONERROR);
+				MessageBoxA(nullptr, utils::string::va("No frame offset for: %llX", handler_address), "Error",
+				            MB_ICONERROR);
 				TerminateProcess(GetCurrentProcess(), 0xBAD);
 				return current_checksum;
 			}
@@ -706,20 +707,6 @@ namespace arxan
 			//restore_debug_functions();
 		}
 
-		void pre_destroy() override
-		{
-			utils::hook::copy(GetWindowTextA, this->window_text_buffer_, sizeof(this->window_text_buffer_));
-			nt_query_system_information_hook.clear();
-			nt_query_information_process_hook.clear();
-			nt_close_hook.clear();
-			create_mutex_ex_a_hook.clear();
-			create_thread_hook.clear();
-			open_process_hook.clear();
-			get_thread_context_hook.clear();
-			zw_terminate_process_hook.clear();
-			get_proc_address_hook.clear();
-		}
-
 		int priority() override
 		{
 			return 9999;
diff --git a/src/client/main.cpp b/src/client/main.cpp
index 45cc4484..8b8e3438 100644
--- a/src/client/main.cpp
+++ b/src/client/main.cpp
@@ -13,10 +13,10 @@ namespace
 {
 	std::pair<void**, void*> g_original_import{};
 
-	DECLSPEC_NORETURN void WINAPI exit_hook(const int code)
+	DECLSPEC_NORETURN void WINAPI exit_hook(const uint32_t code)
 	{
 		component_loader::pre_destroy();
-		exit(code);
+		ExitProcess(code);
 	}
 
 	std::pair<void**, void*> patch_steam_import(const std::string& func, void* function)
diff --git a/src/common/utils/nt.cpp b/src/common/utils/nt.cpp
index d57187be..1b3673b8 100644
--- a/src/common/utils/nt.cpp
+++ b/src/common/utils/nt.cpp
@@ -12,10 +12,11 @@ namespace utils::nt
 		return library::load(path.generic_string());
 	}
 
-	library library::get_by_address(void* address)
+	library library::get_by_address(const void* address)
 	{
 		HMODULE handle = nullptr;
-		GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, static_cast<LPCSTR>(address), &handle);
+		GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+		                   static_cast<LPCSTR>(address), &handle);
 		return library(handle);
 	}
 
@@ -183,18 +184,20 @@ namespace utils::nt
 
 				while (original_thunk_data->u1.AddressOfData)
 				{
-				  if(thunk_data->u1.Function == (uint64_t)target_function) {
-				    return reinterpret_cast<void**>(&thunk_data->u1.Function);
-				  }
+					if (thunk_data->u1.Function == (uint64_t)target_function)
+					{
+						return reinterpret_cast<void**>(&thunk_data->u1.Function);
+					}
 
 					const size_t ordinal_number = original_thunk_data->u1.AddressOfData & 0xFFFFFFF;
 
-					if (ordinal_number <= 0xFFFF) {
-					  if (GetProcAddress(other_module.module_, reinterpret_cast<char*>(ordinal_number)) ==
-					    target_function)
-					  {
-					    return reinterpret_cast<void**>(&thunk_data->u1.Function);
-					  }
+					if (ordinal_number <= 0xFFFF)
+					{
+						if (GetProcAddress(other_module.module_, reinterpret_cast<char*>(ordinal_number)) ==
+							target_function)
+						{
+							return reinterpret_cast<void**>(&thunk_data->u1.Function);
+						}
 					}
 
 					++original_thunk_data;
@@ -210,6 +213,17 @@ namespace utils::nt
 		return nullptr;
 	}
 
+	bool is_shutdown_in_progress()
+	{
+		static auto* shutdown_in_progress = []
+		{
+			const library ntdll("ntdll.dll");
+			return ntdll.get_proc<BOOLEAN(*)()>("RtlDllShutdownInProgress");
+		}();
+
+		return shutdown_in_progress();
+	}
+
 	void raise_hard_exception()
 	{
 		int data = false;
diff --git a/src/common/utils/nt.hpp b/src/common/utils/nt.hpp
index dce0cbf1..0c8ef8b8 100644
--- a/src/common/utils/nt.hpp
+++ b/src/common/utils/nt.hpp
@@ -23,7 +23,7 @@ namespace utils::nt
 	public:
 		static library load(const std::string& name);
 		static library load(const std::filesystem::path& path);
-		static library get_by_address(void* address);
+		static library get_by_address(const void* address);
 
 		library();
 		explicit library(const std::string& name);
@@ -165,6 +165,8 @@ namespace utils::nt
 		HANDLE handle_{InvalidHandle};
 	};
 
+	bool is_shutdown_in_progress();
+
 	__declspec(noreturn) void raise_hard_exception();
 	std::string load_resource(int id);