WindowServer: Close all menus belonging to a client when it disconnects

Previously we would be left with a menu stack containing nulled-out
WeakPtr's to menus in the now-disconnected clients.

This was tripping up an assertion when clicking anywhere after shutting
down a program while it had a menu open.

3 files changed, 12 insertions, 0 deletions

diff --git a/Servers/WindowServer/WSClientConnection.cpp b/Servers/WindowServer/WSClientConnection.cpp
index 78e6003174..98793c379e 100644
--- a/Servers/WindowServer/WSClientConnection.cpp
+++ b/Servers/WindowServer/WSClientConnection.cpp
@@ -49,6 +49,7 @@ WSClientConnection::WSClientConnection(CLocalSocket& client_socket, int client_i
 
 WSClientConnection::~WSClientConnection()
 {
+    WSWindowManager::the().menu_manager().close_all_menus_from_client({}, *this);
     auto windows = move(m_windows);
 }
 
diff --git a/Servers/WindowServer/WSMenuManager.cpp b/Servers/WindowServer/WSMenuManager.cpp
index 5c9088837a..f40b571527 100644
--- a/Servers/WindowServer/WSMenuManager.cpp
+++ b/Servers/WindowServer/WSMenuManager.cpp
@@ -171,6 +171,15 @@ void WSMenuManager::set_needs_window_resize()
     m_needs_window_resize = true;
 }
 
+void WSMenuManager::close_all_menus_from_client(Badge<WSClientConnection>, WSClientConnection & client)
+{
+    if (m_open_menu_stack.is_empty())
+        return;
+    if (m_open_menu_stack.first()->client() != &client)
+        return;
+    close_everyone();
+}
+
 void WSMenuManager::close_everyone()
 {
     for (auto& menu : m_open_menu_stack) {
diff --git a/Servers/WindowServer/WSMenuManager.h b/Servers/WindowServer/WSMenuManager.h
index a5a76cf42e..3fb68f133e 100644
--- a/Servers/WindowServer/WSMenuManager.h
+++ b/Servers/WindowServer/WSMenuManager.h
@@ -34,6 +34,8 @@ public:
     void close_everyone_not_in_lineage(WSMenu&);
     void close_menu_and_descendants(WSMenu&);
 
+    void close_all_menus_from_client(Badge<WSClientConnection>, WSClientConnection&);
+
     void add_applet(WSWindow&);
     void remove_applet(WSWindow&);
     void invalidate_applet(const WSWindow&, const Rect&);