From a120f8ff4de31ef4fe34306016e23a94aa2e665d Mon Sep 17 00:00:00 2001
From: Liam <byteslice@airmail.cc>
Date: Fri, 19 Jan 2024 13:52:19 -0500
Subject: [PATCH] nvservices: close map handles on session close

---
 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 ++
 3 files changed, 16 insertions(+)

diff --git a/src/core/hle/service/nvdrv/core/container.cpp b/src/core/hle/service/nvdrv/core/container.cpp
index 21ef57d277..dc1b4d5be6 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 1b59c6b15a..bc1c033c66 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 d7f6958450..b8be599ae9 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`