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();