Merge pull request #2053 from smblott-github/use-port-for-frames

Use port for frames

4 files changed, 84 insertions, 104 deletions

diff --git a/background_scripts/main.coffee b/background_scripts/main.coffee
index 81a694d0..1a67f2b2 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.onMessage.addListener portHandlers[port.name] port.sender, port
 
 chrome.runtime.onMessage.addListener((request, sender, sendResponse) ->
   if (sendRequestHandlers[request.handler])
@@ -102,20 +90,6 @@ logMessage = do ->
 #
 getCurrentTabUrl = (request, sender) -> sender.tab.url
 
-#
-# Checks the user's preferences in local storage to determine if Vimium is enabled for the given URL, and
-# whether any keys should be passed through to the underlying page.
-# The source frame also informs us whether or not it has the focus, which allows us to track the URL of the
-# active frame.
-#
-root.isEnabledForUrl = isEnabledForUrl = (request, sender) ->
-  urlForTab[sender.tab.id] = request.url if request.frameIsFocused
-  rule = Exclusions.getRule(request.url)
-  {
-    isEnabledForUrl: not rule or rule.passKeys
-    passKeys: rule?.passKeys or ""
-  }
-
 onURLChange = (details) ->
   chrome.tabs.sendMessage details.tabId, name: "checkEnabledAfterURLChange"
 
@@ -377,33 +351,46 @@ openOptionsPageInNewTab = ->
   chrome.tabs.getSelected(null, (tab) ->
     chrome.tabs.create({ url: chrome.runtime.getURL("pages/options.html"), index: tab.index + 1 }))
 
-registerFrame = (request, sender) ->
-  (frameIdsForTab[sender.tab.id] ?= []).push request.frameId
-
-unregisterFrame = (request, sender) ->
-  # When a tab is closing, Chrome sometimes passes messages without sender.tab.  Therefore, we guard against
-  # this.
-  tabId = sender.tab?.id
-  if tabId? and frameIdsForTab[tabId]?
-    if request.tab_is_closing
-      delete frameIdsForTab[tabId]
-    else
-      frameIdsForTab[tabId] = frameIdsForTab[tabId].filter (id) -> id != request.frameId
+Frames =
+  onConnect: (sender, port) ->
+    [tabId, frameId] = [sender.tab.id, sender.frameId]
+    frameIdsForTab[tabId] ?= []
+    frameIdsForTab[tabId].push frameId unless frameId in frameIdsForTab[tabId]
+    port.postMessage handler: "registerFrameId", chromeFrameId: frameId
+
+    port.onDisconnect.addListener listener = ->
+      # Unregister the frame.  However, we never unregister the main/top frame.  If the tab is navigating to
+      # another page, then there'll be a new top frame with the same Id soon.  If the tab is closing, then
+      # we tidy up in the chrome.tabs.onRemoved listener.  This elides any dependency on the order in which
+      # events happen (e.g. on navigation, a new top frame registers before the old one unregisters).
+      if tabId of frameIdsForTab and frameId != 0
+        frameIdsForTab[tabId] = frameIdsForTab[tabId].filter (fId) -> fId != frameId
+
+    # Return our onMessage handler for this port.
+    (request, port) =>
+      this[request.handler] {request, tabId, frameId, port}
+
+  isEnabledForUrl: ({request, tabId, port}) ->
+    urlForTab[tabId] = request.url if request.frameIsFocused
+    rule = Exclusions.getRule request.url
+    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 = sender.tab.id
-  # Cycle frameIdsForTab to the focused frame.  However, also ensure that we don't inadvertently register a
-  # frame which wasn't previously registered (such as a frameset).
-  if frameIdsForTab[tabId]?
-    frameIdsForTab[tabId] = cycleToFrame frameIdsForTab[tabId], request.frameId
+  [tabId, frameId] = [sender.tab.id, sender.frameId]
+  frameIdsForTab[tabId] ?= []
+  frameIdsForTab[tabId] = cycleToFrame frameIdsForTab[tabId], frameId
   # Inform all frames that a frame has received the focus.
-  chrome.tabs.sendMessage sender.tab.id,
-    name: "frameFocused"
-    focusFrameId: request.frameId
+  chrome.tabs.sendMessage tabId, name: "frameFocused", focusFrameId: frameId
 
 # Rotate through frames to the frame count places after frameId.
 cycleToFrame = (frames, frameId, count = 0) ->
-  frames ||= []
   # We can't always track which frame chrome has focussed, but here we learn that it's frameId; so add an
   # additional offset such that we do indeed start from frameId.
   count = (count + Math.max 0, frames.indexOf frameId) % frames.length
@@ -420,6 +407,7 @@ bgLog = (request, sender) ->
 # Port handler mapping
 portHandlers =
   completions: handleCompletions
+  frames: Frames.onConnect.bind Frames
 
 sendRequestHandlers =
   runBackgroundCommand: runBackgroundCommand
@@ -428,13 +416,10 @@ sendRequestHandlers =
   openUrlInIncognito: TabOperations.openUrlInIncognito
   openUrlInCurrentTab: TabOperations.openUrlInCurrentTab
   openOptionsPageInNewTab: openOptionsPageInNewTab
-  registerFrame: registerFrame
-  unregisterFrame: unregisterFrame
   frameFocused: handleFrameFocused
   nextFrame: (request) -> BackgroundCommands.nextFrame 1, request.frameId
   copyToClipboard: copyToClipboard
   pasteFromClipboard: pasteFromClipboard
-  isEnabledForUrl: isEnabledForUrl
   selectSpecificTab: selectSpecificTab
   createMark: Marks.create.bind(Marks)
   gotoMark: Marks.goto.bind(Marks)
@@ -446,9 +431,11 @@ sendRequestHandlers =
 # We always remove chrome.storage.local/findModeRawQueryListIncognito on startup.
 chrome.storage.local.remove "findModeRawQueryListIncognito"
 
-# Remove chrome.storage.local/findModeRawQueryListIncognito if there are no remaining incognito-mode windows.
-# Since the common case is that there are none to begin with, we first check whether the key is set at all.
+# Tidy up tab caches when tabs are removed.  Also remove chrome.storage.local/findModeRawQueryListIncognito if
+# there are no remaining incognito-mode windows.  Since the common case is that there are none to begin with,
+# we first check whether the key is set at all.
 chrome.tabs.onRemoved.addListener (tabId) ->
+  delete cache[tabId] for cache in [frameIdsForTab, urlForTab]
   chrome.storage.local.get "findModeRawQueryListIncognito", (items) ->
     if items.findModeRawQueryListIncognito
       chrome.windows.getAll null, (windows) ->
@@ -457,11 +444,6 @@ chrome.tabs.onRemoved.addListener (tabId) ->
         # There are no remaining incognito-mode tabs, and findModeRawQueryListIncognito is set.
         chrome.storage.local.remove "findModeRawQueryListIncognito"
 
-# Tidy up tab caches when tabs are removed.  We cannot rely on unregisterFrame because Chrome does not always
-# provide sender.tab there.
-chrome.tabs.onRemoved.addListener (tabId) ->
-  delete cache[tabId] for cache in [ frameIdsForTab, urlForTab ]
-
 # Convenience function for development use.
 window.runTests = -> open(chrome.runtime.getURL('tests/dom_tests/dom_tests.html'))
 
@@ -502,3 +484,4 @@ chrome.runtime.onInstalled.addListener ({reason}) ->
 
 root.TabOperations = TabOperations
 root.logMessage = logMessage
+root.Frames = Frames
diff --git a/content_scripts/vimium_frontend.coffee b/content_scripts/vimium_frontend.coffee
index 037d01d3..13d1377d 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
@@ -30,10 +27,8 @@ textInputXPath = (->
   DomUtils.makeXPath(inputElements)
 )()
 
-#
-# Give this frame a unique (non-zero) id.
-#
-frameId = 1 + Math.floor(Math.random()*999999999)
+# This is set by Frame.registerFrameId(). A frameId of 0 indicates that this is the top frame in the tab.
+frameId = null
 
 # For debugging only. This logs to the console on the background page.
 bgLog = (args...) ->
@@ -138,6 +133,7 @@ window.initializeModes = ->
 # Complete initialization work that sould be done prior to DOMReady.
 #
 initializePreDomReady = ->
+  Frame.init()
   checkIfEnabledForUrl()
 
   requestHandlers =
@@ -209,30 +205,28 @@ onFocus = (event) ->
 window.addEventListener "focus", onFocus
 window.addEventListener "hashchange", onFocus
 
-#
-# Initialization tasks that must wait for the document to be ready.
-#
 initializeOnDomReady = ->
-  # Tell the background page we're in the dom ready 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
-
-registerFrame = ->
-  # Don't register frameset containers; focusing them is no use.
-  unless document.body?.tagName.toLowerCase() == "frameset"
-    chrome.runtime.sendMessage
-      handler: "registerFrame"
-      frameId: frameId
-
-# Unregister the frame if we're going to exit.
-unregisterFrame = ->
-  chrome.runtime.sendMessage
-    handler: "unregisterFrame"
-    frameId: frameId
-    tab_is_closing: DomUtils.isTopFrame()
+  # Tell the background page we're in the domReady state.
+  Frame.postMessage "domReady"
+
+Frame =
+  port: null
+  listeners: {}
+
+  addEventListener: (handler, callback) -> @listeners[handler] = callback
+  postMessage: (handler, request = {}) -> @port.postMessage extend request, {handler}
+  registerFrameId: ({chromeFrameId}) -> frameId = window.frameId = chromeFrameId
+
+  init: (callback) ->
+    @port = chrome.runtime.connect name: "frames"
+
+    @port.onMessage.addListener (request) =>
+      (@listeners[request.handler] ? this[request.handler]) request
+
+    @port.onDisconnect.addListener ->
+      # We disable the content scripts when we lose contact with the background page.
+      isEnabledForUrl = false
+      window.removeEventListener "focus", onFocus
 
 handleShowHUDforDuration = ({ text, duration }) ->
   if DomUtils.isTopFrame()
@@ -271,12 +265,12 @@ DomUtils.documentReady ->
 # Called from the backend in order to change frame focus.
 #
 focusThisFrame = (request) ->
-  if window.innerWidth < 3 or window.innerHeight < 3
-    # This frame is too small to focus. Cancel and tell the background frame to focus the next one instead.
-    # This affects sites like Google Inbox, which have many tiny iframes. See #1317.
-    # Here we're assuming that there is at least one frame large enough to focus.
-    chrome.runtime.sendMessage({ handler: "nextFrame", frameId: frameId })
-    return
+  unless request.forceFocusThisFrame
+    if window.innerWidth < 3 or window.innerHeight < 3 or document.body?.tagName.toLowerCase() == "frameset"
+      # This frame is too small to focus or it's a frameset. Cancel and tell the background page to focus the
+      # next frame instead.  This affects sites like Google Inbox, which have many tiny iframes. See #1317.
+      chrome.runtime.sendMessage handler: "nextFrame", frameId: frameId
+      return
   window.focus()
   flashFrame() if request.highlight
 
@@ -317,7 +311,7 @@ extend window,
   goToRoot: ->
     window.location.href = window.location.origin
 
-  mainFrame: -> focusThisFrame highlight: true
+  mainFrame: -> focusThisFrame highlight: true, forceFocusThisFrame: true
 
   toggleViewSource: ->
     chrome.runtime.sendMessage { handler: "getCurrentTabUrl" }, (url) ->
@@ -451,10 +445,9 @@ initializeTopFrame = (request = null) ->
 
 # Checks if Vimium should be enabled or not in this frame.  As a side effect, it also informs the background
 # page whether this frame has the focus, allowing the background page to track the active frame's URL.
-checkIfEnabledForUrl = (frameIsFocused = windowIsFocused()) ->
-  url = window.location.toString()
-  chrome.runtime.sendMessage { handler: "isEnabledForUrl", url: url, frameIsFocused: frameIsFocused }, (response) ->
-    { isEnabledForUrl, passKeys } = response
+checkIfEnabledForUrl = do ->
+  Frame.addEventListener "isEnabledForUrl", (response) ->
+    {isEnabledForUrl, passKeys, frameIsFocused} = response
     installListeners() # But only if they have not been installed already.
     # Initialize UI components. We only initialize these once we know that Vimium is enabled; see #1838.
     if isEnabledForUrl
@@ -472,14 +465,15 @@ checkIfEnabledForUrl = (frameIsFocused = windowIsFocused()) ->
           if isEnabledForUrl and not passKeys then "enabled"
           else if isEnabledForUrl then "partial"
           else "disabled"
-    null
+
+  (frameIsFocused = windowIsFocused()) ->
+    Frame.postMessage "isEnabledForUrl", {frameIsFocused, url: window.location.toString()}
 
 # When we're informed by the background page that a URL in this tab has changed, we check if we have the
 # correct enabled state (but only if this frame has the focus).
 checkEnabledAfterURLChange = ->
   checkIfEnabledForUrl() if windowIsFocused()
 
-
 window.handleEscapeForFindMode = ->
   document.body.classList.remove("vimiumFindMode")
   # removing the class does not re-color existing selections. we recreate the current selection so it reverts
@@ -662,11 +656,10 @@ window.HelpDialog ?=
 
 initializePreDomReady()
 DomUtils.documentReady initializeOnDomReady
-DomUtils.documentReady registerFrame
-window.addEventListener "unload", unregisterFrame
 
 root = exports ? window
 root.handlerStack = handlerStack
 root.frameId = frameId
+root.Frame = Frame
 root.windowIsFocused = windowIsFocused
 root.bgLog = bgLog
diff --git a/tests/dom_tests/dom_tests.coffee b/tests/dom_tests/dom_tests.coffee
index c58d75d6..53e76fc3 100644
--- a/tests/dom_tests/dom_tests.coffee
+++ b/tests/dom_tests/dom_tests.coffee
@@ -2,6 +2,7 @@
 # Install frontend event handlers.
 installListeners()
 HUD.init()
+Frame.registerFrameId chromeFrameId: 0
 
 installListener = (element, event, callback) ->
   element.addEventListener event, (-> callback.apply(this, arguments)), true
diff --git a/tests/unit_tests/exclusion_test.coffee b/tests/unit_tests/exclusion_test.coffee
index 0e4b87bc..649dfd1f 100644
--- a/tests/unit_tests/exclusion_test.coffee
+++ b/tests/unit_tests/exclusion_test.coffee
@@ -19,6 +19,9 @@ 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, port: postMessage: (request) -> request}
+
 # These tests cover only the most basic aspects of excluded URLs and passKeys.
 #
 context "Excluded URLs and pass keys",