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
202
203
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) {
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) {
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();
|
