WebDriver: Defer removing closed window handles until no longer needed

WebDriver::Session::close_window may invoke Session::stop, which needs
the WebContent connection to still exist. Do not remove the window's
handle (thus destroying the connection) until it is no longer needed.

2 files changed, 15 insertions, 8 deletions

diff --git a/Userland/Services/WebDriver/Session.cpp b/Userland/Services/WebDriver/Session.cpp
index ba89a6b0f9..392edeeece 100644
--- a/Userland/Services/WebDriver/Session.cpp
+++ b/Userland/Services/WebDriver/Session.cpp
@@ -11,6 +11,7 @@
 #include "Session.h"
 #include "Client.h"
 #include <AK/JsonObject.h>
+#include <AK/ScopeGuard.h>
 #include <LibCore/LocalServer.h>
 #include <LibCore/StandardPaths.h>
 #include <LibCore/System.h>
@@ -117,13 +118,17 @@ Web::WebDriver::Response Session::stop()
 // 11.2 Close Window, https://w3c.github.io/webdriver/#dfn-close-window
 Web::WebDriver::Response Session::close_window()
 {
-    // 3. Close the current top-level browsing context.
-    TRY(web_content_connection().close_window());
-    m_windows.remove(m_current_window_handle);
+    {
+        // Defer removing the window handle from this session until after we know we are done with its connection.
+        ScopeGuard guard { [this] { m_windows.remove(m_current_window_handle); } };
 
-    // 4. If there are no more open top-level browsing contexts, then close the session.
-    if (m_windows.is_empty())
-        TRY(stop());
+        // 3. Close the current top-level browsing context.
+        TRY(web_content_connection().close_window());
+
+        // 4. If there are no more open top-level browsing contexts, then close the session.
+        if (m_windows.size() == 1)
+            TRY(stop());
+    }
 
     // 5. Return the result of running the remote end steps for the Get Window Handles command.
     return get_window_handles();
diff --git a/Userland/Services/WebDriver/Session.h b/Userland/Services/WebDriver/Session.h
index 49467af9ad..a42e70589a 100644
--- a/Userland/Services/WebDriver/Session.h
+++ b/Userland/Services/WebDriver/Session.h
@@ -36,8 +36,10 @@ public:
 
     WebContentConnection& web_content_connection() const
     {
-        auto const& current_window = m_windows.get(m_current_window_handle).value();
-        return current_window.web_content_connection;
+        auto current_window = m_windows.get(m_current_window_handle);
+        VERIFY(current_window.has_value());
+
+        return current_window->web_content_connection;
     }
 
     String const& current_window_handle() const