Merge pull request #2047.

Fix sporadic DNT header test failures by having Badger's INITIALIZED flag account for listener initialization and then checking INITIALIZED before performing the test.
EFForg · Sep 7, 2018 · 1dfb0f8 · 1dfb0f8
2 parents d308e64 + 79452a6
commit 1dfb0f8
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 29 deletions.
diff --git a/src/js/background.js b/src/js/background.js
@@ -46,13 +46,19 @@ function Badger() {
   });
 
   self.storage = new pbStorage.BadgerPen(function(thisStorage) {
-    if (self.INITIALIZED) { return; }
+    if (self.INITIALIZED) {
+      return;
+    }
+
     self.heuristicBlocking = new HeuristicBlocking.HeuristicBlocker(thisStorage);
     self.updateTabList();
     self.initializeDefaultSettings();
+
     try {
       self.runMigrations();
     } finally {
+      // TODO "await" to set INITIALIZED until both below async functions resolve?
+      // see TODO in qunit_config.js
       self.loadFirstRunSeedData();
       self.initializeYellowlist();
       self.initializeDNT();
@@ -74,7 +80,15 @@ function Badger() {
       }
     });
 
-    // TODO: register all privacy badger listeners here in the storage callback
+    // start all the listeners
+    incognito.startListeners();
+    webrequest.startListeners();
+    HeuristicBlocking.startListeners();
+    FirefoxAndroid.startListeners();
+    startBackgroundListeners();
+
+    console.log("Privacy Badger is ready to rock!");
+    console.log("Set DEBUG=1 to view console messages.");
 
     self.INITIALIZED = true;
   });
@@ -856,21 +870,4 @@ function startBackgroundListeners() {
   }
 }
 
-/**
- * lets get this party started
- */
-console.log('Loading badgers into the pen.');
 var badger = window.badger = new Badger();
-
-/**
-* Start all the listeners
-*/
-incognito.startListeners();
-webrequest.startListeners();
-HeuristicBlocking.startListeners();
-FirefoxAndroid.startListeners();
-startBackgroundListeners();
-
-// TODO move listeners and this message behind INITIALIZED
-console.log('Privacy badger is ready to rock!');
-console.log('Set DEBUG=1 to view console messages.');
diff --git a/src/js/heuristicblocking.js b/src/js/heuristicblocking.js
@@ -515,11 +515,7 @@ function startListeners() {
    * Adds heuristicBlockingAccounting as listened to onBeforeSendHeaders request
    */
   chrome.webRequest.onBeforeSendHeaders.addListener(function(details) {
-    if (badger) {
-      return badger.heuristicBlocking.heuristicBlockingAccounting(details);
-    } else {
-      return {};
-    }
+    return badger.heuristicBlocking.heuristicBlockingAccounting(details);
   }, {urls: ["<all_urls>"]}, ["requestHeaders"]);
 
   /**
@@ -534,11 +530,7 @@ function startListeners() {
       }
     }
     if (hasSetCookie) {
-      if (badger) {
-        return badger.heuristicBlocking.heuristicBlockingAccounting(details);
-      } else {
-        return {};
-      }
+      return badger.heuristicBlocking.heuristicBlockingAccounting(details);
     }
   },
   {urls: ["<all_urls>"]}, ["responseHeaders"]);

diff --git a/tests/selenium/dnt_test.py b/tests/selenium/dnt_test.py
@@ -232,6 +232,11 @@ def test_should_not_record_nontracking_domains(self):
 
     def test_first_party_dnt_header(self):
         TEST_URL = "https://httpbin.org/get"
+
+        self.load_url(self.bg_url)
+        # wait until DNT-injecting webRequest listeners have been registered
+        self.wait_for_script("return badger.INITIALIZED")
+
         headers = retry_until(partial(self.get_first_party_headers, TEST_URL),
                               times=8)
         self.assertTrue(headers is not None, "It seems we failed to get DNT headers")