Simplify domReady handling.

3 files changed, 23 insertions, 30 deletions

diff --git a/background_scripts/main.coffee b/background_scripts/main.coffee
index 6b845ff1..25aaef7b 100644
--- a/background_scripts/main.coffee
+++ b/background_scripts/main.coffee
@@ -63,20 +63,8 @@ handleCompletions = (sender) -> (request, port) ->
   completionHandlers[request.handler] completers[request.name], request, port
 
 chrome.runtime.onConnect.addListener (port, name) ->
-  sender = port.sender
-  senderTabId = sender.tab?.id
-  # If this is a tab we've been waiting to open, execute any "tab loaded" handlers, e.g. to restore
-  # the tab's scroll position. Wait until domReady before doing this; otherwise operations like restoring
-  # the scroll position will not be possible.
-  if (port.name == "domReady" && senderTabId != null)
-    if (tabLoadedHandlers[senderTabId])
-      toCall = tabLoadedHandlers[senderTabId]
-      # Delete first to be sure there's no circular events.
-      delete tabLoadedHandlers[senderTabId]
-      toCall.call()
-
   if (portHandlers[port.name])
-    port.onMessage.addListener portHandlers[port.name] sender, port
+    port.onMessage.addListener portHandlers[port.name] port.sender, port
 
 chrome.runtime.onMessage.addListener((request, sender, sendResponse) ->
   if (sendRequestHandlers[request.handler])
@@ -382,17 +370,20 @@ Frames =
 
     # Return our onMessage handler for this port.
     (request, port) =>
-      response = this[request.handler] {request, tabId, frameId, port}
-      port.postMessage response if response != false
+      this[request.handler] {request, tabId, frameId, port}
 
-  isEnabledForUrl: ({request, tabId}) ->
+  isEnabledForUrl: ({request, tabId, port}) ->
     urlForTab[tabId] = request.url if request.frameIsFocused
     rule = Exclusions.getRule request.url
-    # Send a response...
-    extend request,
+    port.postMessage extend request,
       isEnabledForUrl: not rule or 0 < rule.passKeys.length
       passKeys: rule?.passKeys ? ""
 
+  domReady: ({tabId, frameId}) ->
+    if frameId == 0
+      tabLoadedHandlers[tabId]?()
+      delete tabLoadedHandlers[tabId]
+
 handleFrameFocused = (request, sender) ->
   [tabId, frameId] = [sender.tab.id, sender.frameId]
   # This might be the first time we've heard from this tab.
diff --git a/content_scripts/vimium_frontend.coffee b/content_scripts/vimium_frontend.coffee
index 70ff19f5..36ca7834 100644
--- a/content_scripts/vimium_frontend.coffee
+++ b/content_scripts/vimium_frontend.coffee
@@ -1,8 +1,5 @@
 #
-# This content script takes input from its webpage and executes commands locally on behalf of the background
-# page. It must be run prior to domReady so that we perform some operations very early. We tell the
-# background page that we're in domReady and ready to accept normal commands by connectiong to a port named
-# "domReady".
+# This content script  must be run prior to domReady so that we perform some operations very early.
 #
 
 isEnabledForUrl = true
@@ -208,13 +205,9 @@ onFocus = (event) ->
 window.addEventListener "focus", onFocus
 window.addEventListener "hashchange", onFocus
 
-DomUtils.documentReady ->
+initializeOnDomReady = ->
   # Tell the background page we're in the domReady state.
-  chrome.runtime.connect({name: "domReady"}).onDisconnect.addListener ->
-    # We disable content scripts when we lose contact with the background page.
-    isEnabledForUrl = false
-    chrome.runtime.sendMessage = ->
-    window.removeEventListener "focus", onFocus
+  Frame.postMessage "domReady"
 
 Frame =
   port: null
@@ -227,10 +220,17 @@ Frame =
   init: (callback) ->
     @addEventListener "registerFrameId", Frame.registerFrameId
     @port = chrome.runtime.connect name: "frames"
-    @port.onDisconnect.addListener => @port.postMessage = ->
+
     @port.onMessage.addListener (request) =>
       handler request for handler in @listeners[request.handler]
 
+    @port.onDisconnect.addListener =>
+      # We disable content scripts when we lose contact with the background page.
+      isEnabledForUrl = false
+      chrome.runtime.sendMessage = ->
+      window.removeEventListener "focus", onFocus
+      @port.postMessage = ->
+
 handleShowHUDforDuration = ({ text, duration }) ->
   if DomUtils.isTopFrame()
     DomUtils.documentReady -> HUD.showForDuration text, duration
@@ -658,6 +658,7 @@ window.HelpDialog ?=
     if @showing then @hide() else @show html
 
 initializePreDomReady()
+DomUtils.documentReady initializeOnDomReady
 
 root = exports ? window
 root.handlerStack = handlerStack
diff --git a/tests/unit_tests/exclusion_test.coffee b/tests/unit_tests/exclusion_test.coffee
index e6a47779..33b04777 100644
--- a/tests/unit_tests/exclusion_test.coffee
+++ b/tests/unit_tests/exclusion_test.coffee
@@ -19,7 +19,8 @@ extend(global, require "../../background_scripts/exclusions.js")
 extend(global, require "../../background_scripts/commands.js")
 extend(global, require "../../background_scripts/main.js")
 
-isEnabledForUrl = (request) -> Frames.isEnabledForUrl {request, tabId: 0}
+isEnabledForUrl = (request) ->
+  Frames.isEnabledForUrl {request, tabId: 0, port: postMessage: (id) -> id}
 
 # These tests cover only the most basic aspects of excluded URLs and passKeys.
 #