-
Notifications
You must be signed in to change notification settings - Fork 54
/
classifierHelper.js
173 lines (147 loc) · 5.57 KB
/
classifierHelper.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
// Created from toolkit/components/url-classifier/tests/mochitest/classifierHelper.js
// Unfortunately, browser tests cannot load that script as it is too reliant on
// being loaded in the content process.
let dbService = Cc["@mozilla.org/url-classifier/dbservice;1"]
.getService(Ci.nsIUrlClassifierDBService);
if (typeof(classifierHelper) == "undefined") {
var classifierHelper = {};
}
const HASHLEN = 32;
const PREFS = {
PROVIDER_LISTS: "browser.safebrowsing.provider.mozilla.lists",
DISALLOW_COMPLETIONS: "urlclassifier.disallow_completions",
PROVIDER_GETHASHURL: "browser.safebrowsing.provider.mozilla.gethashURL",
};
classifierHelper._curAddChunkNum = 1;
// Keep urls added to database, those urls should be automatically
// removed after test complete.
classifierHelper._updatesToCleanup = [];
// This function returns a Promise resolved when SafeBrowsing.jsm is initialized.
// SafeBrowsing.jsm is initialized after mozEntries are added. Add observer
// to receive "finished" event. For the case when this function is called
// after the event had already been notified, we lookup entries to see if
// they are already added to database.
classifierHelper.waitForInit = function() {
// This url must sync with the table, url in SafeBrowsing.jsm addMozEntries
const table = "test-phish-simple";
const url = "http://itisatrap.org/firefox/its-a-trap.html";
let principal = Services.scriptSecurityManager.createCodebasePrincipal(
Services.io.newURI(url), {});
return new Promise(function(resolve, reject) {
Services.obs.addObserver(function() {
resolve();
}, "mozentries-update-finished");
let listener = {
QueryInterface: ChromeUtils.generateQI(["nsIUrlClassifierUpdateObserver"]),
handleEvent(value) {
if (value === table) {
resolve();
}
},
};
dbService.lookup(principal, table, listener);
});
};
// 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 (let list of lists) {
// Add test db to provider
let pref = Services.prefs.getCharPref(PREFS.PROVIDER_LISTS);
pref += "," + list;
Services.prefs.setCharPref(PREFS.PROVIDER_LISTS, pref);
// Rename test db so we will not disallow it from completions
pref = Services.prefs.getCharPref(PREFS.DISALLOW_COMPLETIONS);
pref = pref.replace(list, list + "-backup");
Services.prefs.setCharPref(PREFS.DISALLOW_COMPLETIONS, pref);
}
// Set get hash url
Services.prefs.setCharPref(PREFS.PROVIDER_GETHASHURL, url);
};
// Pass { url: ..., db: ... } to add url to database,
// Returns a Promise.
classifierHelper.addUrlToDB = function(updateData) {
let testUpdate = "";
for (let update of updateData) {
let LISTNAME = update.db;
let CHUNKDATA = update.url;
let CHUNKLEN = CHUNKDATA.length;
let 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;
}
return classifierHelper._update(testUpdate);
};
// This API is used to expire all add/sub chunks we have updated
// by using addUrlToDB.
// Returns a Promise.
classifierHelper.resetDatabase = function() {
var testUpdate = "";
for (var update of classifierHelper._updatesToCleanup) {
testUpdate +=
"n:1000\n" +
"i:" + update.db + "\n" +
"ad:" + update.addChunk + "\n";
}
return classifierHelper._update(testUpdate);
};
classifierHelper.reloadDatabase = function() {
dbService.reloadDatabase();
};
classifierHelper._update = function(update) {
return (async function() {
// beginUpdate may fail if there's an existing update in progress
// retry until success or testcase timeout.
let success = false;
while (!success) {
try {
await new Promise((resolve, reject) => {
let listener = {
QueryInterface: ChromeUtils.generateQI(["nsIUrlClassifierUpdateObserver"]),
updateUrlRequested(url) { },
streamFinished(status) { },
updateError(errorCode) {
reject(errorCode);
},
updateSuccess(requestedTimeout) {
resolve();
},
};
dbService.beginUpdate(listener, "", "");
dbService.beginStream("", "");
dbService.updateStream(update);
dbService.finishStream();
dbService.finishUpdate();
});
success = true;
} catch (e) {
// Wait 1 second before trying again.
await new Promise(resolve => setTimeout(resolve, 1000));
}
}
})();
};
classifierHelper._cleanup = function() {
// Clean all the preferences that may have been touched by classifierHelper
for (var pref in PREFS) {
Services.prefs.clearUserPref(pref);
}
if (!classifierHelper._updatesToCleanup) {
return Promise.resolve();
}
return classifierHelper.resetDatabase();
};
// Cleanup will be called at end of each testcase to remove all the urls added to database.
registerCleanupFunction(classifierHelper._cleanup);