Bug 1921353 - Ensure up-to-date flags in DNR cache a=dmeehan

Original Revision: https://phabricator.services.mozilla.com/D225005

Differential Revision: https://phabricator.services.mozilla.com/D227535

toolkit/components/extensions/ExtensionDNRStore.sys.mjs

@@ -743,6 +743,9 @@ class RulesetsStore {
         );
       }
 
+      // Note: when dataPromise resolves, this._data and this._dataPromises are
+      // set. Keep this logic in sync with the end of #initExtension().
+
       this.unloadOnShutdown(extension);
       dataPromise = this.#readData(extension);
       this._dataPromises.set(extension.uuid, dataPromise);
@@ -1027,6 +1030,7 @@ class RulesetsStore {
       async () => {
         const staticRulesets = await this.getEnabledStaticRulesets(extension);
 
+        // Note: if the outcome changes, call #setStartupFlag to update this!
         return staticRulesets.size;
       }
     );
@@ -1036,6 +1040,7 @@ class RulesetsStore {
       async () => {
         const dynamicRuleset = await this.getDynamicRules(extension);
 
+        // Note: if the outcome changes, call #setStartupFlag to update this!
         return dynamicRuleset.length;
       }
     );
@@ -1052,9 +1057,40 @@ class RulesetsStore {
         updateStaticRulesets: hasEnabledStaticRules,
         updateDynamicRuleset: hasDynamicRules,
       });
+    } else if (
+      !extension.hasShutdown &&
+      !this._dataPromises.has(extension.uuid)
+    ) {
+      // #getDataPromise() initializes _dataPromises and _data (via #readData).
+      // This may be called when the StartupCache is not populated, but if they
+      // were, then these methods are not called. All other logic expects these
+      // to be initialized when #initExtension() returns, see e.g. bug 1921353.
+      let storeData = this.#getDefaults(extension);
+      this._data.set(extension.uuid, storeData);
+      this._dataPromises.set(extension.uuid, Promise.resolve(storeData));
+      this.unloadOnShutdown(extension);
     }
   }
 
