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
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
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();
|