Rework UI component init sequence.

In testing global link hints, it seems that UI components can be opened
before they're actually ready, and Vimium hangs with a broken UI
component displayed.  This issue seems to be in 1.54 too.

It's not clear what's causing this -- and it's hard to reproduce
systematically.  However, it only seems to happen on navigation.
Therefore, it seems likely that there is an issue with the
initialization sequence.

Here, we wait for:

- the DOM content to be ready, and
- the port to be provided

before registering the UI component as ready.

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