Bug 866444 - Use get/setIcon to restore tab icons and remove 'image' from xulAttributes; r=yoric

browser/base/content/browser-places.js

@@ -621,8 +621,8 @@ HistoryMenu.prototype = {
       let m = document.createElement("menuitem");
       m.setAttribute("label", menuLabel);
       let selectedTab = undoItem.tabs[undoItem.selected - 1];
-      if (selectedTab.attributes.image) {
-        let iconURL = selectedTab.attributes.image;
+      if (selectedTab.image) {
+        let iconURL = selectedTab.image;
         // don't initiate a connection just to fetch a favicon (see bug 467828)
         if (/^https?:/.test(iconURL))
           iconURL = "moz-anno:favicon:" + iconURL;

browser/components/sessionstore/src/SessionStore.jsm

@@ -248,10 +248,6 @@ let SessionStoreInternal = {
     Ci.nsISupportsWeakReference
   ]),
 
-  // xul:tab attributes to (re)store (extensions might want to hook in here);
-  // the favicon is always saved for the about:sessionrestore page
-  xulAttributes: {"image": true},
-
   // set default load state
   _loadState: STATE_STOPPED,
 
@@ -1300,7 +1296,7 @@ let SessionStoreInternal = {
       this._windows[aWindow.__SSi]._closedTabs.unshift({
         state: tabState,
         title: tabTitle,
-        image: aTab.getAttribute("image"),
+        image: tabbrowser.getIcon(aTab),
         pos: aTab._tPos
       });
       var length = this._windows[aWindow.__SSi]._closedTabs.length;
@@ -1696,11 +1692,9 @@ let SessionStoreInternal = {
   },
 
   persistTabAttribute: function ssi_persistTabAttribute(aName) {
-    if (aName in this.xulAttributes)
-      return; // this attribute is already being tracked
-
-    this.xulAttributes[aName] = true;
-    this.saveStateDelayed();
+    if (TabAttributes.persist(aName)) {
+      this.saveStateDelayed();
+    }
   },
 
   /**
@@ -1990,11 +1984,12 @@ let SessionStoreInternal = {
     else if (tabData.disallow)
       delete tabData.disallow;
 
-    tabData.attributes = {};
-    for (let name in this.xulAttributes) {
-      if (aTab.hasAttribute(name))
-        tabData.attributes[name] = aTab.getAttribute(name);
-    }
+    // Save tab attributes.
+    tabData.attributes = TabAttributes.get(aTab);
+
+    // Store the tab icon.
+    let tabbrowser = aTab.ownerDocument.defaultView.gBrowser;
+    tabData.image = tabbrowser.getIcon(aTab);
 
     if (aTab.__SS_extdata)
       tabData.extData = aTab.__SS_extdata;
@@ -2983,8 +2978,10 @@ let SessionStoreInternal = {
       else
         tabbrowser.showTab(tab);
 
-      for (let name in tabData.attributes)
-        this.xulAttributes[name] = true;
+      if ("attributes" in tabData) {
+        // Ensure that we persist tab attributes restored from previous sessions.
+        Object.keys(tabData.attributes).forEach(a => TabAttributes.persist(a));
+      }
 
       // keep the data around to prevent dataloss in case
       // a tab gets closed before it's been properly restored
@@ -3107,10 +3104,15 @@ let SessionStoreInternal = {
     for (let cap of gDocShellCapabilities(browser.docShell))
       browser.docShell["allow" + cap] = !disallow.has(cap);
 
-    for (let name in this.xulAttributes)
-      tab.removeAttribute(name);
-    for (let name in tabData.attributes)
-      tab.setAttribute(name, tabData.attributes[name]);
+    // Restore tab attributes.
+    if ("attributes" in tabData) {
+      TabAttributes.set(tab, tabData.attributes);
+    }
+
+    // Restore the tab icon.
+    if ("image" in tabData) {
+      aWindow.gBrowser.setIcon(tab, tabData.image);
+    }
 
     if (tabData.storage && browser.docShell instanceof Ci.nsIDocShell)
       SessionStorage.deserialize(browser.docShell, tabData.storage);
@@ -4622,6 +4624,52 @@ let DyingWindowCache = {
   }
 };
 
+// A set of tab attributes to persist. We will read a given list of tab
+// attributes when collecting tab data and will re-set those attributes when
+// the given tab data is restored to a new tab.
+let TabAttributes = {
+  _attrs: new Set(),
+
+  // We never want to directly read or write those attributes.
+  // 'image' should not be accessed directly but handled by using the
+  //         gBrowser.getIcon()/setIcon() methods.
+  // 'pending' is used internal by sessionstore and managed accordingly.
+  _skipAttrs: new Set(["image", "pending"]),
+
+  persist: function (name) {
+    if (this._attrs.has(name) || this._skipAttrs.has(name)) {
+      return false;
+    }
+
+    this._attrs.add(name);
+    return true;
+  },
+
+  get: function (tab) {
+    let data = {};
+
+    for (let name of this._attrs) {
+      if (tab.hasAttribute(name)) {
+        data[name] = tab.getAttribute(name);
+      }
+    }
+
+    return data;
+  },
+
+  set: function (tab, data = {}) {
+    // Clear attributes.
+    for (let name of this._attrs) {
+      tab.removeAttribute(name);
+    }
+
+    // Set attributes.
+    for (let name in data) {
+      tab.setAttribute(name, data[name]);
+    }
+  }
+};
+
 // This is used to help meter the number of restoring tabs. This is the control
 // point for telling the next tab to restore. It gets attached to each gBrowser
 // via gBrowser.addTabsProgressListener

browser/components/sessionstore/test/Makefile.in

@@ -21,6 +21,7 @@ DISABLED_XPCSHELL_TESTS = \
 
 MOCHITEST_BROWSER_FILES = \
 	head.js \
+	browser_attributes.js \
 	browser_capabilities.js \
 	browser_dying_cache.js \
 	browser_form_restore_events.js \

browser/components/sessionstore/test/browser_attributes.js

@@ -0,0 +1,72 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+function test() {
+  TestRunner.run();
+}
+
+/**
+ * This test makes sure that we correctly preserve tab attributes when storing
+ * and restoring tabs. It also ensures that we skip special attributes like
+ * 'image' and 'pending' that need to be handled differently or internally.
+ */
+
+const PREF = "browser.sessionstore.restore_on_demand";
+
+function runTests() {
+  Services.prefs.setBoolPref(PREF, true)
+  registerCleanupFunction(() => Services.prefs.clearUserPref(PREF));
+
+  // Add a new tab with a nice icon.
+  let tab = gBrowser.addTab("about:robots");
+  yield whenBrowserLoaded(tab.linkedBrowser);
+
+  // Check that the tab has an 'image' attribute.
+  ok(tab.hasAttribute("image"), "tab.image exists");
+
+  // Make sure we do not persist 'image' attributes.
+  ss.persistTabAttribute("image");
+  let {attributes} = JSON.parse(ss.getTabState(tab));
+  ok(!("image" in attributes), "'image' attribute not saved");
+  ok(!("custom" in attributes), "'custom' attribute not saved");
+
+  // Test persisting a custom attribute.
+  tab.setAttribute("custom", "foobar");
+  ss.persistTabAttribute("custom");
+
+  let {attributes} = JSON.parse(ss.getTabState(tab));
+  is(attributes.custom, "foobar", "'custom' attribute is correct");
+
+  // Make sure we're backwards compatible and restore old 'image' attributes.
+  let state = {
+    entries: [{url: "about:mozilla"}],
+    attributes: {custom: "foobaz", image: gBrowser.getIcon(tab)}
+  };
+
+  // Prepare a pending tab waiting to be restored.
+  whenTabRestoring(tab);
+  yield ss.setTabState(tab, JSON.stringify(state));
+
+  ok(tab.hasAttribute("pending"), "tab is pending");
+  is(gBrowser.getIcon(tab), state.attributes.image, "tab has correct icon");
+
+  // Let the pending tab load.
+  gBrowser.selectedTab = tab;
+  yield whenBrowserLoaded(tab.linkedBrowser);
+
+  // Ensure no 'image' or 'pending' attributes are stored.
+  let {attributes} = JSON.parse(ss.getTabState(tab));
+  ok(!("image" in attributes), "'image' attribute not saved");
+  ok(!("pending" in attributes), "'pending' attribute not saved");
+  is(attributes.custom, "foobaz", "'custom' attribute is correct");
+
+  // Clean up.
+  gBrowser.removeTab(tab);
+}
+
+function whenTabRestoring(tab) {
+  tab.addEventListener("SSTabRestoring", function onRestoring() {
+    tab.removeEventListener("SSTabRestoring", onRestoring);
+    executeSoon(next);
+  });
+}