1 files changed, 201 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 000000000..973f0c2c4
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/mochitest/classifierHelper.js
@@ -0,0 +1,201 @@
+if (typeof(classifierHelper) == "undefined") {
+  var classifierHelper = {};
+}
+
+const CLASSIFIER_COMMON_URL = SimpleTest.getTestFileURL("classifierCommon.js");
+var gScript = SpecialPowers.loadChromeScript(CLASSIFIER_COMMON_URL);
+
+const ADD_CHUNKNUM = 524;
+const SUB_CHUNKNUM = 523;
+const HASHLEN = 32;
+
+const PREFS = {
+  PROVIDER_LISTS : "browser.safebrowsing.provider.mozilla.lists",
+  DISALLOW_COMPLETIONS : "urlclassifier.disallow_completions",
+  PROVIDER_GETHASHURL : "browser.safebrowsing.provider.mozilla.gethashURL"
+};
+
+// addUrlToDB & removeUrlFromDB are 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 = function(lists, url) {
+  for (var list of lists) {
+    // Add test db to provider
+    var pref = SpecialPowers.getCharPref(PREFS.PROVIDER_LISTS);
+    pref += "," + list;
+    SpecialPowers.setCharPref(PREFS.PROVIDER_LISTS, pref);
+
+    // Rename test db so we will not disallow it from completions
+    pref = SpecialPowers.getCharPref(PREFS.DISALLOW_COMPLETIONS);
+    pref = pref.replace(list, list + "-backup");
+    SpecialPowers.setCharPref(PREFS.DISALLOW_COMPLETIONS, pref);
+  }
+
+  // Set get hash url
+  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;
+
+      classifierHelper._updatesToCleanup.push(update);
+      testUpdate +=
+        "n:1000\n" +
+        "i:" + LISTNAME + "\n" +
+        "ad:1\n" +
+        "a:" + ADD_CHUNKNUM + ":" + HASHLEN + ":" + CHUNKLEN + "\n" +
+        CHUNKDATA;
+    }
+
+    classifierHelper._update(testUpdate, resolve, reject);
+  });
+}
+
+// Pass { url: ..., db: ... } to remove url from database,
+// onsuccess/onerror will be called when update complete.
+classifierHelper.removeUrlFromDB = function(updateData) {
+  return new Promise(function(resolve, reject) {
+    var testUpdate = "";
+    for (var update of updateData) {
+      var LISTNAME = update.db;
+      var CHUNKDATA = ADD_CHUNKNUM + ":" + update.url;
+      var CHUNKLEN = CHUNKDATA.length;
+      var HASHLEN = update.len ? update.len : 32;
+
+      testUpdate +=
+        "n:1000\n" +
+        "i:" + LISTNAME + "\n" +
+        "s:" + SUB_CHUNKNUM + ":" + HASHLEN + ":" + CHUNKLEN + "\n" +
+        CHUNKDATA;
+    }
+
+    classifierHelper._updatesToCleanup =
+      classifierHelper._updatesToCleanup.filter((v) => {
+        return updateData.indexOf(v) == -1;
+      });
+
+    classifierHelper._update(testUpdate, resolve, reject);
+  });
+};
+
+// This API is used to expire all add/sub chunks we have updated
+// by using addUrlToDB and removeUrlFromDB.
+classifierHelper.resetDB = function() {
+  return new Promise(function(resolve, reject) {
+    var testUpdate = "";
+    for (var update of classifierHelper._updatesToCleanup) {
+      if (testUpdate.includes(update.db))
+        continue;
+
+      testUpdate +=
+        "n:1000\n" +
+        "i:" + update.db + "\n" +
+        "ad:" + ADD_CHUNKNUM + "\n" +
+        "sd:" + SUB_CHUNKNUM + "\n"
+    }
+
+    classifierHelper._update(testUpdate, resolve, reject);
+  });
+};
+
+classifierHelper.reloadDatabase = function() {
+  return new Promise(function(resolve, reject) {
+    gScript.addMessageListener("reloadSuccess", function handler() {
+      gScript.removeMessageListener('reloadSuccess', handler);
+      resolve();
+    });
+
+    gScript.sendAsyncMessage("doReload");
+  });
+}
+
+classifierHelper._update = function(testUpdate, onsuccess, onerror) {
+  // Queue the task if there is still an on-going update
+  classifierHelper._updates.push({"data": testUpdate,
+                                  "onsuccess": onsuccess,
+                                  "onerror": 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
+  for (var pref in PREFS) {
+    SpecialPowers.clearUserPref(pref);
+  }
+
+  if (!classifierHelper._updatesToCleanup) {
+    return Promise.resolve();
+  }
+
+  return classifierHelper.resetDB();
+};
+
+classifierHelper._setup();