1 files changed, 307 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/UrlClassifierTestUtils.sys.mjs b/toolkit/components/url-classifier/tests/UrlClassifierTestUtils.sys.mjs
new file mode 100644
index 0000000000..c69d0c24b4
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/UrlClassifierTestUtils.sys.mjs
@@ -0,0 +1,307 @@
+const ANNOTATION_TABLE_NAME = "mochitest1-track-simple";
+const ANNOTATION_TABLE_PREF = "urlclassifier.trackingAnnotationTable";
+const ANNOTATION_ENTITYLIST_TABLE_NAME = "mochitest1-trackwhite-simple";
+const ANNOTATION_ENTITYLIST_TABLE_PREF =
+  "urlclassifier.trackingAnnotationWhitelistTable";
+
+const TRACKING_TABLE_NAME = "mochitest2-track-simple";
+const TRACKING_TABLE_PREF = "urlclassifier.trackingTable";
+const ENTITYLIST_TABLE_NAME = "mochitest2-trackwhite-simple";
+const ENTITYLIST_TABLE_PREF = "urlclassifier.trackingWhitelistTable";
+
+const SOCIAL_ANNOTATION_TABLE_NAME = "mochitest3-track-simple";
+const SOCIAL_ANNOTATION_TABLE_PREF =
+  "urlclassifier.features.socialtracking.annotate.blacklistTables";
+const SOCIAL_TRACKING_TABLE_NAME = "mochitest4-track-simple";
+const SOCIAL_TRACKING_TABLE_PREF =
+  "urlclassifier.features.socialtracking.blacklistTables";
+const EMAIL_TRACKING_TABLE_NAME = "mochitest5-track-simple";
+const EMAIL_TRACKING_TABLE_PREF =
+  "urlclassifier.features.emailtracking.blocklistTables";
+
+let timer = Cc["@mozilla.org/timer;1"].createInstance(Ci.nsITimer);
+
+export var UrlClassifierTestUtils = {
+  addTestTrackers() {
+    // Add some URLs to the tracking databases
+    let annotationURL1 = "tracking.example.org/"; // only for annotations
+    let annotationURL2 = "itisatracker.org/";
+    let annotationURL3 = "trackertest.org/";
+    let annotationURL4 = "another-tracking.example.net/";
+    let annotationURL5 = "tlsresumptiontest.example.org/";
+    let annotationEntitylistedURL = "itisatrap.org/?resource=example.org";
+    let trackingURL1 = "tracking.example.com/"; // only for TP
+    let trackingURL2 = "itisatracker.org/";
+    let trackingURL3 = "trackertest.org/";
+    let entitylistedURL = "itisatrap.org/?resource=itisatracker.org";
+    let socialTrackingURL = "social-tracking.example.org/";
+    let emailTrackingURL = "email-tracking.example.org/";
+
+    let annotationUpdate =
+      "n:1000\ni:" +
+      ANNOTATION_TABLE_NAME +
+      "\nad:5\n" +
+      "a:1:32:" +
+      annotationURL1.length +
+      "\n" +
+      annotationURL1 +
+      "\n" +
+      "a:2:32:" +
+      annotationURL2.length +
+      "\n" +
+      annotationURL2 +
+      "\n" +
+      "a:3:32:" +
+      annotationURL3.length +
+      "\n" +
+      annotationURL3 +
+      "\n" +
+      "a:4:32:" +
+      annotationURL4.length +
+      "\n" +
+      annotationURL4 +
+      "\n" +
+      "a:5:32:" +
+      annotationURL5.length +
+      "\n" +
+      annotationURL5 +
+      "\n";
+    let socialAnnotationUpdate =
+      "n:1000\ni:" +
+      SOCIAL_ANNOTATION_TABLE_NAME +
+      "\nad:1\n" +
+      "a:1:32:" +
+      socialTrackingURL.length +
+      "\n" +
+      socialTrackingURL +
+      "\n";
+    let annotationEntitylistUpdate =
+      "n:1000\ni:" +
+      ANNOTATION_ENTITYLIST_TABLE_NAME +
+      "\nad:1\n" +
+      "a:1:32:" +
+      annotationEntitylistedURL.length +
+      "\n" +
+      annotationEntitylistedURL +
+      "\n";
+    let trackingUpdate =
+      "n:1000\ni:" +
+      TRACKING_TABLE_NAME +
+      "\nad:3\n" +
+      "a:1:32:" +
+      trackingURL1.length +
+      "\n" +
+      trackingURL1 +
+      "\n" +
+      "a:2:32:" +
+      trackingURL2.length +
+      "\n" +
+      trackingURL2 +
+      "\n" +
+      "a:3:32:" +
+      trackingURL3.length +
+      "\n" +
+      trackingURL3 +
+      "\n";
+    let socialTrackingUpdate =
+      "n:1000\ni:" +
+      SOCIAL_TRACKING_TABLE_NAME +
+      "\nad:1\n" +
+      "a:1:32:" +
+      socialTrackingURL.length +
+      "\n" +
+      socialTrackingURL +
+      "\n";
+    let emailTrackingUpdate =
+      "n:1000\ni:" +
+      EMAIL_TRACKING_TABLE_NAME +
+      "\nad:1\n" +
+      "a:1:32:" +
+      emailTrackingURL.length +
+      "\n" +
+      emailTrackingURL +
+      "\n";
+    let entitylistUpdate =
+      "n:1000\ni:" +
+      ENTITYLIST_TABLE_NAME +
+      "\nad:1\n" +
+      "a:1:32:" +
+      entitylistedURL.length +
+      "\n" +
+      entitylistedURL +
+      "\n";
+
+    var tables = [
+      {
+        pref: ANNOTATION_TABLE_PREF,
+        name: ANNOTATION_TABLE_NAME,
+        update: annotationUpdate,
+      },
+      {
+        pref: SOCIAL_ANNOTATION_TABLE_PREF,
+        name: SOCIAL_ANNOTATION_TABLE_NAME,
+        update: socialAnnotationUpdate,
+      },
+      {
+        pref: ANNOTATION_ENTITYLIST_TABLE_PREF,
+        name: ANNOTATION_ENTITYLIST_TABLE_NAME,
+        update: annotationEntitylistUpdate,
+      },
+      {
+        pref: TRACKING_TABLE_PREF,
+        name: TRACKING_TABLE_NAME,
+        update: trackingUpdate,
+      },
+      {
+        pref: SOCIAL_TRACKING_TABLE_PREF,
+        name: SOCIAL_TRACKING_TABLE_NAME,
+        update: socialTrackingUpdate,
+      },
+      {
+        pref: EMAIL_TRACKING_TABLE_PREF,
+        name: EMAIL_TRACKING_TABLE_NAME,
+        update: emailTrackingUpdate,
+      },
+      {
+        pref: ENTITYLIST_TABLE_PREF,
+        name: ENTITYLIST_TABLE_NAME,
+        update: entitylistUpdate,
+      },
+    ];
+
+    let tableIndex = 0;
+    let doOneUpdate = () => {
+      if (tableIndex == tables.length) {
+        return Promise.resolve();
+      }
+      return this.useTestDatabase(tables[tableIndex]).then(
+        () => {
+          tableIndex++;
+          return doOneUpdate();
+        },
+        aErrMsg => {
+          dump("Rejected: " + aErrMsg + ". Retry later.\n");
+          return new Promise(resolve => {
+            timer.initWithCallback(resolve, 100, Ci.nsITimer.TYPE_ONE_SHOT);
+          }).then(doOneUpdate);
+        }
+      );
+    };
+
+    return doOneUpdate();
+  },
+
+  cleanupTestTrackers() {
+    Services.prefs.clearUserPref(ANNOTATION_TABLE_PREF);
+    Services.prefs.clearUserPref(SOCIAL_ANNOTATION_TABLE_PREF);
+    Services.prefs.clearUserPref(ANNOTATION_ENTITYLIST_TABLE_PREF);
+    Services.prefs.clearUserPref(TRACKING_TABLE_PREF);
+    Services.prefs.clearUserPref(SOCIAL_TRACKING_TABLE_PREF);
+    Services.prefs.clearUserPref(EMAIL_TRACKING_TABLE_PREF);
+    Services.prefs.clearUserPref(ENTITYLIST_TABLE_PREF);
+  },
+
+  /**
+   * Add some entries to a test tracking protection database, and resets
+   * back to the default database after the test ends.
+   *
+   * @return {Promise}
+   */
+  useTestDatabase(table) {
+    Services.prefs.setCharPref(table.pref, table.name);
+
+    return new Promise((resolve, reject) => {
+      let dbService = Cc["@mozilla.org/url-classifier/dbservice;1"].getService(
+        Ci.nsIUrlClassifierDBService
+      );
+      let listener = {
+        QueryInterface: iid => {
+          if (
+            iid.equals(Ci.nsISupports) ||
+            iid.equals(Ci.nsIUrlClassifierUpdateObserver)
+          ) {
+            return listener;
+          }
+
+          throw Components.Exception("", Cr.NS_ERROR_NO_INTERFACE);
+        },
+        updateUrlRequested: url => {},
+        streamFinished: status => {},
+        updateError: errorCode => {
+          reject("Got updateError when updating " + table.name);
+        },
+        updateSuccess: requestedTimeout => {
+          resolve();
+        },
+      };
+
+      try {
+        dbService.beginUpdate(listener, table.name, "");
+        dbService.beginStream("", "");
+        dbService.updateStream(table.update);
+        dbService.finishStream();
+        dbService.finishUpdate();
+      } catch (e) {
+        reject("Failed to update with dbService: " + table.name);
+      }
+    });
+  },
+
+  /**
+   * Handle the next "urlclassifier-before-block-channel" event.
+   * @param {Object} options
+   * @param {String} [options.filterOrigin] - Only handle event for channels
+   * with matching origin.
+   * @param {function} [options.onBeforeBlockChannel] - Optional callback for
+   * the event. Called before acting on the channel.
+   * @param {("allow"|"replace")} [options.action] - Whether to allow or replace
+   * the channel.
+   * @returns {Promise} - Resolves once event has been handled.
+   */
+  handleBeforeBlockChannel({
+    filterOrigin = null,
+    onBeforeBlockChannel,
+    action,
+  }) {
+    if (action && action != "allow" && action != "replace") {
+      throw new Error("Invalid action " + action);
+    }
+    let channelClassifierService = Cc[
+      "@mozilla.org/url-classifier/channel-classifier-service;1"
+    ].getService(Ci.nsIChannelClassifierService);
+
+    let resolver;
+    let promise = new Promise(resolve => {
+      resolver = resolve;
+    });
+
+    let observer = {
+      observe(subject, topic) {
+        if (topic != "urlclassifier-before-block-channel") {
+          return;
+        }
+        let channel = subject.QueryInterface(Ci.nsIUrlClassifierBlockedChannel);
+
+        if (filterOrigin) {
+          let { url } = channel;
+          let { origin } = new URL(url);
+          if (filterOrigin != origin) {
+            return;
+          }
+        }
+
+        if (onBeforeBlockChannel) {
+          onBeforeBlockChannel(channel);
+        }
+        if (action) {
+          channel[action]();
+        }
+
+        channelClassifierService.removeListener(observer);
+        resolver();
+      },
+    };
+    channelClassifierService.addListener(observer);
+    return promise;
+  },
+};