8 files changed, 45 insertions, 10 deletions

diff --git a/src/core/hle/service/nvdrv/core/nvmap.cpp b/src/core/hle/service/nvdrv/core/nvmap.cpp
index fbd8a74a5..a51ca5444 100644
--- a/src/core/hle/service/nvdrv/core/nvmap.cpp
+++ b/src/core/hle/service/nvdrv/core/nvmap.cpp
@@ -255,15 +255,16 @@ std::optional<NvMap::FreeInfo> NvMap::FreeHandle(Handle::Id handle, bool interna
             .address = handle_description->address,
             .size = handle_description->size,
             .was_uncached = handle_description->flags.map_uncached.Value() != 0,
+            .can_unlock = true,
         };
     } else {
         return std::nullopt;
     }
 
-    // Handle hasn't been freed from memory, set address to 0 to mark that the handle wasn't freed
+    // If the handle hasn't been freed from memory, mark that
     if (!hWeak.expired()) {
         LOG_DEBUG(Service_NVDRV, "nvmap handle: {} wasn't freed as it is still in use", handle);
-        freeInfo.address = 0;
+        freeInfo.can_unlock = false;
     }
 
     return freeInfo;
diff --git a/src/core/hle/service/nvdrv/core/nvmap.h b/src/core/hle/service/nvdrv/core/nvmap.h
index b9dd3801f..a8e573890 100644
--- a/src/core/hle/service/nvdrv/core/nvmap.h
+++ b/src/core/hle/service/nvdrv/core/nvmap.h
@@ -105,6 +105,7 @@ public:
         u64 address;       //!< Address the handle referred to before deletion
         u64 size;          //!< Page-aligned handle size
         bool was_uncached; //!< If the handle was allocated as uncached
+        bool can_unlock;   //!< If the address region is ready to be unlocked
     };
 
     explicit NvMap(Tegra::Host1x::Host1x& host1x);
