1 files changed, 204 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/mochitest/classifierHelper.js b/toolkit/components/url-classifier/tests/mochitest/classifierHelper.js
new file mode 100644
index 0000000000..157e5d54f7
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/mochitest/classifierHelper.js
@@ -0,0 +1,204 @@
+if (typeof classifierHelper == "undefined") {
+  var classifierHelper = {};
+}
+
+const CLASSIFIER_COMMON_URL = SimpleTest.getTestFileURL("classifierCommon.js");
+var gScript = SpecialPowers.loadChromeScript(CLASSIFIER_COMMON_URL);
+var gOriginalGetHashURL;
+
+const PREFS = {
+  PROVIDER_LISTS: "browser.safebrowsing.provider.mozilla.lists",
+  DISALLOW_COMPLETIONS: "urlclassifier.disallow_completions",
+  PROVIDER_GETHASHURL: "browser.safebrowsing.provider.mozilla.gethashURL",
+};
+
+classifierHelper._curAddChunkNum = 1;
+
+// addUrlToDB is asynchronous, queue the task to ensure
+// the callback follow correct order.
+classifierHelper._updates = [];
+
+// Keep urls added to database, those urls should be automatically
+// removed after test complete.
+classifierHelper._updatesToCleanup = [];
+
+classifierHelper._initsCB = [];
+
+// This function return a Promise, promise is resolved when SafeBrowsing.jsm
+// is initialized.
+classifierHelper.waitForInit = function () {
+  return new Promise(function (resolve, reject) {
+    classifierHelper._initsCB.push(resolve);
+    gScript.sendAsyncMessage("waitForInit");
+  });
+};
+
+// This function is used to allow completion for specific "list",
+// some lists like "test-malware-simple" is default disabled to ask for complete.
+// "list" is the db we would like to allow it
+// "url" is the completion server
+classifierHelper.allowCompletion = async function (lists, url) {
+  for (var list of lists) {
+    // Add test db to provider
+    var pref = await SpecialPowers.getParentCharPref(PREFS.PROVIDER_LISTS);
+    pref += "," + list;
+    await SpecialPowers.setCharPref(PREFS.PROVIDER_LISTS, pref);
+
+    // Rename test db so we will not disallow it from completions
+    pref = await SpecialPowers.getParentCharPref(PREFS.DISALLOW_COMPLETIONS);
+    pref = pref.replace(list, list + "-backup");
+    await SpecialPowers.setCharPref(PREFS.DISALLOW_COMPLETIONS, pref);
+  }
+
+  // Store the original get hash URL in order to reset it back during clean up.
+  gOriginalGetHashURL = SpecialPowers.getCharPref(PREFS.PROVIDER_GETHASHURL);
+
+  // Set get hash url
+  await SpecialPowers.setCharPref(PREFS.PROVIDER_GETHASHURL, url);
+};
+
+// Pass { url: ..., db: ... } to add url to database,
+// onsuccess/onerror will be called when update complete.
+classifierHelper.addUrlToDB = function (updateData) {
+  return new Promise(function (resolve, reject) {
+    var testUpdate = "";
+    for (var update of updateData) {
+      var LISTNAME = update.db;
+      var CHUNKDATA = update.url;
+      var CHUNKLEN = CHUNKDATA.length;
+      var HASHLEN = update.len ? update.len : 32;
+
+      update.addChunk = classifierHelper._curAddChunkNum;
+      classifierHelper._curAddChunkNum += 1;
+
+      classifierHelper._updatesToCleanup.push(update);
+      testUpdate +=
+        "n:1000\n" +
+        "i:" +
+        LISTNAME +
+        "\n" +
+        "ad:1\n" +
+        "a:" +
+        update.addChunk +
+        ":" +
+        HASHLEN +
+        ":" +
+        CHUNKLEN +
+        "\n" +
+        CHUNKDATA;
+    }
+
+    classifierHelper._update(testUpdate, resolve, reject);
+  });
+};
+
+// This API is used to expire all add/sub chunks we have updated
+// by using addUrlToDB.
+classifierHelper.resetDatabase = function () {
+  function removeDatabase() {
+    return new Promise(function (resolve, reject) {
+      var testUpdate = "";
+      for (var update of classifierHelper._updatesToCleanup) {
+        testUpdate +=
+          "n:1000\ni:" + update.db + "\nad:" + update.addChunk + "\n";
+      }
+
+      classifierHelper._update(testUpdate, resolve, reject);
+    });
+  }
+
+  // Remove and then reload will ensure both database and cache will
+  // be cleared.
+  return Promise.resolve()
+    .then(removeDatabase)
+    .then(classifierHelper.reloadDatabase);
+};
+
+classifierHelper.reloadDatabase = function () {
+  return new Promise(function (resolve, reject) {
+    gScript.addMessageListener("reloadSuccess", function handler() {
+      gScript.removeMessageListener("reloadSuccess", handler);
+      resolve();
+    });
+
+    gScript.sendAsyncMessage("doReload");
+  });
+};
+
+classifierHelper.getTables = function () {
+  return new Promise(resolve => {
+    gScript.addMessageListener("GetTableSuccess", function handler(tables) {
+      gScript.removeMessageListener("GetTableSuccess", handler);
+      resolve(tables);
+    });
+
+    gScript.sendAsyncMessage("doGetTables");
+  });
+};
+
+classifierHelper._update = function (testUpdate, onsuccess, onerror) {
+  // Queue the task if there is still an on-going update
+  classifierHelper._updates.push({
+    data: testUpdate,
+    onsuccess,
+    onerror,
+  });
+  if (classifierHelper._updates.length != 1) {
+    return;
+  }
+
+  gScript.sendAsyncMessage("doUpdate", { testUpdate });
+};
+
+classifierHelper._updateSuccess = function () {
+  var update = classifierHelper._updates.shift();
+  update.onsuccess();
+
+  if (classifierHelper._updates.length) {
+    var testUpdate = classifierHelper._updates[0].data;
+    gScript.sendAsyncMessage("doUpdate", { testUpdate });
+  }
+};
+
+classifierHelper._updateError = function (errorCode) {
+  var update = classifierHelper._updates.shift();
+  update.onerror(errorCode);
+
+  if (classifierHelper._updates.length) {
+    var testUpdate = classifierHelper._updates[0].data;
+    gScript.sendAsyncMessage("doUpdate", { testUpdate });
+  }
+};
+
+classifierHelper._inited = function () {
+  classifierHelper._initsCB.forEach(function (cb) {
+    cb();
+  });
+  classifierHelper._initsCB = [];
+};
+
+classifierHelper._setup = function () {
+  gScript.addMessageListener("updateSuccess", classifierHelper._updateSuccess);
+  gScript.addMessageListener("updateError", classifierHelper._updateError);
+  gScript.addMessageListener("safeBrowsingInited", classifierHelper._inited);
+
+  // cleanup will be called at end of each testcase to remove all the urls added to database.
+  SimpleTest.registerCleanupFunction(classifierHelper._cleanup);
+};
+
+classifierHelper._cleanup = function () {
+  // clean all the preferences may touch by helper
+  Object.values(PREFS).map(pref => SpecialPowers.clearUserPref(pref));
+
+  // Set the getHashURL back, the original value isn't the same as the default
+  // pref value.
+  SpecialPowers.setCharPref(PREFS.PROVIDER_GETHASHURL, gOriginalGetHashURL);
+
+  if (!classifierHelper._updatesToCleanup) {
+    return Promise.resolve();
+  }
+
+  return classifierHelper.resetDatabase();
+};
+
+classifierHelper._setup();