From 1dda774518489ab4727c4c3bdd17113fb5d962ae Mon Sep 17 00:00:00 2001
From: George Roman <george.roman.99@gmail.com>
Date: Sun, 9 Jun 2019 22:25:41 +0300
Subject: [PATCH] Fix already borrowed error in webdriver

---
 components/script/script_thread.rs      | 28 +++++++-----
 components/script/webdriver_handlers.rs | 59 ++++++++++++-------------
 2 files changed, 46 insertions(+), 41 deletions(-)

diff --git a/components/script/script_thread.rs b/components/script/script_thread.rs
index f284e14dd0c..7adaef5d352 100644
--- a/components/script/script_thread.rs
+++ b/components/script/script_thread.rs
@@ -1833,6 +1833,22 @@ impl ScriptThread {
     }
 
     fn handle_webdriver_msg(&self, pipeline_id: PipelineId, msg: WebDriverScriptCommand) {
+        // https://github.com/servo/servo/issues/23535
+        // These two messages need different treatment since the JS script might mutate
+        // `self.documents`, which would conflict with the immutable borrow of it that
+        // occurs for the rest of the messages
+        match msg {
+            WebDriverScriptCommand::ExecuteScript(script, reply) => {
+                let window = { self.documents.borrow().find_window(pipeline_id) };
+                return webdriver_handlers::handle_execute_script(window, script, reply);
+            },
+            WebDriverScriptCommand::ExecuteAsyncScript(script, reply) => {
+                let window = { self.documents.borrow().find_window(pipeline_id) };
+                return webdriver_handlers::handle_execute_async_script(window, script, reply);
+            },
+            _ => (),
+        }
+
         let documents = self.documents.borrow();
         match msg {
             WebDriverScriptCommand::AddCookie(params, reply) => {
@@ -1841,9 +1857,6 @@ impl ScriptThread {
             WebDriverScriptCommand::DeleteCookies(reply) => {
                 webdriver_handlers::handle_delete_cookies(&*documents, pipeline_id, reply)
             },
-            WebDriverScriptCommand::ExecuteScript(script, reply) => {
-                webdriver_handlers::handle_execute_script(&*documents, pipeline_id, script, reply)
-            },
             WebDriverScriptCommand::FindElementCSS(selector, reply) => {
                 webdriver_handlers::handle_find_element_css(
                     &*documents,
@@ -1936,14 +1949,7 @@ impl ScriptThread {
             WebDriverScriptCommand::GetTitle(reply) => {
                 webdriver_handlers::handle_get_title(&*documents, pipeline_id, reply)
             },
-            WebDriverScriptCommand::ExecuteAsyncScript(script, reply) => {
-                webdriver_handlers::handle_execute_async_script(
-                    &*documents,
-                    pipeline_id,
-                    script,
-                    reply,
-                )
-            },
+            _ => (),
         }
     }
 
diff --git a/components/script/webdriver_handlers.rs b/components/script/webdriver_handlers.rs
index 9980e460f49..81f45283c9a 100644
--- a/components/script/webdriver_handlers.rs
+++ b/components/script/webdriver_handlers.rs
@@ -23,6 +23,7 @@ use crate::dom::htmliframeelement::HTMLIFrameElement;
 use crate::dom::htmlinputelement::HTMLInputElement;
 use crate::dom::htmloptionelement::HTMLOptionElement;
 use crate::dom::node::{window_from_node, Node, ShadowIncluding};
+use crate::dom::window::Window;
 use crate::script_thread::Documents;
 use cookie::Cookie;
 use euclid::{Point2D, Rect, Size2D};
@@ -87,53 +88,51 @@ pub unsafe fn jsval_to_webdriver(cx: *mut JSContext, val: HandleValue) -> WebDri
 
 #[allow(unsafe_code)]
 pub fn handle_execute_script(
-    documents: &Documents,
-    pipeline: PipelineId,
+    window: Option<DomRoot<Window>>,
     eval: String,
     reply: IpcSender<WebDriverJSResult>,
 ) {
-    let window = match documents.find_window(pipeline) {
-        Some(window) => window,
+    match window {
+        Some(window) => {
+            let result = unsafe {
+                let cx = window.get_cx();
+                rooted!(in(cx) let mut rval = UndefinedValue());
+                window
+                    .upcast::<GlobalScope>()
+                    .evaluate_js_on_global_with_result(&eval, rval.handle_mut());
+                jsval_to_webdriver(cx, rval.handle())
+            };
+
+            reply.send(result).unwrap();
+        },
         None => {
-            return reply
+            reply
                 .send(Err(WebDriverJSError::BrowsingContextNotFound))
                 .unwrap();
         },
-    };
-
-    let result = unsafe {
-        let cx = window.get_cx();
-        rooted!(in(cx) let mut rval = UndefinedValue());
-        window
-            .upcast::<GlobalScope>()
-            .evaluate_js_on_global_with_result(&eval, rval.handle_mut());
-        jsval_to_webdriver(cx, rval.handle())
-    };
-
-    reply.send(result).unwrap();
+    }
 }
 
 pub fn handle_execute_async_script(
-    documents: &Documents,
-    pipeline: PipelineId,
+    window: Option<DomRoot<Window>>,
     eval: String,
     reply: IpcSender<WebDriverJSResult>,
 ) {
-    let window = match documents.find_window(pipeline) {
-        Some(window) => window,
+    match window {
+        Some(window) => {
+            let cx = window.get_cx();
+            window.set_webdriver_script_chan(Some(reply));
+            rooted!(in(cx) let mut rval = UndefinedValue());
+            window
+                .upcast::<GlobalScope>()
+                .evaluate_js_on_global_with_result(&eval, rval.handle_mut());
+        },
         None => {
-            return reply
+            reply
                 .send(Err(WebDriverJSError::BrowsingContextNotFound))
                 .unwrap();
         },
-    };
-
-    let cx = window.get_cx();
-    window.set_webdriver_script_chan(Some(reply));
-    rooted!(in(cx) let mut rval = UndefinedValue());
-    window
-        .upcast::<GlobalScope>()
-        .evaluate_js_on_global_with_result(&eval, rval.handle_mut());
+    }
 }
 
 pub fn handle_get_browsing_context_id(