From 34fb9081e724f02fed879cdb0f2befbda8637575 Mon Sep 17 00:00:00 2001
From: yuzubot <yuzu@yuzu-emu.org>
Date: Thu, 25 Jan 2024 01:01:11 +0000
Subject: [PATCH] Merge yuzu-emu#12759

---
 src/core/hle/kernel/k_page_table_base.cpp     | 13 +++++++++----
 src/core/hle/service/nvdrv/core/container.cpp |  1 +
 src/core/hle/service/nvdrv/core/nvmap.cpp     | 13 +++++++++++++
 src/core/hle/service/nvdrv/core/nvmap.h       |  2 ++
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/core/hle/kernel/k_page_table_base.cpp b/src/core/hle/kernel/k_page_table_base.cpp
index 3f0a39d33..1dd86fb3c 100644
--- a/src/core/hle/kernel/k_page_table_base.cpp
+++ b/src/core/hle/kernel/k_page_table_base.cpp
@@ -69,9 +69,14 @@ public:
 };
 
 template <typename AddressType>
-void InvalidateInstructionCache(KernelCore& kernel, AddressType addr, u64 size) {
+void InvalidateInstructionCache(KernelCore& kernel, KPageTableBase* table, AddressType addr,
+                                u64 size) {
     // TODO: lock the process list
     for (auto& process : kernel.GetProcessList()) {
+        if (std::addressof(process->GetPageTable().GetBasePageTable()) != table) {
+            continue;
+        }
+
         for (size_t i = 0; i < Core::Hardware::NUM_CPU_CORES; i++) {
             auto* interface = process->GetArmInterface(i);
             if (interface) {
@@ -1302,7 +1307,7 @@ Result KPageTableBase::UnmapCodeMemory(KProcessAddress dst_address, KProcessAddr
     bool reprotected_pages = false;
     SCOPE_EXIT({
         if (reprotected_pages && any_code_pages) {
-            InvalidateInstructionCache(m_kernel, dst_address, size);
+            InvalidateInstructionCache(m_kernel, this, dst_address, size);
         }
     });
 
@@ -2036,7 +2041,7 @@ Result KPageTableBase::SetProcessMemoryPermission(KProcessAddress addr, size_t s
         for (const auto& block : pg) {
             StoreDataCache(GetHeapVirtualPointer(m_kernel, block.GetAddress()), block.GetSize());
         }
-        InvalidateInstructionCache(m_kernel, addr, size);
+        InvalidateInstructionCache(m_kernel, this, addr, size);
     }
 
     R_SUCCEED();
@@ -3277,7 +3282,7 @@ Result KPageTableBase::WriteDebugMemory(KProcessAddress dst_address, KProcessAdd
     R_TRY(PerformCopy());
 
     // Invalidate the instruction cache, as this svc allows modifying executable pages.
-    InvalidateInstructionCache(m_kernel, dst_address, size);
+    InvalidateInstructionCache(m_kernel, this, dst_address, size);
 
     R_SUCCEED();
 }
diff --git a/src/core/hle/service/nvdrv/core/container.cpp b/src/core/hle/service/nvdrv/core/container.cpp
index 21ef57d27..dc1b4d5be 100644
--- a/src/core/hle/service/nvdrv/core/container.cpp
+++ b/src/core/hle/service/nvdrv/core/container.cpp
@@ -112,6 +112,7 @@ SessionId Container::OpenSession(Kernel::KProcess* process) {
 
 void Container::CloseSession(SessionId session_id) {
     std::scoped_lock lk(impl->session_guard);
+    impl->file.UnmapAllHandles(session_id);
     auto& session = impl->sessions[session_id.id];
     auto& smmu = impl->host1x.MemoryManager();
     if (session.has_preallocated_area) {
diff --git a/src/core/hle/service/nvdrv/core/nvmap.cpp b/src/core/hle/service/nvdrv/core/nvmap.cpp
index 1b59c6b15..bc1c033c6 100644
--- a/src/core/hle/service/nvdrv/core/nvmap.cpp
+++ b/src/core/hle/service/nvdrv/core/nvmap.cpp
@@ -326,4 +326,17 @@ std::optional<NvMap::FreeInfo> NvMap::FreeHandle(Handle::Id handle, bool interna
     return freeInfo;
 }
 
+void NvMap::UnmapAllHandles(NvCore::SessionId session_id) {
+    auto handles_copy = [&] {
+        std::scoped_lock lk{handles_lock};
+        return handles;
+    }();
+
+    for (auto& [id, handle] : handles_copy) {
+        if (handle->session_id.id == session_id.id) {
+            FreeHandle(id, false);
+        }
+    }
+}
+
 } // namespace Service::Nvidia::NvCore
diff --git a/src/core/hle/service/nvdrv/core/nvmap.h b/src/core/hle/service/nvdrv/core/nvmap.h
index d7f695845..b8be599ae 100644
--- a/src/core/hle/service/nvdrv/core/nvmap.h
+++ b/src/core/hle/service/nvdrv/core/nvmap.h
@@ -152,6 +152,8 @@ public:
      */
     std::optional<FreeInfo> FreeHandle(Handle::Id handle, bool internal_session);
 
+    void UnmapAllHandles(NvCore::SessionId session_id);
+
 private:
     std::list<std::shared_ptr<Handle>> unmap_queue{};
     std::mutex unmap_queue_lock{}; //!< Protects access to `unmap_queue`