Bug 1274461 - Part 2: restore tabs should preserve userContextId. r=mikedeboer

browser/components/sessionstore/SessionStore.jsm

@@ -2923,12 +2923,26 @@ var SessionStoreInternal = {
     let numVisibleTabs = 0;
 
     for (var t = 0; t < newTabCount; t++) {
-      tabs.push(t < openTabCount ?
-                tabbrowser.tabs[t] :
-                tabbrowser.addTab("about:blank", {
-                  skipAnimation: true,
-                  forceNotRemote: true,
-                }));
+      // When trying to restore into existing tab, we also take the userContextId
+      // into account if present.
+      let userContextId = winData.tabs[t].userContextId;
+      let reuseExisting = t < openTabCount &&
+                          (tabbrowser.tabs[t].getAttribute("usercontextid") == (userContextId || ""));
+      let tab = reuseExisting ? tabbrowser.tabs[t] :
+                                tabbrowser.addTab("about:blank",
+                                                  {skipAnimation: true,
+                                                   forceNotRemote: true,
+                                                   userContextId});
+
+      // If we inserted a new tab because the userContextId didn't match with the
+      // open tab, even though `t < openTabCount`, we need to remove that open tab
+      // and put the newly added tab in its place.
+      if (!reuseExisting && t < openTabCount) {
+        tabbrowser.removeTab(tabbrowser.tabs[t]);
+        tabbrowser.moveTabTo(tab, t);
+      }
+
+      tabs.push(tab);
 
       if (winData.tabs[t].pinned)
         tabbrowser.pinTab(tabs[t]);

browser/components/sessionstore/test/browser_sessionStoreContainer.js

@@ -2,6 +2,8 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
+requestLongerTimeout(3);
+
 add_task(function* () {
   for (let i = 0; i < 3; ++i) {
     let tab = gBrowser.addTab("http://example.com/", { userContextId: i });
@@ -66,3 +68,70 @@ add_task(function* () {
 
   yield promiseRemoveTab(tab2);
 });
+
+add_task(function* () {
+  let win = window.openDialog(location, "_blank", "chrome,all,dialog=no");
+  yield promiseWindowLoaded(win);
+
+  // Create 4 tabs with different userContextId.
+  for (let userContextId = 1; userContextId < 5; userContextId++) {
+    let tab = win.gBrowser.addTab("http://example.com/", {userContextId});
+    yield promiseBrowserLoaded(tab.linkedBrowser);
+    yield TabStateFlusher.flush(tab.linkedBrowser);
+  }
+
+  // Move the default tab of window to the end.
+  // We want the 1st tab to have non-default userContextId, so later when we
+  // restore into win2 we can test restore into an existing tab with different
+  // userContextId.
+  win.gBrowser.moveTabTo(win.gBrowser.tabs[0], win.gBrowser.tabs.length - 1);
+
+  let winState = JSON.parse(ss.getWindowState(win));
+
+  for (let i = 0; i < 4; i++) {
+    Assert.equal(winState.windows[0].tabs[i].userContextId, i + 1,
+                 "1st Window: tabs[" + i + "].userContextId should exist.");
+  }
+
+  let win2 = window.openDialog(location, "_blank", "chrome,all,dialog=no");
+  yield promiseWindowLoaded(win2);
+
+  // Create tabs with different userContextId, but this time we create them with
+  // fewer tabs and with different order with win.
+  for (let userContextId = 3; userContextId > 0; userContextId--) {
+    let tab = win2.gBrowser.addTab("http://example.com/", {userContextId});
+    yield promiseBrowserLoaded(tab.linkedBrowser);
+    yield TabStateFlusher.flush(tab.linkedBrowser);
+  }
+
+  ss.setWindowState(win2, JSON.stringify(winState), true);
+
+  for (let i = 0; i < 4; i++) {
+    let browser = win2.gBrowser.tabs[i].linkedBrowser;
+    yield ContentTask.spawn(browser, { expectedId: i + 1 }, function* (args) {
+      Assert.equal(docShell.getOriginAttributes().userContextId,
+                   args.expectedId,
+                   "The docShell has the correct userContextId");
+
+      Assert.equal(content.document.nodePrincipal.originAttributes.userContextId,
+                   args.expectedId,
+                   "The document has the correct userContextId");
+    });
+  }
+
+  // Test the last tab, which doesn't have userContextId.
+  let browser = win2.gBrowser.tabs[4].linkedBrowser;
+  yield ContentTask.spawn(browser, { expectedId: 0 }, function* (args) {
+    Assert.equal(docShell.getOriginAttributes().userContextId,
+                 args.expectedId,
+                 "The docShell has the correct userContextId");
+
+    Assert.equal(content.document.nodePrincipal.originAttributes.userContextId,
+                 args.expectedId,
+                 "The document has the correct userContextId");
+  });
+
+  yield BrowserTestUtils.closeWindow(win);
+  yield BrowserTestUtils.closeWindow(win2);
+});
+