From a8b01049235bffa13d18a010311c16c8b9c316b4 Mon Sep 17 00:00:00 2001
From: Feng Chen <vonchenplus@gmail.com>
Date: Mon, 25 Oct 2021 18:55:20 +0800
Subject: Fix memory leak

---
 src/core/hle/kernel/k_handle_table.cpp |  1 +
 src/core/hle/kernel/kernel.cpp         | 23 +++++++++++++++++++++++
 src/core/hle/kernel/kernel.h           |  8 ++++++++
 src/core/hle/kernel/svc.cpp            |  6 ++++++
 4 files changed, 38 insertions(+)

diff --git a/src/core/hle/kernel/k_handle_table.cpp b/src/core/hle/kernel/k_handle_table.cpp
index 44d13169f..e90fc0628 100644
--- a/src/core/hle/kernel/k_handle_table.cpp
+++ b/src/core/hle/kernel/k_handle_table.cpp
@@ -56,6 +56,7 @@ bool KHandleTable::Remove(Handle handle) {
     }
 
     // Close the object.
+    kernel.UnregisterInUseObject(obj);
     obj->Close();
     return true;
 }
diff --git a/src/core/hle/kernel/kernel.cpp b/src/core/hle/kernel/kernel.cpp
index bea945301..d054a7925 100644
--- a/src/core/hle/kernel/kernel.cpp
+++ b/src/core/hle/kernel/kernel.cpp
@@ -170,6 +170,17 @@ struct KernelCore::Impl {
         // Next host thead ID to use, 0-3 IDs represent core threads, >3 represent others
         next_host_thread_id = Core::Hardware::NUM_CPU_CORES;
 
+        // Close kernel objects that were not freed on shutdown
+        {
+            std::lock_guard lk(registered_in_use_objects_lock);
+            if (registered_in_use_objects.size()) {
+                for (auto thread : registered_in_use_objects) {
+                    thread->Close();
+                }
+                registered_in_use_objects.clear();
+            }
+        }
+
         // Track kernel objects that were not freed on shutdown
         {
             std::lock_guard lk(registered_objects_lock);
@@ -714,9 +725,11 @@ struct KernelCore::Impl {
     std::unordered_set<KServerPort*> server_ports;
     std::unordered_set<KServerSession*> server_sessions;
     std::unordered_set<KAutoObject*> registered_objects;
+    std::unordered_set<KAutoObject*> registered_in_use_objects;
     std::mutex server_ports_lock;
     std::mutex server_sessions_lock;
     std::mutex registered_objects_lock;
+    std::mutex registered_in_use_objects_lock;
 
     std::unique_ptr<Core::ExclusiveMonitor> exclusive_monitor;
     std::vector<Kernel::PhysicalCore> cores;
@@ -928,6 +941,16 @@ void KernelCore::UnregisterKernelObject(KAutoObject* object) {
     impl->registered_objects.erase(object);
 }
 
+void KernelCore::RegisterInUseObject(KAutoObject* object) {
+    std::lock_guard lk(impl->registered_in_use_objects_lock);
+    impl->registered_in_use_objects.insert(object);
+}
+
+void KernelCore::UnregisterInUseObject(KAutoObject* object) {
+    std::lock_guard lk(impl->registered_in_use_objects_lock);
+    impl->registered_in_use_objects.erase(object);
+}
+
 bool KernelCore::IsValidNamedPort(NamedPortTable::const_iterator port) const {
     return port != impl->named_ports.cend();
 }
diff --git a/src/core/hle/kernel/kernel.h b/src/core/hle/kernel/kernel.h
index b6658b437..d2ceae950 100644
--- a/src/core/hle/kernel/kernel.h
+++ b/src/core/hle/kernel/kernel.h
@@ -204,6 +204,14 @@ public:
     /// destroyed during the current emulation session.
     void UnregisterKernelObject(KAutoObject* object);
 
+    /// Registers kernel objects with guest in use state, this is purely for close
+    /// after emulation has been shutdown.
+    void RegisterInUseObject(KAutoObject* object);
+
+    /// Unregisters a kernel object previously registered with RegisterInUseObject when it was
+    /// destroyed during the current emulation session.
+    void UnregisterInUseObject(KAutoObject* object);
+
     /// Determines whether or not the given port is a valid named port.
     bool IsValidNamedPort(NamedPortTable::const_iterator port) const;
 
diff --git a/src/core/hle/kernel/svc.cpp b/src/core/hle/kernel/svc.cpp
index f98f24a60..d30755b7e 100644
--- a/src/core/hle/kernel/svc.cpp
+++ b/src/core/hle/kernel/svc.cpp
@@ -427,11 +427,15 @@ static ResultCode WaitSynchronization(Core::System& system, s32* index, VAddr ha
         R_UNLESS(handle_table.GetMultipleObjects<KSynchronizationObject>(objs.data(), handles,
                                                                          num_handles),
                  ResultInvalidHandle);
+        for (const auto& obj : objs) {
+            kernel.RegisterInUseObject(obj);
+        }
     }
 
     // Ensure handles are closed when we're done.
     SCOPE_EXIT({
         for (u64 i = 0; i < num_handles; ++i) {
+            kernel.UnregisterInUseObject(objs[i]);
             objs[i]->Close();
         }
     });
@@ -1544,6 +1548,7 @@ static ResultCode StartThread(Core::System& system, Handle thread_handle) {
 
     // If we succeeded, persist a reference to the thread.
     thread->Open();
+    system.Kernel().RegisterInUseObject(thread.GetPointerUnsafe());
 
     return ResultSuccess;
 }
@@ -1559,6 +1564,7 @@ static void ExitThread(Core::System& system) {
     auto* const current_thread = system.Kernel().CurrentScheduler()->GetCurrentThread();
     system.GlobalSchedulerContext().RemoveThread(current_thread);
     current_thread->Exit();
+    system.Kernel().UnregisterInUseObject(current_thread);
 }
 
 static void ExitThread32(Core::System& system) {
-- 
cgit v1.2.3