From bb5fd73b9eda75da67376bdb8328a5b2ae319b5b Mon Sep 17 00:00:00 2001
From: "Dr. Chat" <arkolbed@gmail.com>
Date: Sat, 26 Dec 2015 16:37:43 -0600
Subject: [PATCH] Register user-modules with the kernel in UserModule::Restore

---
 src/xenia/kernel/kernel_state.cc | 25 +++++++++++++++++++++++++
 src/xenia/kernel/kernel_state.h  |  2 ++
 src/xenia/kernel/user_module.cc  |  9 +++++++--
 src/xenia/kernel/xmodule.cc      |  6 +++++-
 4 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/src/xenia/kernel/kernel_state.cc b/src/xenia/kernel/kernel_state.cc
index 420fe12d6..c2d559194 100644
--- a/src/xenia/kernel/kernel_state.cc
+++ b/src/xenia/kernel/kernel_state.cc
@@ -159,6 +159,31 @@ void KernelState::RegisterModule(XModule* module) {}
 
 void KernelState::UnregisterModule(XModule* module) {}
 
+bool KernelState::RegisterUserModule(object_ref<UserModule> module) {
+  auto lock = global_critical_region_.Acquire();
+
+  for (auto user_module : user_modules_) {
+    if (user_module->path() == module->path()) {
+      // Already loaded.
+      return false;
+    }
+  }
+
+  user_modules_.push_back(module);
+  return true;
+}
+
+void KernelState::UnregisterUserModule(UserModule* module) {
+  auto lock = global_critical_region_.Acquire();
+
+  for (auto it = user_modules_.begin(); it != user_modules_.end(); it++) {
+    if ((*it)->path() == module->path()) {
+      user_modules_.erase(it);
+      return;
+    }
+  }
+}
+
 bool KernelState::IsKernelModule(const char* name) {
   if (!name) {
     // Executing module isn't a kernel module.
diff --git a/src/xenia/kernel/kernel_state.h b/src/xenia/kernel/kernel_state.h
index aa59cba6b..053e52851 100644
--- a/src/xenia/kernel/kernel_state.h
+++ b/src/xenia/kernel/kernel_state.h
@@ -120,6 +120,8 @@ class KernelState {
 
   void RegisterModule(XModule* module);
   void UnregisterModule(XModule* module);
+  bool RegisterUserModule(object_ref<UserModule> module);
+  void UnregisterUserModule(UserModule* module);
   bool IsKernelModule(const char* name);
   object_ref<XModule> GetModule(const char* name, bool user_only = false);
 
diff --git a/src/xenia/kernel/user_module.cc b/src/xenia/kernel/user_module.cc
index ddbc627d9..ba38bf551 100644
--- a/src/xenia/kernel/user_module.cc
+++ b/src/xenia/kernel/user_module.cc
@@ -325,14 +325,19 @@ object_ref<UserModule> UserModule::Restore(KernelState* kernel_state,
   // XModule::Save took care of this earlier...
   // TODO: Find a nicer way to represent that here.
   if (!module->RestoreObject(stream)) {
-    return false;
+    return nullptr;
   }
 
   auto result = module->LoadFromFile(path);
   if (XFAILED(result)) {
     XELOGD("UserModule::Restore LoadFromFile(%s) FAILED - code %.8X",
            path.c_str(), result);
-    return false;
+    return nullptr;
+  }
+
+  if (!kernel_state->RegisterUserModule(retain_object(module))) {
+    // Already loaded?
+    assert_always();
   }
 
   return object_ref<UserModule>(module);
diff --git a/src/xenia/kernel/xmodule.cc b/src/xenia/kernel/xmodule.cc
index 86bbd78d1..9b567fb80 100644
--- a/src/xenia/kernel/xmodule.cc
+++ b/src/xenia/kernel/xmodule.cc
@@ -114,7 +114,7 @@ bool XModule::Save(ByteStream* stream) {
 object_ref<XModule> XModule::Restore(KernelState* kernel_state,
                                      ByteStream* stream) {
   if (stream->Read<uint32_t>() != 'XMOD') {
-    return false;
+    return nullptr;
   }
 
   auto path = stream->Read<std::string>();
@@ -123,6 +123,10 @@ object_ref<XModule> XModule::Restore(KernelState* kernel_state,
   // Can only save user modules at the moment, so just redirect.
   // TODO: Find a way to call RestoreObject here before UserModule::Restore.
   auto module = UserModule::Restore(kernel_state, stream, path);
+  if (!module) {
+    return nullptr;
+  }
+
   XELOGD("XModule %.8X (%s)", module->handle(), module->path_.c_str());
 
   module->hmodule_ptr_ = hmodule_ptr;