diff --git a/src/core/hle/service/nvdrv/devices/nvmap.cpp b/src/core/hle/service/nvdrv/devices/nvmap.cpp
index b60679021..44388655d 100644
--- a/src/core/hle/service/nvdrv/devices/nvmap.cpp
+++ b/src/core/hle/service/nvdrv/devices/nvmap.cpp
@@ -251,10 +251,12 @@ NvResult nvmap::IocFree(const std::vector<u8>& input, std::vector<u8>& output) {
     }
 
     if (auto freeInfo{file.FreeHandle(params.handle, false)}) {
-        ASSERT(system.CurrentProcess()
-                   ->PageTable()
-                   .UnlockForDeviceAddressSpace(freeInfo->address, freeInfo->size)
-                   .IsSuccess());
+        if (freeInfo->can_unlock) {
+            ASSERT(system.CurrentProcess()
+                       ->PageTable()
+                       .UnlockForDeviceAddressSpace(freeInfo->address, freeInfo->size)
+                       .IsSuccess());
+        }
         params.address = freeInfo->address;
         params.size = static_cast<u32>(freeInfo->size);
         params.flags.raw = 0;
diff --git a/src/core/hle/service/nvflinger/buffer_queue_producer.cpp b/src/core/hle/service/nvflinger/buffer_queue_producer.cpp
index 77ddbb6ef..41ba44b21 100644
--- a/src/core/hle/service/nvflinger/buffer_queue_producer.cpp
+++ b/src/core/hle/service/nvflinger/buffer_queue_producer.cpp
@@ -742,6 +742,13 @@ Status BufferQueueProducer::Disconnect(NativeWindowApi api) {
             return Status::NoError;
         }
 
+        // HACK: We are not Android. Remove handle for items in queue, and clear queue.
+        // Allows synchronous destruction of nvmap handles.
+        for (auto& item : core->queue) {
+            nvmap.FreeHandle(item.graphic_buffer->BufferId(), true);
+        }
+        core->queue.clear();
+
         switch (api) {
         case NativeWindowApi::Egl:
         case NativeWindowApi::Cpu:
diff --git a/src/core/hle/service/nvflinger/nvflinger.cpp b/src/core/hle/service/nvflinger/nvflinger.cpp
index dad93b38e..c3af12c90 100644
--- a/src/core/hle/service/nvflinger/nvflinger.cpp
+++ b/src/core/hle/service/nvflinger/nvflinger.cpp
@@ -138,6 +138,19 @@ std::optional<u64> NVFlinger::OpenDisplay(std::string_view name) {
     return itr->GetID();
 }
 
+bool NVFlinger::CloseDisplay(u64 display_id) {
+    const auto lock_guard = Lock();
+    auto* const display = FindDisplay(display_id);
+
+    if (display == nullptr) {
+        return false;
+    }
+
+    display->Reset();
+
+    return true;
+}
+
 std::optional<u64> NVFlinger::CreateLayer(u64 display_id) {
     const auto lock_guard = Lock();
     auto* const display = FindDisplay(display_id);
diff --git a/src/core/hle/service/nvflinger/nvflinger.h b/src/core/hle/service/nvflinger/nvflinger.h
index b8191c595..460bef976 100644
--- a/src/core/hle/service/nvflinger/nvflinger.h
+++ b/src/core/hle/service/nvflinger/nvflinger.h
@@ -58,6 +58,11 @@ public:
     /// If an invalid display name is provided, then an empty optional is returned.
     [[nodiscard]] std::optional<u64> OpenDisplay(std::string_view name);
 
+    /// Closes the specified display by its ID.
+    ///
+    /// Returns false if an invalid display ID is provided.
+    [[nodiscard]] bool CloseDisplay(u64 display_id);
+
     /// Creates a layer on the specified display and returns the layer ID.
     ///
     /// If an invalid display ID is specified, then an empty optional is returned.
diff --git a/src/core/hle/service/vi/display/vi_display.h b/src/core/hle/service/vi/display/vi_display.h
index 33d5f398c..0b65a65da 100644
--- a/src/core/hle/service/vi/display/vi_display.h
+++ b/src/core/hle/service/vi/display/vi_display.h
@@ -106,6 +106,12 @@ public:
     ///
     void CloseLayer(u64 layer_id);
 
+    /// Resets the display for a new connection.
+    void Reset() {
+        layers.clear();
+        got_vsync_event = false;
+    }
+
     /// Attempts to find a layer with the given ID.
     ///
     /// @param layer_id The layer ID.
diff --git a/src/core/hle/service/vi/vi.cpp b/src/core/hle/service/vi/vi.cpp
index 9c917cacf..bb283e74e 100644
--- a/src/core/hle/service/vi/vi.cpp
+++ b/src/core/hle/service/vi/vi.cpp
@@ -324,10 +324,10 @@ private:
         IPC::RequestParser rp{ctx};
         const u64 display = rp.Pop<u64>();
 
-        LOG_WARNING(Service_VI, "(STUBBED) called. display=0x{:016X}", display);
+        const Result rc = nv_flinger.CloseDisplay(display) ? ResultSuccess : ResultUnknown;
 
         IPC::ResponseBuilder rb{ctx, 2};
-        rb.Push(ResultSuccess);
+        rb.Push(rc);
     }
 
     void CreateManagedLayer(Kernel::HLERequestContext& ctx) {
@@ -508,10 +508,10 @@ private:
         IPC::RequestParser rp{ctx};
         const u64 display_id = rp.Pop<u64>();
 
-        LOG_WARNING(Service_VI, "(STUBBED) called. display_id=0x{:016X}", display_id);
+        const Result rc = nv_flinger.CloseDisplay(display_id) ? ResultSuccess : ResultUnknown;
 
         IPC::ResponseBuilder rb{ctx, 2};
-        rb.Push(ResultSuccess);
+        rb.Push(rc);
     }
 
     // This literally does nothing internally in the actual service itself,