Use Chrome frameIds.

3 files changed, 34 insertions, 35 deletions

diff --git a/background_scripts/main.coffee b/background_scripts/main.coffee
index 2e0330ee..6fddb6e3 100644
--- a/background_scripts/main.coffee
+++ b/background_scripts/main.coffee
@@ -76,7 +76,7 @@ chrome.runtime.onConnect.addListener (port, name) ->
       toCall.call()
 
   if (portHandlers[port.name])
-    port.onMessage.addListener portHandlers[port.name] sender
+    port.onMessage.addListener portHandlers[port.name] sender, port
 
 chrome.runtime.onMessage.addListener((request, sender, sendResponse) ->
   if (sendRequestHandlers[request.handler])
@@ -378,31 +378,28 @@ openOptionsPageInNewTab = ->
     chrome.tabs.create({ url: chrome.runtime.getURL("pages/options.html"), index: tab.index + 1 }))
 
 Frames =
-  onConnect: (sender) ->
-    (request, port) =>
-      this[request.handler] request, sender, port
+  onConnect: (sender, port) ->
+    [tabId, frameId] = [sender.tab.id, sender.frameId]
+    (frameIdsForTab[tabId] ?= []).push frameId
+    port.postMessage name: "registerFrameId", chromeFrameId: frameId
 
-  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
+      if frameId == 0 # This is the top frame in the tab.
         delete frameIdsForTab[tabId]
       else
         if tabId of frameIdsForTab
           frameIdsForTab[tabId] = frameIdsForTab[tabId].filter (fId) -> fId != frameId
+          # Sub-frames can connect before the top frame; so we can't rely on seeing the top frame at all.
+          delete frameIdsForTab[tabId] if frameIdsForTab[tabId].length == 0
 
 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]
+  # This might be the first time we've heard from this tab.
+  frameIdsForTab[tabId] ?= []
+  # Cycle frameIdsForTab to the focused frame.
+  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) ->
diff --git a/content_scripts/vimium_frontend.coffee b/content_scripts/vimium_frontend.coffee
index c948f843..ce01695a 100644
--- a/content_scripts/vimium_frontend.coffee
+++ b/content_scripts/vimium_frontend.coffee
@@ -30,10 +30,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...) ->
@@ -170,7 +168,7 @@ initializePreDomReady = ->
 # Wrapper to install event listeners.  Syntactic sugar.
 installListener = (element, event, callback) ->
   element.addEventListener(event, ->
-    if isEnabledForUrl then callback.apply(this, arguments) else true
+    if isEnabledForUrl and frameId? then callback.apply(this, arguments) else true
   , true)
 
 #
@@ -209,18 +207,20 @@ 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.
-  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
+Frame =
+  port: null
+
+  init: ->
+    # Tell the background page we're in the domReady state.
+    @port = chrome.runtime.connect name: "domReady"
+    @port.onMessage.addListener (request) => this[request.name] request
+    @port.onDisconnect.addListener ->
+      # We disable content scripts when we lose contact with the background page.
+      isEnabledForUrl = false
+      chrome.runtime.sendMessage = ->
+      window.removeEventListener "focus", onFocus
+
+  registerFrameId: ({chromeFrameId}) -> frameId = window.frameId = chromeFrameId
 
 handleShowHUDforDuration = ({ text, duration }) ->
   if DomUtils.isTopFrame()
@@ -649,10 +649,11 @@ window.HelpDialog ?=
     if @showing then @hide() else @show html
 
 initializePreDomReady()
-DomUtils.documentReady initializeOnDomReady
+DomUtils.documentReady Frame.init.bind Frame
 
 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