Bug 1204159: Switch to existing system add-ons when they match the wanted set. r=rhelmer

Checks the set that balrog gives us against the existing sets and switches
without needing to download again.

toolkit/mozapps/extensions/internal/XPIProvider.jsm

@@ -347,6 +347,12 @@ function findMatchingStaticBlocklistItem(aAddon) {
   return null;
 }
 
+/**
+ * Converts an iterable of addon objects into a map with the add-on's ID as key.
+ */
+function addonMap(addons) {
+  return new Map([for (a of addons) [a.id, a]]);
+}
 
 /**
  * Sets permissions on a file
@@ -2802,13 +2808,48 @@ this.XPIProvider = {
       return;
     }
 
-    addonList = [for (spec of addonList) { spec, path: null, addon: null }];
+    addonList = new Map([for (spec of addonList) [spec.id, { spec, path: null, addon: null }]]);
 
-    // Bug 1204159: If this matches the current set in the profile or app locations
+    let getAddonsInLocation = (location) => {
-    // then just switch to those
+      return new Promise(resolve => {
+        XPIDatabase.getAddonsInLocation(location, resolve);
+      });
+    };
+
+    let setMatches = (wanted, existing) => {
+      if (wanted.size != existing.size)
+        return false;
+
+      for (let [id, addon] of existing) {
+        let wantedInfo = wanted.get(id);
+
+        if (!wantedInfo)
+          return false;
+        if (wantedInfo.spec.version != addon.version)
+          return false;
+      }
+
+      return true;
+    };
 
     let systemAddonLocation = XPIProvider.installLocationsByName[KEY_APP_SYSTEM_ADDONS];
 
+    // If this matches the current set in the profile location then do nothing.
+    let updatedAddons = addonMap(yield getAddonsInLocation(KEY_APP_SYSTEM_ADDONS));
+    if (setMatches(addonList, updatedAddons)) {
+      logger.info("Retaining existing updated system add-ons.");
+      return;
+    }
+
+    // If this matches the current set in the default location then reset the
+    // updated set.
+    let defaultAddons = addonMap(yield getAddonsInLocation(KEY_APP_SYSTEM_DEFAULTS));
+    if (setMatches(addonList, defaultAddons)) {
+      logger.info("Resetting system add-ons.");
+      systemAddonLocation.resetAddonSet();
+      return;
+    }
+
     // Download all the add-ons
     // Bug 1204158: If we already have some of these locally then just use those
     let downloadAddon = Task.async(function*(item) {
@@ -2820,7 +2861,7 @@ this.XPIProvider = {
         logger.error(`Failed to download system add-on ${item.spec.id}`, e);
       }
     });
-    yield Promise.all([for (item of addonList) downloadAddon(item)]);
+    yield Promise.all([for (item of addonList.values()) downloadAddon(item)]);
 
     // The download promises all resolve regardless, now check if they all
     // succeeded
@@ -2842,21 +2883,21 @@ this.XPIProvider = {
     }
 
     try {
-      if (!addonList.every(item => item.path && item.addon && validateAddon(item))) {
+      if (!Array.from(addonList.values()).every(item => item.path && item.addon && validateAddon(item))) {
         throw new Error("Rejecting updated system add-on set that either could not " +
                         "be downloaded or contained unusable add-ons.");
       }
 
       // Install into the install location
       logger.info("Installing new system add-on set");
-      yield systemAddonLocation.installAddonSet([for (item of addonList) item.addon]);
+      yield systemAddonLocation.installAddonSet([for (item of addonList.values()) item.addon]);
 
       // Bug 1204156: Switch to the new system add-ons without requiring a restart
     }
     finally {
       // Delete the temporary files
       logger.info("Deleting temporary files");
-      for (let item of addonList) {
+      for (let item of addonList.values()) {
         // If this item downloaded delete the temporary file.
         if (item.path) {
           try {
@@ -7414,13 +7455,10 @@ Object.assign(MutableDirectoryInstallLocation.prototype, {
 function SystemAddonInstallLocation(aName, aDirectory, aScope, aResetSet) {
   this._baseDir = aDirectory;
 
-  if (aResetSet) {
+  if (aResetSet)
-    this._addonSet = { schema: 1, addons: {} };
+    this.resetAddonSet();
-    this._saveAddonSet(this._addonSet);
+
-  }
+  this._addonSet = this._loadAddonSet();
-  else {
-    this._addonSet = this._loadAddonSet();
-  }
 
   this._directory = null;
   if (this._addonSet.directory) {
@@ -7525,6 +7563,13 @@ Object.assign(SystemAddonInstallLocation.prototype, {
     return true;
   },
 
+  /**
+   * Resets the add-on set so on the next startup the default set will be used.
+   */
+  resetAddonSet: function() {
+    this._saveAddonSet({ schema: 1, addons: {} });
+  },
+
   /**
    * Installs a new set of system add-ons into the location and updates the
    * add-on set in prefs. We wait to switch state until a restart.

toolkit/mozapps/extensions/internal/XPIProviderUtils.js

@@ -1125,6 +1125,18 @@ this.XPIDatabase = {
                                       makeSafe(aCallback)));
   },
 
+  /**
+   * Asynchronously get all the add-ons in a particular install location.
+   *
+   * @param  aLocation
+   *         The name of the install location
+   * @param  aCallback
+   *         A callback to pass the array of DBAddonInternals to
+   */
+  getAddonsInLocation: function XPIDB_getAddonsInLocation(aLocation, aCallback) {
+    this.getAddonList(aAddon => aAddon._installLocation.name == aLocation, aCallback);
+  },
+
   /**
    * Asynchronously gets the add-on with the specified ID that is visible.
    *

toolkit/mozapps/extensions/test/xpcshell/test_system_update.js

@@ -343,11 +343,15 @@ add_task(function* setup() {
 })
 
 function* setup_conditions(setup) {
+  do_print("Clearing existing database.");
+  Services.prefs.clearUserPref(PREF_SYSTEM_ADDON_SET);
+  distroDir.leafName = "empty";
+  startupManager(false);
+  yield promiseShutdownManager();
+
   do_print("Setting up conditions.");
   yield setup.setup();
 
-  // Blow away the cache to force a rescan of the filesystem
-  Services.prefs.clearUserPref(PREF_XPI_STATE);
   startupManager(false);
 
   // Make sure the initial state is correct
@@ -435,8 +439,23 @@ add_task(function* test_app_update_disabled() {
   yield promiseShutdownManager();
 });
 
-// Tests that a set that matches the hidden default set works
+// Tests that a set that matches the default set does nothing
 add_task(function* test_match_default() {
+  yield setup_conditions(TEST_CONDITIONS.withAppSet);
+
+  yield install_system_addons(yield build_xml([
+    { id: "system2@tests.mozilla.org", version: "2.0", path: "system2_2.xpi" },
+    { id: "system3@tests.mozilla.org", version: "2.0", path: "system3_2.xpi" }
+  ]));
+
+  // Shouldn't have installed an updated set
+  yield verify_state(TEST_CONDITIONS.withAppSet.initialState);
+
+  yield promiseShutdownManager();
+});
+
+// Tests that a set that matches the hidden default set works
+add_task(function* test_match_default_revert() {
   yield setup_conditions(TEST_CONDITIONS.withBothSets);
 
   yield install_system_addons(yield build_xml([
@@ -444,10 +463,9 @@ add_task(function* test_match_default() {
     { id: "system2@tests.mozilla.org", version: "1.0", path: "system2_1.xpi" }
   ]));
 
-  // Bug 1204159: This should revert to the default set instead of installing
+  // This should revert to the default set instead of installing new versions
-  // new versions into the updated set.
+  // into an updated set.
-  //yield verify_state([false, "1.0", "1.0", null, null, null]);
+  yield verify_state([false, "1.0", "1.0", null, null, null]);
-  yield verify_state([true, "1.0", "1.0", null, null, null]);
 
   yield promiseShutdownManager();
 });
@@ -461,12 +479,11 @@ add_task(function* test_match_current() {
     { id: "system3@tests.mozilla.org", version: "2.0", path: "system3_2.xpi" }
   ]));
 
-  // Bug 1204159: This should remain with the current set instead of creating
+  // This should remain with the current set instead of creating a new copy
-  // a new copy
+  let set = JSON.parse(Services.prefs.getCharPref(PREF_SYSTEM_ADDON_SET));
-  //let set = JSON.parse(Services.prefs.getCharPref(PREF_SYSTEM_ADDON_SET));
+  do_check_eq(set.directory, "prefilled");
-  //do_check_eq(set.directory, "prefilled");
 
-  yield verify_state([true, null, "2.0", "2.0", null, null]);
+  yield verify_state(TEST_CONDITIONS.withBothSets.initialState);
 
   yield promiseShutdownManager();
 });