+  /**
+   * Update the flags that record the (non-)existence of static/dynamic rules.
+   * These flags are used by #initExtension.
+   * "StartupCache" here refers to the general StartupCache, NOT the one from
+   * #getCacheFilePath().
+   */
+  #setStartupFlag(extension, name, value) {
+    // The StartupCache.set method is async, but we do not wait because in
+    // practice the "async" part of it completes very quickly because the
+    // underlying StartupCache data has already been read when an extension is
+    // starting.
+    // And any writes is scheduled with an AsyncShutdown blocker, which ensures
+    // that the writes complete before the browser shuts down.
+    StartupCache.general.set(
+      [extension.id, extension.version, "dnr", name],
+      value
+    );
+  }
+
   #promiseStartupCacheLoaded() {
     if (!this._ensureCacheLoaded) {
       if (this._data.size) {
@@ -1137,7 +1173,11 @@ class RulesetsStore {
    *
    * @returns {Promise<StoreData>}
    */
-  async #readData(extension) {
+  #readData(extension) {
+    // This just forwards to the actual implementation.
+    return this._readData(extension);
+  }
+  async _readData(extension) {
     const startTime = Cu.now();
     try {
       let result;
@@ -1560,6 +1600,7 @@ class RulesetsStore {
     this._data.get(extension.uuid).updateRulesets({
       dynamicRuleset: validatedRules,
     });
+    this.#setStartupFlag(extension, "hasDynamicRules", validatedRules.length);
     await this.save(extension);
     // updateRulesetManager calls ruleManager.setDynamicRules using the
     // validated rules assigned above to this._data.
@@ -1653,6 +1694,11 @@ class RulesetsStore {
     this._data.get(extension.uuid).updateRulesets({
       staticRulesets: updatedEnabledRulesets,
     });
+    this.#setStartupFlag(
+      extension,
+      "hasEnabledStaticRules",
+      updatedEnabledRulesets.size
+    );
     await this.save(extension);
     this.updateRulesetManager(extension, {
       updateDynamicRuleset: false,

toolkit/components/extensions/test/xpcshell/test_ext_dnr_dynamic_rules.js

@@ -8,6 +8,7 @@ ChromeUtils.defineESModuleGetters(this, {
   ExtensionDNRLimits: "resource://gre/modules/ExtensionDNRLimits.sys.mjs",
   ExtensionDNRStore: "resource://gre/modules/ExtensionDNRStore.sys.mjs",
   TestUtils: "resource://testing-common/TestUtils.sys.mjs",
+  sinon: "resource://testing-common/Sinon.sys.mjs",
 });
 
 AddonTestUtils.init(this);
@@ -919,6 +920,122 @@ add_task(async function test_tabId_conditions_invalid_in_dynamic_rules() {
   });
 });
 
+// Regression test for https://bugzilla.mozilla.org/show_bug.cgi?id=1921353
+// StartupCache contains hasDynamicRules and hasEnabledStaticRules as an
+// optimization, to avoid unnecessarily trying to read and parse DNR rules from
+// disk when an extension does knowingly not have any. This had two bugs:
+// 1. When rule reading was skipped, necessary internal state was not correctly
+//    initialized, and consequently rules could not be registered or updated.
+// 2. The hasDynamicRules/hasEnabledStaticRules flags were not cleared upon
+//    update, and consequently the flag stayed in the initial state (false),
+//    and the previously stored rules did not apply after a browser restart.
+// This is a regression test to verify that these issues do not occur any more.
+//
+// See also test_enable_disabled_static_rules_after_restart in
+// test_ext_dnr_static_rules.js for the similar test with static rules.
+add_task(async function test_register_dynamic_rules_after_restart() {
+  // Through this test, we confirm that the underlying "expensive" rule data
+  // storage is accessed when needed, and skipped when no relevant rules had
+  // been detected at the previous session. Caching too much or too little in
+  // StartupCache will trigger test failures in assertStoreReadsSinceLastCall.
+  const sandboxStoreSpies = sinon.createSandbox();
+  const dnrStore = ExtensionDNRStore._getStoreForTesting();
+  const spyReadDNRStore = sandboxStoreSpies.spy(dnrStore, "_readData");
+
+  function assertStoreReadsSinceLastCall(expectedCount, description) {
+    equal(spyReadDNRStore.callCount, expectedCount, description);
+    spyReadDNRStore.resetHistory();
+  }
+
+  let { extension } = await runAsDNRExtension({
+    unloadTestAtEnd: false,
+    id: "test-dnr-restart-and-rule-registration-changes@test-extension",
+    background: () => {
+      const dnr = browser.declarativeNetRequest;
+
+      browser.runtime.onInstalled.addListener(async () => {
+        browser.test.assertDeepEq(
+          await dnr.getDynamicRules(),
+          [],
+          "No DNR rules at install time"
+        );
+        browser.test.sendMessage("onInstalled");
+      });
+      browser.test.onMessage.addListener(async msg => {
+        if (msg === "get_rule_count") {
+          browser.test.sendMessage(
+            "get_rule_count:done",
+            (await dnr.getDynamicRules()).length
+          );
+          return;
+        } else if (msg === "zero_to_one_rule") {
+          await dnr.updateDynamicRules({
+            addRules: [{ id: 1, condition: {}, action: { type: "block" } }],
+          });
+        } else if (msg === "one_to_zero_rules") {
+          await dnr.updateDynamicRules({ removeRuleIds: [1] });
+        } else {
+          browser.test.fail(`Unexpected message: ${msg}`);
+        }
+        browser.test.sendMessage(`${msg}:done`);
+      });
+      browser.runtime.onStartup.addListener(async () => {
+        browser.test.sendMessage("onStartup");
+      });
+    },
+  });
+  await extension.awaitMessage("onInstalled");
+  assertStoreReadsSinceLastCall(1, "Read once at initial startup");
+  await promiseRestartManager();
+  await extension.awaitMessage("onStartup");
+  assertStoreReadsSinceLastCall(0, "Read skipped due to hasDynamicRules=false");
+
+  // Regression test 1: before bug 1921353 was fixed, the test got stuck here
+  // because the getDynamicRules() call after a restart unexpectedly failed
+  // with "this._data.get(...) is undefined".
+  equal(
+    await callTestMessageHandler(extension, "get_rule_count"),
+    0,
+    "Started and without rules"
+  );
+
+  await callTestMessageHandler(extension, "zero_to_one_rule");
+
+  assertStoreReadsSinceLastCall(0, "No further reads before restart");
+  await promiseRestartManager();
+  await extension.awaitMessage("onStartup");
+
+  // Regression test 2: before bug 1921353 was fixed, even with a fix to the
+  // previous bug, the dynamic rules would not be read at startup due to the
+  // wrong cached hasDynamicRules flag in StartupCache.
+  assertStoreReadsSinceLastCall(1, "Read due to hasDynamicRules=true");
+  equal(
+    await callTestMessageHandler(extension, "get_rule_count"),
+    1,
+    "Started and should see the previously registered dynamic rule"
+  );
+
+  // For full coverage, also verify that when there are really 0 rules, that
+  // the initialization is skipped as expected.
+
+  await callTestMessageHandler(extension, "one_to_zero_rules");
+  await promiseRestartManager();
+  await extension.awaitMessage("onStartup");
+  assertStoreReadsSinceLastCall(0, "Read skipped because rules were cleared");
+  equal(
+    await callTestMessageHandler(extension, "get_rule_count"),
+    0,
+    "Started without rules, should not see rules"
+  );
+  // declarativeNetRequest.getDynamicRules() queries in-memory state, so we do
+  // not expect another read from disk.
+  assertStoreReadsSinceLastCall(0, "Read still skipped despite API call");
+
+  await extension.unload();
+
+  sandboxStoreSpies.restore();
+});
+
 add_task(async function test_dynamic_rules_telemetry() {
   resetTelemetryData();
 

toolkit/components/extensions/test/xpcshell/test_ext_dnr_static_rules.js

@@ -8,6 +8,7 @@ ChromeUtils.defineESModuleGetters(this, {
   ExtensionDNRLimits: "resource://gre/modules/ExtensionDNRLimits.sys.mjs",
   ExtensionDNRStore: "resource://gre/modules/ExtensionDNRStore.sys.mjs",
   TestUtils: "resource://testing-common/TestUtils.sys.mjs",
+  sinon: "resource://testing-common/Sinon.sys.mjs",
 });
 
 AddonTestUtils.init(this);
@@ -416,6 +417,104 @@ add_task(async function test_load_static_rules() {
   );
 });
 
+// As an optimization, the hasEnabledStaticRules flag in StartupCache is used
+// to avoid unnecessarily trying to read and parse DNR rules from disk when an
+// extension does knowingly not have any.
+//
+// 1. When rule reading was skipped, necessary internal state was not correctly
+//    initialized, and consequently updateStaticRules() would reject.
+// 2. The hasDynamicRules/hasEnabledStaticRules flags were not cleared upon
+//    update, and consequently the flag stayed in the initial state (false),
+//    and the previously stored rules did not apply after a browser restart.
+//
+// See also test_register_dynamic_rules_after_restart in
+// test_ext_dnr_dynamic_rules.js for the similar test with dynamic rules.
+add_task(async function test_enable_disabled_static_rules_after_restart() {
+  // Through this test, we confirm that the underlying "expensive" rule data
+  // storage is accessed when needed, and skipped when no relevant rules had
+  // been detected at the previous session. Caching too much or too little in
+  // StartupCache will trigger test failures in assertStoreReadsSinceLastCall.
+  const sandboxStoreSpies = sinon.createSandbox();
+  const dnrStore = ExtensionDNRStore._getStoreForTesting();
+  const spyReadDNRStore = sandboxStoreSpies.spy(dnrStore, "_readData");
+
+  function assertStoreReadsSinceLastCall(expectedCount, description) {
+    equal(spyReadDNRStore.callCount, expectedCount, description);
+    spyReadDNRStore.resetHistory();
+  }
+
+  const rule_resources = [
+    {
+      id: "ruleset_initially_disabled",
+      enabled: false,
+      path: "ruleset_initially_disabled.json",
+    },
+  ];
+
+  const files = {
+    "ruleset_initially_disabled.json": JSON.stringify([getDNRRule()]),
+  };
+
+  const extension = ExtensionTestUtils.loadExtension(
+    getDNRExtension({ rule_resources, files, id: "dnr@initially-disabled" })
+  );
+  await extension.startup();
+  await extension.awaitMessage("bgpage:ready");
+
+  assertStoreReadsSinceLastCall(1, "Read once at initial startup");
+  await AddonTestUtils.promiseRestartManager();
+  // Note: ordinarily, event pages do not wake up after a browser restart,
+  // unless a relevant event such as runtime.onStartup is triggered. But as
+  // noted in bug 1822735, "event pages without any event listeners" will be
+  // awakened after a restart, so we can expect the bgpage:ready message here:
+  await extension.awaitMessage("bgpage:ready");
+  assertStoreReadsSinceLastCall(
+    0,
+    "Read skipped due to hasEnabledStaticRules=false"
+  );
+
+  // Now transition from zero rulesets to one (zero_to_one_rule).
+  extension.sendMessage("updateEnabledRulesets", {
+    enableRulesetIds: ["ruleset_initially_disabled"],
+  });
+  info("Trying to enable ruleset_initially_disabled...");
+  await extension.awaitMessage("updateEnabledRulesets:done");
+  await assertDNRGetEnabledRulesets(extension, ["ruleset_initially_disabled"]);
+
+  assertStoreReadsSinceLastCall(0, "No further reads before restart");
+  await AddonTestUtils.promiseRestartManager();
+  await extension.awaitMessage("bgpage:ready");
+
+  // Regression test 2: before bug 1921353 was fixed, even with a fix to the
+  // previous bug, the static rules would not be read at startup due to the
+  // wrong cached hasEnabledStaticRules flag in StartupCache.
+  // Verify that the static rules are still enabled.
+  assertStoreReadsSinceLastCall(1, "Read due to hasEnabledStaticRules=true");
+  await assertDNRGetEnabledRulesets(extension, ["ruleset_initially_disabled"]);
+
+  // For full coverage, also verify that when all static rules are disabled,
+  // that the initialization is skipped as expected.
+
+  extension.sendMessage("updateEnabledRulesets", {
+    disableRulesetIds: ["ruleset_initially_disabled"],
+  });
+  info("Trying to disable ruleset_initially_disabled...");
+  await extension.awaitMessage("updateEnabledRulesets:done");
+  await assertDNRGetEnabledRulesets(extension, []);
+
+  await AddonTestUtils.promiseRestartManager();
+  await extension.awaitMessage("bgpage:ready");
+  assertStoreReadsSinceLastCall(0, "Read skipped because rules were disabled");
+  await assertDNRGetEnabledRulesets(extension, []);
+  // declarativeNetRequest.getEnabledStaticRulesets() queries in-memory state,
+  // so we do not expect another read from disk.
+  assertStoreReadsSinceLastCall(0, "Read still skipped despite API call");
+
+  await extension.unload();
+
+  sandboxStoreSpies.restore();
+});
+
 add_task(async function test_load_from_corrupted_data() {
   const ruleset1Data = [
     getDNRRule({