Wait for seed DNR rule registration on first run

To avoid a race condition when we remove tracker data
immediately following initialization in tests and Badger Sett.

Author: ghostwords

src/js/background.js

@@ -131,6 +131,8 @@ function Badger(from_qunit) {
     await pbconfigPromise;
     // also get ready to create dynamic DNR rules
     subscribeToActionMapUpdates();
+    // get ready to await initial DNR rules creation
+    let dnrRegistrationPromise = self.subscribeToDnrUpdates().catch(console.error);
     // now async load seed data
     let seedDataPromise = self.updateTrackerData().catch(console.error);
 
@@ -164,6 +166,8 @@ function Badger(from_qunit) {
       globalThis.DATA_LOAD_IN_PROGRESS = false;
     }
 
+    await dnrRegistrationPromise;
+
     log("Initialization complete");
     self.INITIALIZED = true;
     globalThis.DEBUG = false;
@@ -386,6 +390,42 @@ Badger.prototype = {
     }
   },
 
+  /**
+   * Resolves upon registration of the initial set of DNR rules on install.
+   *
+   * @returns {Promise}
+   */
+  subscribeToDnrUpdates: function () {
+    let self = this;
+
+    return new Promise((resolve) => {
+      function onDynamicRulesUpdate(stats) {
+        let sm = badger.storage.getStore('snitch_map'),
+          num_blocked = sm.keys().filter(d =>
+            sm.getItem(d).length >= constants.TRACKING_THRESHOLD).length;
+
+        // we are done if we registered a block of rules at least as big
+        // as the count of (cookie)blocked eTLD+1 domains in seed data
+        //
+        // this check is here so that we don't think we are done too early,
+        // when we register a smaller set of rules before we get to seed rules
+        //
+        // this assumes seed rules get generated in a single round
+        if (stats.numAdded >= num_blocked) {
+          return resolve();
+        }
+
+        dnrUtils.updateDynamicRules.subscribeToNextUpdate(onDynamicRulesUpdate);
+      }
+
+      if (!self.isFirstRun) {
+        return resolve();
+      }
+
+      dnrUtils.updateDynamicRules.subscribeToNextUpdate(onDynamicRulesUpdate);
+    });
+  },
+
   /**
    * Loads seed dataset.
    *

src/lib/dnr/utils.js

@@ -349,7 +349,8 @@ let updateSessionAllowRules = utils.debounce(async function (tempAllowlist) {
  * Debounced version of chrome.declarativeNetRequest.updateDynamicRules()
  */
 let updateDynamicRules = (function () {
-  let queue = [];
+  let queue = [],
+    subscribers = [];
 
   let _update = utils.debounce(function () {
     let opts = {
@@ -367,19 +368,35 @@ let updateDynamicRules = (function () {
     }
     queue = [];
 
-    log("[DNR] updateDynamicRules: addRules=%s, removeRuleIds=%s",
-      opts.addRules.length, opts.removeRuleIds.length);
+    let num_added = opts.addRules.length,
+      num_removed = opts.removeRuleIds.length;
 
-    chrome.declarativeNetRequest.updateDynamicRules(opts);
+    log("[DNR] updateDynamicRules: addRules=%s, removeRuleIds=%s",
+      num_added, num_removed);
+
+    chrome.declarativeNetRequest.updateDynamicRules(opts, function () {
+      // notify all subscribers
+      let fns = subscribers.slice(0);
+      // avoid an infinite loop in case of immediate re-subscription
+      subscribers = [];
+      for (let fn of fns) {
+        fn({ numAdded: num_added, numRemoved: num_removed });
+      }
+    });
   }, 100);
 
   let ret = function (opts) {
     queue.push(opts);
     _update();
   };
 
+  ret.subscribeToNextUpdate = function (fn) {
+    subscribers.push(fn);
+  };
+
   ret.clearQueue = function () {
     queue = [];
+    subscribers = [];
   };
 
   return ret;