Merge pull request #2097 from smblott-github/ui-component-init-sequence

Rework UI component init sequence.

2 files changed, 20 insertions, 8 deletions

diff --git a/content_scripts/ui_component.coffee b/content_scripts/ui_component.coffee
index 92640eb2..d7bdf2a1 100644
--- a/content_scripts/ui_component.coffee
+++ b/content_scripts/ui_component.coffee
@@ -45,7 +45,8 @@ class UIComponent
         # Get vimiumSecret so the iframe can determine that our message isn't the page impersonating us.
         chrome.storage.local.get "vimiumSecret", ({ vimiumSecret }) =>
           { port1, port2 } = new MessageChannel
-          port1.onmessage = (event) => @handleMessage event
+          port1.onmessage = (event) =>
+            if event?.data == "uiComponentIsReady" then @uiComponentIsReady = true else @handleMessage event
           @iframeElement.contentWindow.postMessage vimiumSecret, chrome.runtime.getURL(""), [ port2 ]
           setIframePort port1
 
@@ -54,8 +55,6 @@ class UIComponent
         @postMessage name: "frameFocused", focusFrameId: request.focusFrameId
       false # Free up the sendResponse handler.
 
-    @styleSheetGetter.use => @iframePort.use => Utils.nextTick => @uiComponentIsReady = true
-
   # Posts a message (if one is provided), then calls continuation (if provided).  The continuation is only
   # ever called *after* the message has been posted.
   postMessage: (message = null, continuation = null) ->
diff --git a/pages/ui_component_server.coffee b/pages/ui_component_server.coffee
index 8b43095b..4210a60e 100644
--- a/pages/ui_component_server.coffee
+++ b/pages/ui_component_server.coffee
@@ -1,6 +1,6 @@
 
-# Fetch the Vimium secret, register the port recieved from the parent window, and stop listening for messages
-# on the window object. vimiumSecret is accessible only within the current instantion of Vimium.  So a
+# Fetch the Vimium secret, register the port received from the parent window, and stop listening for messages
+# on the window object. vimiumSecret is accessible only within the current instance of Vimium.  So a
 # malicious host page trying to register its own port can do no better than guessing.
 registerPort = (event) ->
   chrome.storage.local.get "vimiumSecret", ({vimiumSecret: secret}) ->
@@ -15,13 +15,26 @@ UIComponentServer =
   handleMessage: null
 
   portOpen: (@ownerPagePort) ->
-    @ownerPagePort.onmessage = (event) =>
-      @handleMessage event if @handleMessage
+    @ownerPagePort.onmessage = (event) => @handleMessage? event
+    @registerIsReady()
 
   registerHandler: (@handleMessage) ->
 
   postMessage: (message) ->
-    @ownerPagePort.postMessage message if @ownerPagePort
+    @ownerPagePort?.postMessage message
+
+  # We require both that the DOM is ready and that the port has been opened before the UI component is ready.
+  # These events can happen in either order.  We count them, and notify the content script when we've seen
+  # both.
+  registerIsReady: do ->
+    uiComponentIsReadyCount =
+      if document.readyState == "loading"
+        window.addEventListener "DOMContentLoaded", -> UIComponentServer.registerIsReady()
+        0
+      else
+        1
+
+    -> @postMessage "uiComponentIsReady" if ++uiComponentIsReadyCount == 2
 
 root = exports ? window
 root.UIComponentServer = UIComponentServer