1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
2 /* This Source Code Form is subject to the terms of the Mozilla Public
3 * License, v. 2.0. If a copy of the MPL was not distributed with this
4 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
5
6 #include "LookupCacheV4.h"
7 #include "mozilla/Preferences.h"
8 #include <mozilla/RefPtr.h>
9 #include "nsAppDirectoryServiceDefs.h"
10 #include "nsClassHashtable.h"
11 #include "nsString.h"
12 #include "VariableLengthPrefixSet.h"
13
14 #include "Common.h"
15
16 // Create fullhash by appending random characters.
CreateFullHash(const nsACString & in)17 static nsCString CreateFullHash(const nsACString& in) {
18 nsCString out(in);
19 out.SetLength(32);
20 for (size_t i = in.Length(); i < 32; i++) {
21 out.SetCharAt(char(rand() % 256), i);
22 }
23
24 return out;
25 }
26
27 // This function generate N prefixes with size between MIN and MAX.
28 // The output array will not be cleared, random result will append to it
RandomPrefixes(uint32_t N,uint32_t MIN,uint32_t MAX,_PrefixArray & array)29 static void RandomPrefixes(uint32_t N, uint32_t MIN, uint32_t MAX,
30 _PrefixArray& array) {
31 array.SetCapacity(array.Length() + N);
32
33 uint32_t range = (MAX - MIN + 1);
34
35 for (uint32_t i = 0; i < N; i++) {
36 uint32_t prefixSize = (rand() % range) + MIN;
37 _Prefix prefix;
38 prefix.SetLength(prefixSize);
39
40 bool added = false;
41 while (!added) {
42 char* dst = prefix.BeginWriting();
43 for (uint32_t j = 0; j < prefixSize; j++) {
44 dst[j] = rand() % 256;
45 }
46
47 if (!array.Contains(prefix)) {
48 array.AppendElement(prefix);
49 added = true;
50 }
51 }
52 }
53 }
54
55 // This test loops through all the prefixes and converts each prefix to
56 // fullhash by appending random characters, each converted fullhash
57 // should at least match its original length in the prefixSet.
DoExpectedLookup(LookupCacheV4 * cache,_PrefixArray & array)58 static void DoExpectedLookup(LookupCacheV4* cache, _PrefixArray& array) {
59 uint32_t matchLength = 0;
60 for (uint32_t i = 0; i < array.Length(); i++) {
61 const nsCString& prefix = array[i];
62 Completion complete;
63 complete.Assign(CreateFullHash(prefix));
64
65 // Find match for prefix-generated full hash
66 bool has, confirmed;
67 cache->Has(complete, &has, &matchLength, &confirmed);
68 MOZ_ASSERT(matchLength != 0);
69
70 if (matchLength != prefix.Length()) {
71 // Return match size is not the same as prefix size.
72 // In this case it could be because the generated fullhash match other
73 // prefixes, check if this prefix exist.
74 bool found = false;
75
76 for (uint32_t j = 0; j < array.Length(); j++) {
77 if (array[j].Length() != matchLength) {
78 continue;
79 }
80
81 if (0 == memcmp(complete.buf, array[j].BeginReading(), matchLength)) {
82 found = true;
83 break;
84 }
85 }
86 ASSERT_TRUE(found);
87 }
88 }
89 }
90
DoRandomLookup(LookupCacheV4 * cache,uint32_t N,_PrefixArray & array)91 static void DoRandomLookup(LookupCacheV4* cache, uint32_t N,
92 _PrefixArray& array) {
93 for (uint32_t i = 0; i < N; i++) {
94 // Random 32-bytes test fullhash
95 char buf[32];
96 for (uint32_t j = 0; j < 32; j++) {
97 buf[j] = (char)(rand() % 256);
98 }
99
100 // Get the expected result.
101 nsTArray<uint32_t> expected;
102 for (uint32_t j = 0; j < array.Length(); j++) {
103 const nsACString& str = array[j];
104 if (0 == memcmp(buf, str.BeginReading(), str.Length())) {
105 expected.AppendElement(str.Length());
106 }
107 }
108
109 Completion complete;
110 complete.Assign(nsDependentCSubstring(buf, 32));
111 bool has, confirmed;
112 uint32_t matchLength = 0;
113 cache->Has(complete, &has, &matchLength, &confirmed);
114
115 ASSERT_TRUE(expected.IsEmpty() ? !matchLength
116 : expected.Contains(matchLength));
117 }
118 }
119
SetupLookupCache(const nsACString & aName)120 static already_AddRefed<LookupCacheV4> SetupLookupCache(
121 const nsACString& aName) {
122 nsCOMPtr<nsIFile> rootDir;
123 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(rootDir));
124
125 nsAutoCString provider("test");
126 RefPtr<LookupCacheV4> lookup = new LookupCacheV4(aName, provider, rootDir);
127 lookup->Init();
128
129 return lookup.forget();
130 }
131
132 class UrlClassifierPrefixSetTest : public ::testing::TestWithParam<uint32_t> {
133 protected:
SetUp()134 void SetUp() override {
135 // max_array_size to 0 means we are testing delta algorithm here.
136 static const char prefKey[] =
137 "browser.safebrowsing.prefixset.max_array_size";
138 mozilla::Preferences::SetUint(prefKey, GetParam());
139
140 mCache = SetupLookupCache("test"_ns);
141 }
142
TearDown()143 void TearDown() override {
144 mCache = nullptr;
145 mArray.Clear();
146 mMap.Clear();
147 }
148
SetupPrefixes(_PrefixArray && aArray)149 nsresult SetupPrefixes(_PrefixArray&& aArray) {
150 mArray = std::move(aArray);
151 PrefixArrayToPrefixStringMap(mArray, mMap);
152 return mCache->Build(mMap);
153 }
154
SetupPrefixesAndVerify(_PrefixArray & aArray)155 void SetupPrefixesAndVerify(_PrefixArray& aArray) {
156 mArray = aArray.Clone();
157 PrefixArrayToPrefixStringMap(mArray, mMap);
158
159 ASSERT_TRUE(NS_SUCCEEDED(mCache->Build(mMap)));
160 Verify();
161 }
162
SetupPrefixesAndVerify(_PrefixArray && aArray)163 void SetupPrefixesAndVerify(_PrefixArray&& aArray) {
164 nsresult rv = SetupPrefixes(std::move(aArray));
165 ASSERT_TRUE(NS_SUCCEEDED(rv));
166 Verify();
167 }
168
SetupRandomPrefixesAndVerify(uint32_t N,uint32_t MIN,uint32_t MAX)169 void SetupRandomPrefixesAndVerify(uint32_t N, uint32_t MIN, uint32_t MAX) {
170 srand(time(nullptr));
171 RandomPrefixes(N, MIN, MAX, mArray);
172 PrefixArrayToPrefixStringMap(mArray, mMap);
173
174 ASSERT_TRUE(NS_SUCCEEDED(mCache->Build(mMap)));
175 Verify();
176 }
177
Verify()178 void Verify() {
179 DoExpectedLookup(mCache, mArray);
180 DoRandomLookup(mCache, 1000, mArray);
181 CheckContent(mCache, mArray);
182 }
183
184 RefPtr<LookupCacheV4> mCache;
185 _PrefixArray mArray;
186 PrefixStringMap mMap;
187 };
188
189 // Test setting prefix set with only 4-bytes prefixes
TEST_P(UrlClassifierPrefixSetTest,FixedLengthSet)190 TEST_P(UrlClassifierPrefixSetTest, FixedLengthSet) {
191 SetupPrefixesAndVerify({
192 _Prefix("alph"),
193 _Prefix("brav"),
194 _Prefix("char"),
195 _Prefix("delt"),
196 _Prefix("echo"),
197 _Prefix("foxt"),
198 });
199 }
200
TEST_P(UrlClassifierPrefixSetTest,FixedLengthRandomSet)201 TEST_P(UrlClassifierPrefixSetTest, FixedLengthRandomSet) {
202 SetupRandomPrefixesAndVerify(1500, 4, 4);
203 }
204
TEST_P(UrlClassifierPrefixSetTest,FixedLengthRandomLargeSet)205 TEST_P(UrlClassifierPrefixSetTest, FixedLengthRandomLargeSet) {
206 SetupRandomPrefixesAndVerify(15000, 4, 4);
207 }
208
TEST_P(UrlClassifierPrefixSetTest,FixedLengthTinySet)209 TEST_P(UrlClassifierPrefixSetTest, FixedLengthTinySet) {
210 SetupPrefixesAndVerify({
211 _Prefix("tiny"),
212 });
213 }
214
215 // Test setting prefix set with only 5~32 bytes prefixes
TEST_P(UrlClassifierPrefixSetTest,VariableLengthSet)216 TEST_P(UrlClassifierPrefixSetTest, VariableLengthSet) {
217 SetupPrefixesAndVerify(
218 {_Prefix("bravo"), _Prefix("charlie"), _Prefix("delta"),
219 _Prefix("EchoEchoEchoEchoEcho"), _Prefix("foxtrot"),
220 _Prefix("GolfGolfGolfGolfGolfGolfGolfGolf"), _Prefix("hotel"),
221 _Prefix("november"), _Prefix("oscar"), _Prefix("quebec"),
222 _Prefix("romeo"), _Prefix("sierrasierrasierrasierrasierra"),
223 _Prefix("Tango"), _Prefix("whiskey"), _Prefix("yankee"),
224 _Prefix("ZuluZuluZuluZulu")});
225 }
226
TEST_P(UrlClassifierPrefixSetTest,VariableLengthRandomSet)227 TEST_P(UrlClassifierPrefixSetTest, VariableLengthRandomSet) {
228 SetupRandomPrefixesAndVerify(1500, 5, 32);
229 }
230
231 // Test setting prefix set with both 4-bytes prefixes and 5~32 bytes prefixes
TEST_P(UrlClassifierPrefixSetTest,MixedPrefixSet)232 TEST_P(UrlClassifierPrefixSetTest, MixedPrefixSet) {
233 SetupPrefixesAndVerify(
234 {_Prefix("enus"), _Prefix("apollo"), _Prefix("mars"),
235 _Prefix("Hecatonchires cyclopes"), _Prefix("vesta"), _Prefix("neptunus"),
236 _Prefix("jupiter"), _Prefix("diana"), _Prefix("minerva"),
237 _Prefix("ceres"), _Prefix("Aidos,Adephagia,Adikia,Aletheia"),
238 _Prefix("hecatonchires"), _Prefix("alcyoneus"), _Prefix("hades"),
239 _Prefix("vulcanus"), _Prefix("juno"), _Prefix("mercury"),
240 _Prefix("Stheno, Euryale and Medusa")});
241 }
242
TEST_P(UrlClassifierPrefixSetTest,MixedRandomPrefixSet)243 TEST_P(UrlClassifierPrefixSetTest, MixedRandomPrefixSet) {
244 SetupRandomPrefixesAndVerify(1500, 4, 32);
245 }
246
247 // Test resetting prefix set
TEST_P(UrlClassifierPrefixSetTest,ResetPrefix)248 TEST_P(UrlClassifierPrefixSetTest, ResetPrefix) {
249 // Base prefix set
250 _PrefixArray oldArray = {
251 _Prefix("Iceland"), _Prefix("Peru"), _Prefix("Mexico"),
252 _Prefix("Australia"), _Prefix("Japan"), _Prefix("Egypt"),
253 _Prefix("America"), _Prefix("Finland"), _Prefix("Germany"),
254 _Prefix("Italy"), _Prefix("France"), _Prefix("Taiwan"),
255 };
256 SetupPrefixesAndVerify(oldArray);
257
258 // New prefix set
259 _PrefixArray newArray = {
260 _Prefix("Pikachu"), _Prefix("Bulbasaur"), _Prefix("Charmander"),
261 _Prefix("Blastoise"), _Prefix("Pidgey"), _Prefix("Mewtwo"),
262 _Prefix("Jigglypuff"), _Prefix("Persian"), _Prefix("Tentacool"),
263 _Prefix("Onix"), _Prefix("Eevee"), _Prefix("Jynx"),
264 };
265 SetupPrefixesAndVerify(newArray);
266
267 // Should not match any of the first prefix set
268 uint32_t matchLength = 0;
269 for (uint32_t i = 0; i < oldArray.Length(); i++) {
270 Completion complete;
271 complete.Assign(CreateFullHash(oldArray[i]));
272
273 // Find match for prefix-generated full hash
274 bool has, confirmed;
275 mCache->Has(complete, &has, &matchLength, &confirmed);
276
277 ASSERT_TRUE(matchLength == 0);
278 }
279 }
280
281 // Test only set one 4-bytes prefix and one full-length prefix
TEST_P(UrlClassifierPrefixSetTest,TinyPrefixSet)282 TEST_P(UrlClassifierPrefixSetTest, TinyPrefixSet) {
283 SetupPrefixesAndVerify({
284 _Prefix("AAAA"),
285 _Prefix("11112222333344445555666677778888"),
286 });
287 }
288
289 // Test empty prefix set and IsEmpty function
TEST_P(UrlClassifierPrefixSetTest,EmptyFixedPrefixSet)290 TEST_P(UrlClassifierPrefixSetTest, EmptyFixedPrefixSet) {
291 ASSERT_TRUE(mCache->IsEmpty());
292
293 SetupPrefixesAndVerify({});
294
295 // Insert an 4-bytes prefix, then IsEmpty should return false
296 SetupPrefixesAndVerify({_Prefix("test")});
297
298 ASSERT_TRUE(!mCache->IsEmpty());
299 }
300
TEST_P(UrlClassifierPrefixSetTest,EmptyVariableLengthPrefixSet)301 TEST_P(UrlClassifierPrefixSetTest, EmptyVariableLengthPrefixSet) {
302 ASSERT_TRUE(mCache->IsEmpty());
303
304 SetupPrefixesAndVerify({});
305
306 // Insert an 5~32 bytes prefix, then IsEmpty should return false
307 SetupPrefixesAndVerify({_Prefix("test variable length")});
308
309 ASSERT_TRUE(!mCache->IsEmpty());
310 }
311
312 // Test prefix size should only between 4~32 bytes
TEST_P(UrlClassifierPrefixSetTest,MinMaxPrefixSet)313 TEST_P(UrlClassifierPrefixSetTest, MinMaxPrefixSet) {
314 // Test prefix set between 4-32 bytes, should success
315 SetupPrefixesAndVerify({_Prefix("1234"), _Prefix("ABCDEFGHIJKKMNOP"),
316 _Prefix("1aaa2bbb3ccc4ddd5eee6fff7ggg8hhh")});
317
318 // Prefix size less than 4-bytes should fail
319 nsresult rv = SetupPrefixes({_Prefix("123")});
320 ASSERT_TRUE(NS_FAILED(rv));
321
322 // Prefix size greater than 32-bytes should fail
323 rv = SetupPrefixes({_Prefix("1aaa2bbb3ccc4ddd5eee6fff7ggg8hhh9")});
324 ASSERT_TRUE(NS_FAILED(rv));
325 }
326
327 // Test save then load prefix set with only 4-bytes prefixes
TEST_P(UrlClassifierPrefixSetTest,LoadSaveFixedLengthPrefixSet)328 TEST_P(UrlClassifierPrefixSetTest, LoadSaveFixedLengthPrefixSet) {
329 nsCOMPtr<nsIFile> file;
330 _PrefixArray array;
331 PrefixStringMap map;
332
333 // Save
334 {
335 RefPtr<LookupCacheV4> save = SetupLookupCache("test-save"_ns);
336
337 RandomPrefixes(10000, 4, 4, array);
338
339 PrefixArrayToPrefixStringMap(array, map);
340 save->Build(map);
341
342 DoExpectedLookup(save, array);
343 DoRandomLookup(save, 1000, array);
344 CheckContent(save, array);
345
346 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
347 file->Append(u"test.vlpset"_ns);
348 save->StoreToFile(file);
349 }
350
351 // Load
352 {
353 RefPtr<LookupCacheV4> load = SetupLookupCache("test-load"_ns);
354 load->LoadFromFile(file);
355
356 DoExpectedLookup(load, array);
357 DoRandomLookup(load, 1000, array);
358 CheckContent(load, array);
359 }
360
361 file->Remove(false);
362 }
363
364 // Test save then load prefix set with only 5~32 bytes prefixes
TEST_P(UrlClassifierPrefixSetTest,LoadSaveVariableLengthPrefixSet)365 TEST_P(UrlClassifierPrefixSetTest, LoadSaveVariableLengthPrefixSet) {
366 nsCOMPtr<nsIFile> file;
367 _PrefixArray array;
368 PrefixStringMap map;
369
370 // Save
371 {
372 RefPtr<LookupCacheV4> save = SetupLookupCache("test-save"_ns);
373
374 RandomPrefixes(10000, 5, 32, array);
375
376 PrefixArrayToPrefixStringMap(array, map);
377 save->Build(map);
378
379 DoExpectedLookup(save, array);
380 DoRandomLookup(save, 1000, array);
381 CheckContent(save, array);
382
383 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
384 file->Append(u"test.vlpset"_ns);
385 save->StoreToFile(file);
386 }
387
388 // Load
389 {
390 RefPtr<LookupCacheV4> load = SetupLookupCache("test-load"_ns);
391 load->LoadFromFile(file);
392
393 DoExpectedLookup(load, array);
394 DoRandomLookup(load, 1000, array);
395 CheckContent(load, array);
396 }
397
398 file->Remove(false);
399 }
400
401 // Test save then load prefix with both 4 bytes prefixes and 5~32 bytes prefixes
TEST_P(UrlClassifierPrefixSetTest,LoadSavePrefixSet)402 TEST_P(UrlClassifierPrefixSetTest, LoadSavePrefixSet) {
403 nsCOMPtr<nsIFile> file;
404 _PrefixArray array;
405 PrefixStringMap map;
406
407 // Save
408 {
409 RefPtr<LookupCacheV4> save = SetupLookupCache("test-save"_ns);
410
411 // Try to simulate the real case that most prefixes are 4bytes
412 RandomPrefixes(20000, 4, 4, array);
413 RandomPrefixes(1000, 5, 32, array);
414
415 PrefixArrayToPrefixStringMap(array, map);
416 save->Build(map);
417
418 DoExpectedLookup(save, array);
419 DoRandomLookup(save, 1000, array);
420 CheckContent(save, array);
421
422 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
423 file->Append(u"test.vlpset"_ns);
424 save->StoreToFile(file);
425 }
426
427 // Load
428 {
429 RefPtr<LookupCacheV4> load = SetupLookupCache("test-load"_ns);
430 load->LoadFromFile(file);
431
432 DoExpectedLookup(load, array);
433 DoRandomLookup(load, 1000, array);
434 CheckContent(load, array);
435 }
436
437 file->Remove(false);
438 }
439
440 // This is for fixed-length prefixset
TEST_P(UrlClassifierPrefixSetTest,LoadSaveNoDelta)441 TEST_P(UrlClassifierPrefixSetTest, LoadSaveNoDelta) {
442 nsCOMPtr<nsIFile> file;
443 _PrefixArray array;
444 PrefixStringMap map;
445
446 for (uint32_t i = 0; i < 100; i++) {
447 // construct a tree without deltas by making the distance
448 // between entries larger than 16 bits
449 uint32_t v = ((1 << 16) + 1) * i;
450 nsCString* ele = array.AppendElement();
451 ele->AppendASCII(reinterpret_cast<const char*>(&v), 4);
452 }
453
454 // Save
455 {
456 RefPtr<LookupCacheV4> save = SetupLookupCache("test-save"_ns);
457
458 PrefixArrayToPrefixStringMap(array, map);
459 save->Build(map);
460
461 DoExpectedLookup(save, array);
462 DoRandomLookup(save, 1000, array);
463 CheckContent(save, array);
464
465 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
466 file->Append(u"test.vlpset"_ns);
467 save->StoreToFile(file);
468 }
469
470 // Load
471 {
472 RefPtr<LookupCacheV4> load = SetupLookupCache("test-load"_ns);
473 load->LoadFromFile(file);
474
475 DoExpectedLookup(load, array);
476 DoRandomLookup(load, 1000, array);
477 CheckContent(load, array);
478 }
479
480 file->Remove(false);
481 }
482
483 // To run the same test for different configurations of
484 // "browser_safebrowsing_prefixset_max_array_size"
485 INSTANTIATE_TEST_CASE_P(UrlClassifierPrefixSetTest, UrlClassifierPrefixSetTest,
486 ::testing::Values(0, UINT32_MAX));
487