Merge pull request #2853 from mrmr1993/pr/link-hints-only-activated-frames

Small tweaks for popup message code

2 files changed, 14 insertions, 12 deletions

diff --git a/background_scripts/main.coffee b/background_scripts/main.coffee
index 2c7b13ba..e3188a26 100644
--- a/background_scripts/main.coffee
+++ b/background_scripts/main.coffee
@@ -301,7 +301,7 @@ for icon in [ENABLED_ICON, DISABLED_ICON, PARTIAL_ICON]
 Frames =
   onConnect: (sender, port) ->
     [tabId, frameId] = [sender.tab.id, sender.frameId]
-    port.onDisconnect.addListener -> Frames.unregisterFrame {tabId, frameId}
+    port.onDisconnect.addListener -> Frames.unregisterFrame {tabId, frameId, port}
     port.postMessage handler: "registerFrameId", chromeFrameId: frameId
     (portsForTab[tabId] ?= {})[frameId] = port
 
@@ -312,11 +312,11 @@ Frames =
   registerFrame: ({tabId, frameId, port}) ->
     frameIdsForTab[tabId].push frameId unless frameId in frameIdsForTab[tabId] ?= []
 
-  unregisterFrame: ({tabId, frameId}) ->
-    # FrameId 0 is the top/main frame.  We never unregister that frame.  If the tab is closing, then we tidy
-    # up elsewhere.  If the tab is navigating to a new page, then a new top frame will be along soon.
-    # This mitigates against the unregister and register messages arriving out of order. See #2125.
-    if 0 < frameId
+  unregisterFrame: ({tabId, frameId, port}) ->
+    # Check that the port trying to unregister the frame hasn't already been replaced by a new frame
+    # registering. See #2125.
+    registeredPort = portsForTab[tabId]?[frameId]
+    if registeredPort == port or not registeredPort
       if tabId of frameIdsForTab
         frameIdsForTab[tabId] = (fId for fId in frameIdsForTab[tabId] when fId != frameId)
       if tabId of portsForTab
@@ -389,7 +389,8 @@ HintCoordinator =
 
   prepareToActivateMode: (tabId, originatingFrameId, {modeIndex, isVimiumHelpDialog}) ->
     @tabState[tabId] = {frameIds: frameIdsForTab[tabId][..], hintDescriptors: {}, originatingFrameId, modeIndex}
-    @tabState[tabId].ports = extend {}, portsForTab[tabId]
+    @tabState[tabId].ports = {}
+    frameIdsForTab[tabId].map (frameId) => @tabState[tabId].ports[frameId] = portsForTab[tabId][frameId]
     @sendMessage "getHintDescriptors", tabId, {modeIndex, isVimiumHelpDialog}
 
   # Receive hint descriptors from all frames and activate link-hints mode when we have them all.
diff --git a/pages/options.coffee b/pages/options.coffee
index 529743f4..9e95bcd3 100644
--- a/pages/options.coffee
+++ b/pages/options.coffee
@@ -277,11 +277,8 @@ initPopupPage = ->
     exclusions = null
     document.getElementById("optionsLink").setAttribute "href", chrome.runtime.getURL("pages/options.html")
 
-    # As the active URL, we choose the most recently registered URL from a frame in the tab, or the tab's own
-    # URL.
-    url = chrome.extension.getBackgroundPage().urlForTab[tab.id] || tab.url
-
-    unless chrome.extension.getBackgroundPage().portsForTab[tab.id]
+    tabPorts = chrome.extension.getBackgroundPage().portsForTab[tab.id]
+    unless tabPorts and Object.keys(tabPorts).length > 0
       # The browser has disabled Vimium on this page. Place a message explaining this into the popup.
       document.body.innerHTML = """
         <div style="width: 400px; margin: 5px;">
@@ -300,6 +297,10 @@ initPopupPage = ->
       """
       return
 
+    # As the active URL, we choose the most recently registered URL from a frame in the tab, or the tab's own
+    # URL.
+    url = chrome.extension.getBackgroundPage().urlForTab[tab.id] || tab.url
+
     updateState = ->
       rule = bgExclusions.getRule url, exclusions.readValueFromElement()
       $("state").innerHTML = "Vimium will " +