From 38bf6c88221fa724efe2accea3d61033a7ccab37 Mon Sep 17 00:00:00 2001
From: Silent <zdanio95@gmail.com>
Date: Mon, 2 Sep 2019 21:48:51 +0200
Subject: [PATCH] [Kernel] Make XexUnloadImage fully release the image

Previously XexUnloadImage did not cleanup the image fully,
and if XexLoadImage was to be called again on the same module,
it was not initialized fully, leading to a crash when using it.
---
 src/xenia/kernel/kernel_state.cc              | 33 +++++++++++++++++++
 src/xenia/kernel/kernel_state.h               |  2 ++
 src/xenia/kernel/xboxkrnl/xboxkrnl_modules.cc |  3 +-
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/xenia/kernel/kernel_state.cc b/src/xenia/kernel/kernel_state.cc
index 83f10bbdd..9366c1160 100644
--- a/src/xenia/kernel/kernel_state.cc
+++ b/src/xenia/kernel/kernel_state.cc
@@ -405,6 +405,39 @@ object_ref<UserModule> KernelState::LoadUserModule(const char* raw_name,
   return module;
 }
 
+void KernelState::UnloadUserModule(const object_ref<UserModule>& module,
+                                   bool call_entry) {
+  auto global_lock = global_critical_region_.Acquire();
+
+  if (module->is_dll_module() && module->entry_point() && call_entry) {
+    // Call DllMain(DLL_PROCESS_DETACH):
+    // https://msdn.microsoft.com/en-us/library/windows/desktop/ms682583%28v=vs.85%29.aspx
+    uint64_t args[] = {
+        module->handle(),
+        0,  // DLL_PROCESS_DETACH
+        0,  // 0 for now, assume XexUnloadImage is like FreeLibrary
+    };
+    auto thread_state = XThread::GetCurrentThread()->thread_state();
+    processor()->Execute(thread_state, module->entry_point(), args,
+                         xe::countof(args));
+  }
+
+  auto iter = std::find_if(
+      user_modules_.begin(), user_modules_.end(),
+      [&module](const auto& e) { return e->path() == module->path(); });
+  assert_true(iter != user_modules_.end());  // Unloading an unregistered module
+                                             // is probably really bad
+  user_modules_.erase(iter);
+
+  // Ensure this module was not somehow registered twice
+  assert_true(std::find_if(user_modules_.begin(), user_modules_.end(),
+                           [&module](const auto& e) {
+                             return e->path() == module->path();
+                           }) == user_modules_.end());
+
+  object_table()->ReleaseHandle(module->handle());
+}
+
 void KernelState::TerminateTitle() {
   XELOGD("KernelState::TerminateTitle");
   auto global_lock = global_critical_region_.Acquire();
diff --git a/src/xenia/kernel/kernel_state.h b/src/xenia/kernel/kernel_state.h
index 19fbc5825..95ce08fb9 100644
--- a/src/xenia/kernel/kernel_state.h
+++ b/src/xenia/kernel/kernel_state.h
@@ -131,6 +131,8 @@ class KernelState {
   void SetExecutableModule(object_ref<UserModule> module);
   object_ref<UserModule> LoadUserModule(const char* name,
                                         bool call_entry = true);
+  void UnloadUserModule(const object_ref<UserModule>& module,
+                        bool call_entry = true);
 
   object_ref<KernelModule> GetKernelModule(const char* name);
   template <typename T>
diff --git a/src/xenia/kernel/xboxkrnl/xboxkrnl_modules.cc b/src/xenia/kernel/xboxkrnl/xboxkrnl_modules.cc
index f6d93a9f2..5c5ed2177 100644
--- a/src/xenia/kernel/xboxkrnl/xboxkrnl_modules.cc
+++ b/src/xenia/kernel/xboxkrnl/xboxkrnl_modules.cc
@@ -129,7 +129,8 @@ dword_result_t XexUnloadImage(lpvoid_t hmodule) {
     if (--ldr_data->load_count == 0) {
       // No more references, free it.
       module->Release();
-      kernel_state()->object_table()->RemoveHandle(module->handle());
+      kernel_state()->UnloadUserModule(object_ref<UserModule>(
+          reinterpret_cast<UserModule*>(module.release())));
     }
   }