Use ports to track frames.

Use `onConnect()`, the `domReady` port  and `onDisconnect()` to track
the frames within a tab.

2 files changed, 26 insertions, 43 deletions

diff --git a/background_scripts/main.coffee b/background_scripts/main.coffee
index 81a694d0..2e0330ee 100644
--- a/background_scripts/main.coffee
+++ b/background_scripts/main.coffee
@@ -377,18 +377,21 @@ 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) ->
+    (request, port) =>
+      this[request.handler] request, sender, port
+
+  registerFrame: (request, sender, port) ->
+    (frameIdsForTab[sender.tab.id] ?= []).push request.frameId
+
+    [tabId, frameId, isTopFrame] = [sender.tab.id, request.frameId, request.isTopFrame]
+    port.onDisconnect.addListener ->
+      if isTopFrame
+        delete frameIdsForTab[tabId]
+      else
+        if tabId of frameIdsForTab
+          frameIdsForTab[tabId] = frameIdsForTab[tabId].filter (fId) -> fId != frameId
 
 handleFrameFocused = (request, sender) ->
   tabId = sender.tab.id
@@ -420,6 +423,7 @@ bgLog = (request, sender) ->
 # Port handler mapping
 portHandlers =
   completions: handleCompletions
+  domReady: Frames.onConnect.bind Frames
 
 sendRequestHandlers =
   runBackgroundCommand: runBackgroundCommand
@@ -428,8 +432,6 @@ sendRequestHandlers =
   openUrlInIncognito: TabOperations.openUrlInIncognito
   openUrlInCurrentTab: TabOperations.openUrlInCurrentTab
   openOptionsPageInNewTab: openOptionsPageInNewTab
-  registerFrame: registerFrame
-  unregisterFrame: unregisterFrame
   frameFocused: handleFrameFocused
   nextFrame: (request) -> BackgroundCommands.nextFrame 1, request.frameId
   copyToClipboard: copyToClipboard
@@ -457,11 +459,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'))
 
diff --git a/content_scripts/vimium_frontend.coffee b/content_scripts/vimium_frontend.coffee
index 037d01d3..c948f843 100644
--- a/content_scripts/vimium_frontend.coffee
+++ b/content_scripts/vimium_frontend.coffee
@@ -214,26 +214,14 @@ window.addEventListener "hashchange", onFocus
 #
 initializeOnDomReady = ->
   # Tell the background page we're in the dom ready state.
-  chrome.runtime.connect(name: "domReady").onDisconnect.addListener ->
+  port = chrome.runtime.connect name: "domReady"
+  port.postMessage handler: "registerFrame", frameId: frameId, isTopFrame: DomUtils.isTopFrame()
+  port.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()
-
 handleShowHUDforDuration = ({ text, duration }) ->
   if DomUtils.isTopFrame()
     DomUtils.documentReady -> HUD.showForDuration text, duration
@@ -271,12 +259,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 its 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 +305,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) ->
@@ -662,8 +650,6 @@ window.HelpDialog ?=
 
 initializePreDomReady()
 DomUtils.documentReady initializeOnDomReady
-DomUtils.documentReady registerFrame
-window.addEventListener "unload", unregisterFrame
 
 root = exports ? window
 root.handlerStack = handlerStack