From 37d5a68bbb9558ee11d06f7fa8c9f1af7b2cb1f4 Mon Sep 17 00:00:00 2001
From: Dmitry Vorotilin <d.vorotilin@gmail.com>
Date: Tue, 4 Dec 2018 17:36:21 +0300
Subject: [PATCH] Fix setting execution context

---
 .rspec                                     |  2 ++
 lib/capybara/cuprite/browser/page.rb       | 35 ++++++++++++++++++----
 lib/capybara/cuprite/browser/web_socket.rb |  2 +-
 3 files changed, 33 insertions(+), 6 deletions(-)
 create mode 100644 .rspec

diff --git a/.rspec b/.rspec
new file mode 100644
index 0000000..409412e
--- /dev/null
+++ b/.rspec
@@ -0,0 +1,2 @@
+--color
+--format=doc
diff --git a/lib/capybara/cuprite/browser/page.rb b/lib/capybara/cuprite/browser/page.rb
index 601f8aa..ca43099 100644
--- a/lib/capybara/cuprite/browser/page.rb
+++ b/lib/capybara/cuprite/browser/page.rb
@@ -25,14 +25,16 @@ class Browser
     class Page
       extend Forwardable
 
-      delegate [:command, :wait, :subscribe] => :@client
+      delegate [:wait, :subscribe] => :@client
       delegate targets: :@browser
 
-      attr_reader :target_id, :execution_context_id
+      attr_reader :target_id
 
       def initialize(target_id, browser, logger)
         @target_id = target_id
         @browser, @logger = browser, logger
+        @query_root_node = false
+        @execution_context_mutex = Mutex.new
 
         begin
           @session_id = @browser.command("Target.attachToTarget", targetId: @target_id)["sessionId"]
@@ -84,6 +86,22 @@ def evaluate(node, expr)
         )).dig("result", "value")
       end
 
+      def execution_context_id
+        @execution_context_mutex.synchronize { @execution_context_id }
+      end
+
+      def command(*args)
+        if @query_root_node
+          begin
+            @client.command("DOM.getDocument", depth: 0)["root"]
+          ensure
+            @query_root_node = false
+          end
+        end
+
+        @client.command(*args)
+      end
+
       private
 
       def read(filename)
@@ -96,12 +114,20 @@ def subscribe_events
         end
         subscribe("Runtime.executionContextCreated") do |params|
           @execution_context_id ||= params.dig("context", "id")
+          @execution_context_mutex.unlock if @execution_context_mutex.locked?
         end
         subscribe("Runtime.executionContextDestroyed") do |params|
           if @execution_context_id == params["executionContextId"]
+            @execution_context_mutex.lock
             @execution_context_id = nil
           end
         end
+        subscribe("Runtime.executionContextsCleared") do
+          # If we didn't have time to set context id at the beginning we have
+          # to set lock and release it when we set something.
+          @execution_context_id = nil
+          @execution_context_mutex.lock
+        end
         subscribe("Page.windowOpen") { targets.refresh }
         subscribe("Page.frameStartedLoading") do |params|
           # Remember the first frame started loading since it's the main one
@@ -112,9 +138,8 @@ def subscribe_events
           # It returns node with nodeId 1 and nodeType 9 from which descend the
           # tree and we save it in a variable because if we call that again root
           # node will change the id and all subsequent nodes have to change id too.
-          if params["frameId"] == @frame_id
-            command("DOM.getDocument", depth: 0)["root"]
-          end
+          # `command` is not allowed in the block as it will deadlock the process.
+          @query_root_node = true if params["frameId"] == @frame_id
         end
       end
 
diff --git a/lib/capybara/cuprite/browser/web_socket.rb b/lib/capybara/cuprite/browser/web_socket.rb
index 1c0485b..78f434f 100644
--- a/lib/capybara/cuprite/browser/web_socket.rb
+++ b/lib/capybara/cuprite/browser/web_socket.rb
@@ -55,7 +55,7 @@ def on_message(event)
         data = JSON.parse(event.data)
 
         if block = @subscribed[data["method"]]
-          Thread.new { block.call(data["params"]) }
+          Thread.new { block.call(data["params"]) }.join
         end
 
         @messages << data