1 /* Any copyright is dedicated to the Public Domain.
2 * http://creativecommons.org/publicdomain/zero/1.0/ */
3
4 #include "Common.h"
5 #include "Classifier.h"
6 #include "HashStore.h"
7 #include "nsAppDirectoryServiceDefs.h"
8 #include "nsIFile.h"
9 #include "nsIThread.h"
10 #include "string.h"
11 #include "gtest/gtest.h"
12 #include "nsThreadUtils.h"
13
14 using namespace mozilla;
15 using namespace mozilla::safebrowsing;
16
17 typedef nsCString _Prefix;
18 typedef nsTArray<_Prefix> _PrefixArray;
19
20 #define GTEST_SAFEBROWSING_DIR NS_LITERAL_CSTRING("safebrowsing")
21 #define GTEST_TABLE NS_LITERAL_CSTRING("gtest-malware-proto")
22 #define GTEST_PREFIXFILE NS_LITERAL_CSTRING("gtest-malware-proto.pset")
23
24 // This function removes common elements of inArray and outArray from
25 // outArray. This is used by partial update testcase to ensure partial update
26 // data won't contain prefixes we already have.
RemoveIntersection(const _PrefixArray & inArray,_PrefixArray & outArray)27 static void RemoveIntersection(const _PrefixArray& inArray,
28 _PrefixArray& outArray) {
29 for (uint32_t i = 0; i < inArray.Length(); i++) {
30 int32_t idx = outArray.BinaryIndexOf(inArray[i]);
31 if (idx >= 0) {
32 outArray.RemoveElementAt(idx);
33 }
34 }
35 }
36
37 // This fucntion removes elements from outArray by index specified in
38 // removal array.
RemoveElements(const nsTArray<uint32_t> & removal,_PrefixArray & outArray)39 static void RemoveElements(const nsTArray<uint32_t>& removal,
40 _PrefixArray& outArray) {
41 for (int32_t i = removal.Length() - 1; i >= 0; i--) {
42 outArray.RemoveElementAt(removal[i]);
43 }
44 }
45
MergeAndSortArray(const _PrefixArray & array1,const _PrefixArray & array2,_PrefixArray & output)46 static void MergeAndSortArray(const _PrefixArray& array1,
47 const _PrefixArray& array2,
48 _PrefixArray& output) {
49 output.Clear();
50 output.AppendElements(array1);
51 output.AppendElements(array2);
52 output.Sort();
53 }
54
CalculateCheckSum(_PrefixArray & prefixArray,nsCString & checksum)55 static void CalculateCheckSum(_PrefixArray& prefixArray, nsCString& checksum) {
56 prefixArray.Sort();
57
58 nsresult rv;
59 nsCOMPtr<nsICryptoHash> cryptoHash =
60 do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID, &rv);
61
62 cryptoHash->Init(nsICryptoHash::SHA256);
63 for (uint32_t i = 0; i < prefixArray.Length(); i++) {
64 const _Prefix& prefix = prefixArray[i];
65 cryptoHash->Update(
66 reinterpret_cast<uint8_t*>(const_cast<char*>(prefix.get())),
67 prefix.Length());
68 }
69 cryptoHash->Finish(false, checksum);
70 }
71
72 // N: Number of prefixes, MIN/MAX: minimum/maximum prefix size
73 // This function will append generated prefixes to outArray.
CreateRandomSortedPrefixArray(uint32_t N,uint32_t MIN,uint32_t MAX,_PrefixArray & outArray)74 static void CreateRandomSortedPrefixArray(uint32_t N, uint32_t MIN,
75 uint32_t MAX,
76 _PrefixArray& outArray) {
77 outArray.SetCapacity(outArray.Length() + N);
78
79 const uint32_t range = (MAX - MIN + 1);
80
81 for (uint32_t i = 0; i < N; i++) {
82 uint32_t prefixSize = (rand() % range) + MIN;
83 _Prefix prefix;
84 prefix.SetLength(prefixSize);
85
86 while (true) {
87 char* dst = prefix.BeginWriting();
88 for (uint32_t j = 0; j < prefixSize; j++) {
89 dst[j] = rand() % 256;
90 }
91
92 if (!outArray.Contains(prefix)) {
93 outArray.AppendElement(prefix);
94 break;
95 }
96 }
97 }
98
99 outArray.Sort();
100 }
101
102 // N: Number of removal indices, MAX: maximum index
CreateRandomRemovalIndices(uint32_t N,uint32_t MAX,nsTArray<uint32_t> & outArray)103 static void CreateRandomRemovalIndices(uint32_t N, uint32_t MAX,
104 nsTArray<uint32_t>& outArray) {
105 for (uint32_t i = 0; i < N; i++) {
106 uint32_t idx = rand() % MAX;
107 if (!outArray.Contains(idx)) {
108 outArray.InsertElementSorted(idx);
109 }
110 }
111 }
112
113 // Function to generate TableUpdateV4.
GenerateUpdateData(bool fullUpdate,PrefixStringMap & add,nsTArray<uint32_t> * removal,nsCString * checksum,nsTArray<TableUpdate * > & tableUpdates)114 static void GenerateUpdateData(bool fullUpdate, PrefixStringMap& add,
115 nsTArray<uint32_t>* removal, nsCString* checksum,
116 nsTArray<TableUpdate*>& tableUpdates) {
117 TableUpdateV4* tableUpdate = new TableUpdateV4(GTEST_TABLE);
118 tableUpdate->SetFullUpdate(fullUpdate);
119
120 for (auto iter = add.ConstIter(); !iter.Done(); iter.Next()) {
121 nsCString* pstring = iter.Data();
122 std::string str(pstring->BeginReading(), pstring->Length());
123
124 tableUpdate->NewPrefixes(iter.Key(), str);
125 }
126
127 if (removal) {
128 tableUpdate->NewRemovalIndices(removal->Elements(), removal->Length());
129 }
130
131 if (checksum) {
132 std::string stdChecksum;
133 stdChecksum.assign(const_cast<char*>(checksum->BeginReading()),
134 checksum->Length());
135
136 tableUpdate->NewChecksum(stdChecksum);
137 }
138
139 tableUpdates.AppendElement(tableUpdate);
140 }
141
VerifyPrefixSet(PrefixStringMap & expected)142 static void VerifyPrefixSet(PrefixStringMap& expected) {
143 // Verify the prefix set is written to disk.
144 nsCOMPtr<nsIFile> file;
145 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
146
147 file->AppendNative(GTEST_SAFEBROWSING_DIR);
148 file->AppendNative(GTEST_PREFIXFILE);
149
150 RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;
151 load->Init(GTEST_TABLE);
152
153 PrefixStringMap prefixesInFile;
154 load->LoadFromFile(file);
155 load->GetPrefixes(prefixesInFile);
156
157 for (auto iter = expected.ConstIter(); !iter.Done(); iter.Next()) {
158 nsCString* expectedPrefix = iter.Data();
159 nsCString* resultPrefix = prefixesInFile.Get(iter.Key());
160
161 ASSERT_TRUE(*resultPrefix == *expectedPrefix);
162 }
163 }
164
Clear()165 static void Clear() {
166 nsCOMPtr<nsIFile> file;
167 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
168
169 UniquePtr<Classifier> classifier(new Classifier());
170 classifier->Open(*file);
171 classifier->Reset();
172 }
173
testUpdateFail(nsTArray<TableUpdate * > & tableUpdates)174 static void testUpdateFail(nsTArray<TableUpdate*>& tableUpdates) {
175 nsCOMPtr<nsIFile> file;
176 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
177
178 UniquePtr<Classifier> classifier(new Classifier());
179 classifier->Open(*file);
180
181 nsresult rv = SyncApplyUpdates(classifier.get(), &tableUpdates);
182 ASSERT_TRUE(NS_FAILED(rv));
183 }
184
testUpdate(nsTArray<TableUpdate * > & tableUpdates,PrefixStringMap & expected)185 static void testUpdate(nsTArray<TableUpdate*>& tableUpdates,
186 PrefixStringMap& expected) {
187 nsCOMPtr<nsIFile> file;
188 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
189
190 {
191 // Force nsIUrlClassifierUtils loading on main thread
192 // because nsIUrlClassifierDBService will not run in advance
193 // in gtest.
194 nsresult rv;
195 nsCOMPtr<nsIUrlClassifierUtils> dummy =
196 do_GetService(NS_URLCLASSIFIERUTILS_CONTRACTID, &rv);
197 ASSERT_TRUE(NS_SUCCEEDED(rv));
198 }
199
200 UniquePtr<Classifier> classifier(new Classifier());
201 classifier->Open(*file);
202
203 nsresult rv = SyncApplyUpdates(classifier.get(), &tableUpdates);
204 ASSERT_TRUE(rv == NS_OK);
205 VerifyPrefixSet(expected);
206 }
207
testFullUpdate(PrefixStringMap & add,nsCString * checksum)208 static void testFullUpdate(PrefixStringMap& add, nsCString* checksum) {
209 nsTArray<TableUpdate*> tableUpdates;
210
211 GenerateUpdateData(true, add, nullptr, checksum, tableUpdates);
212
213 testUpdate(tableUpdates, add);
214 }
215
testPartialUpdate(PrefixStringMap & add,nsTArray<uint32_t> * removal,nsCString * checksum,PrefixStringMap & expected)216 static void testPartialUpdate(PrefixStringMap& add, nsTArray<uint32_t>* removal,
217 nsCString* checksum, PrefixStringMap& expected) {
218 nsTArray<TableUpdate*> tableUpdates;
219 GenerateUpdateData(false, add, removal, checksum, tableUpdates);
220
221 testUpdate(tableUpdates, expected);
222 }
223
testOpenLookupCache()224 static void testOpenLookupCache() {
225 nsCOMPtr<nsIFile> file;
226 NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
227 file->AppendNative(GTEST_SAFEBROWSING_DIR);
228
229 RunTestInNewThread([&]() -> void {
230 LookupCacheV4 cache(nsCString(GTEST_TABLE), EmptyCString(), file);
231 nsresult rv = cache.Init();
232 ASSERT_EQ(rv, NS_OK);
233
234 rv = cache.Open();
235 ASSERT_EQ(rv, NS_OK);
236 });
237 }
238
239 // Tests start from here.
TEST(UrlClassifierTableUpdateV4,FixLenghtPSetFullUpdate)240 TEST(UrlClassifierTableUpdateV4, FixLenghtPSetFullUpdate) {
241 srand(time(NULL));
242
243 _PrefixArray array;
244 PrefixStringMap map;
245 nsCString checksum;
246
247 CreateRandomSortedPrefixArray(5000, 4, 4, array);
248 PrefixArrayToPrefixStringMap(array, map);
249 CalculateCheckSum(array, checksum);
250
251 testFullUpdate(map, &checksum);
252
253 Clear();
254 }
255
TEST(UrlClassifierTableUpdateV4,VariableLenghtPSetFullUpdate)256 TEST(UrlClassifierTableUpdateV4, VariableLenghtPSetFullUpdate) {
257 _PrefixArray array;
258 PrefixStringMap map;
259 nsCString checksum;
260
261 CreateRandomSortedPrefixArray(5000, 5, 32, array);
262 PrefixArrayToPrefixStringMap(array, map);
263 CalculateCheckSum(array, checksum);
264
265 testFullUpdate(map, &checksum);
266
267 Clear();
268 }
269
270 // This test contain both variable length prefix set and fixed-length prefix set
TEST(UrlClassifierTableUpdateV4,MixedPSetFullUpdate)271 TEST(UrlClassifierTableUpdateV4, MixedPSetFullUpdate) {
272 _PrefixArray array;
273 PrefixStringMap map;
274 nsCString checksum;
275
276 CreateRandomSortedPrefixArray(5000, 4, 4, array);
277 CreateRandomSortedPrefixArray(1000, 5, 32, array);
278 PrefixArrayToPrefixStringMap(array, map);
279 CalculateCheckSum(array, checksum);
280
281 testFullUpdate(map, &checksum);
282
283 Clear();
284 }
285
TEST(UrlClassifierTableUpdateV4,PartialUpdateWithRemoval)286 TEST(UrlClassifierTableUpdateV4, PartialUpdateWithRemoval) {
287 _PrefixArray fArray;
288
289 // Apply a full update first.
290 {
291 PrefixStringMap fMap;
292 nsCString checksum;
293
294 CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
295 CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
296 PrefixArrayToPrefixStringMap(fArray, fMap);
297 CalculateCheckSum(fArray, checksum);
298
299 testFullUpdate(fMap, &checksum);
300 }
301
302 // Apply a partial update with removal.
303 {
304 _PrefixArray pArray, mergedArray;
305 PrefixStringMap pMap, mergedMap;
306 nsCString checksum;
307
308 CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
309 CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
310 RemoveIntersection(fArray, pArray);
311 PrefixArrayToPrefixStringMap(pArray, pMap);
312
313 // Remove 1/5 of elements of original prefix set.
314 nsTArray<uint32_t> removal;
315 CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
316 RemoveElements(removal, fArray);
317
318 // Calculate the expected prefix map.
319 MergeAndSortArray(fArray, pArray, mergedArray);
320 PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
321 CalculateCheckSum(mergedArray, checksum);
322
323 testPartialUpdate(pMap, &removal, &checksum, mergedMap);
324 }
325
326 Clear();
327 }
328
TEST(UrlClassifierTableUpdateV4,PartialUpdateWithoutRemoval)329 TEST(UrlClassifierTableUpdateV4, PartialUpdateWithoutRemoval) {
330 _PrefixArray fArray;
331
332 // Apply a full update first.
333 {
334 PrefixStringMap fMap;
335 nsCString checksum;
336
337 CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
338 CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
339 PrefixArrayToPrefixStringMap(fArray, fMap);
340 CalculateCheckSum(fArray, checksum);
341
342 testFullUpdate(fMap, &checksum);
343 }
344
345 // Apply a partial update without removal
346 {
347 _PrefixArray pArray, mergedArray;
348 PrefixStringMap pMap, mergedMap;
349 nsCString checksum;
350
351 CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
352 CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
353 RemoveIntersection(fArray, pArray);
354 PrefixArrayToPrefixStringMap(pArray, pMap);
355
356 // Calculate the expected prefix map.
357 MergeAndSortArray(fArray, pArray, mergedArray);
358 PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
359 CalculateCheckSum(mergedArray, checksum);
360
361 testPartialUpdate(pMap, nullptr, &checksum, mergedMap);
362 }
363
364 Clear();
365 }
366
367 // Expect failure because partial update contains prefix already
368 // in old prefix set.
TEST(UrlClassifierTableUpdateV4,PartialUpdatePrefixAlreadyExist)369 TEST(UrlClassifierTableUpdateV4, PartialUpdatePrefixAlreadyExist) {
370 _PrefixArray fArray;
371
372 // Apply a full update fist.
373 {
374 PrefixStringMap fMap;
375 nsCString checksum;
376
377 CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
378 PrefixArrayToPrefixStringMap(fArray, fMap);
379 CalculateCheckSum(fArray, checksum);
380
381 testFullUpdate(fMap, &checksum);
382 }
383
384 // Apply a partial update which contains a prefix in previous full update.
385 // This should cause an update error.
386 {
387 _PrefixArray pArray;
388 PrefixStringMap pMap;
389 nsTArray<TableUpdate*> tableUpdates;
390
391 // Pick one prefix from full update prefix and add it to partial update.
392 // This should result a failure when call ApplyUpdates.
393 pArray.AppendElement(fArray[rand() % fArray.Length()]);
394 CreateRandomSortedPrefixArray(200, 4, 32, pArray);
395 PrefixArrayToPrefixStringMap(pArray, pMap);
396
397 GenerateUpdateData(false, pMap, nullptr, nullptr, tableUpdates);
398 testUpdateFail(tableUpdates);
399 }
400
401 Clear();
402 }
403
404 // Test apply partial update directly without applying an full update first.
TEST(UrlClassifierTableUpdateV4,OnlyPartialUpdate)405 TEST(UrlClassifierTableUpdateV4, OnlyPartialUpdate) {
406 _PrefixArray pArray;
407 PrefixStringMap pMap;
408 nsCString checksum;
409
410 CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
411 CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
412 PrefixArrayToPrefixStringMap(pArray, pMap);
413 CalculateCheckSum(pArray, checksum);
414
415 testPartialUpdate(pMap, nullptr, &checksum, pMap);
416
417 Clear();
418 }
419
420 // Test partial update without any ADD prefixes, only removalIndices.
TEST(UrlClassifierTableUpdateV4,PartialUpdateOnlyRemoval)421 TEST(UrlClassifierTableUpdateV4, PartialUpdateOnlyRemoval) {
422 _PrefixArray fArray;
423
424 // Apply a full update first.
425 {
426 PrefixStringMap fMap;
427 nsCString checksum;
428
429 CreateRandomSortedPrefixArray(5000, 4, 4, fArray);
430 CreateRandomSortedPrefixArray(1000, 5, 32, fArray);
431 PrefixArrayToPrefixStringMap(fArray, fMap);
432 CalculateCheckSum(fArray, checksum);
433
434 testFullUpdate(fMap, &checksum);
435 }
436
437 // Apply a partial update without add prefix, only contain removal indices.
438 {
439 _PrefixArray pArray;
440 PrefixStringMap pMap, mergedMap;
441 nsCString checksum;
442
443 // Remove 1/5 of elements of original prefix set.
444 nsTArray<uint32_t> removal;
445 CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
446 RemoveElements(removal, fArray);
447
448 PrefixArrayToPrefixStringMap(fArray, mergedMap);
449 CalculateCheckSum(fArray, checksum);
450
451 testPartialUpdate(pMap, &removal, &checksum, mergedMap);
452 }
453
454 Clear();
455 }
456
457 // Test one tableupdate array contains full update and multiple partial updates.
TEST(UrlClassifierTableUpdateV4,MultipleTableUpdates)458 TEST(UrlClassifierTableUpdateV4, MultipleTableUpdates) {
459 _PrefixArray fArray, pArray, mergedArray;
460 PrefixStringMap fMap, pMap, mergedMap;
461 nsCString checksum;
462
463 nsTArray<TableUpdate*> tableUpdates;
464
465 // Generate first full udpate
466 CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
467 CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
468 PrefixArrayToPrefixStringMap(fArray, fMap);
469 CalculateCheckSum(fArray, checksum);
470
471 GenerateUpdateData(true, fMap, nullptr, &checksum, tableUpdates);
472
473 // Generate second partial update
474 CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
475 CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
476 RemoveIntersection(fArray, pArray);
477 PrefixArrayToPrefixStringMap(pArray, pMap);
478
479 MergeAndSortArray(fArray, pArray, mergedArray);
480 CalculateCheckSum(mergedArray, checksum);
481
482 GenerateUpdateData(false, pMap, nullptr, &checksum, tableUpdates);
483
484 // Generate thrid partial update
485 fArray.AppendElements(pArray);
486 fArray.Sort();
487 pArray.Clear();
488 CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
489 CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
490 RemoveIntersection(fArray, pArray);
491 PrefixArrayToPrefixStringMap(pArray, pMap);
492
493 // Remove 1/5 of elements of original prefix set.
494 nsTArray<uint32_t> removal;
495 CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
496 RemoveElements(removal, fArray);
497
498 MergeAndSortArray(fArray, pArray, mergedArray);
499 PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
500 CalculateCheckSum(mergedArray, checksum);
501
502 GenerateUpdateData(false, pMap, &removal, &checksum, tableUpdates);
503
504 testUpdate(tableUpdates, mergedMap);
505
506 Clear();
507 }
508
509 // Test apply full update first, and then apply multiple partial updates
510 // in one tableupdate array.
TEST(UrlClassifierTableUpdateV4,MultiplePartialUpdateTableUpdates)511 TEST(UrlClassifierTableUpdateV4, MultiplePartialUpdateTableUpdates) {
512 _PrefixArray fArray;
513
514 // Apply a full update first
515 {
516 PrefixStringMap fMap;
517 nsCString checksum;
518
519 // Generate first full udpate
520 CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
521 CreateRandomSortedPrefixArray(3000, 5, 32, fArray);
522 PrefixArrayToPrefixStringMap(fArray, fMap);
523 CalculateCheckSum(fArray, checksum);
524
525 testFullUpdate(fMap, &checksum);
526 }
527
528 // Apply multiple partial updates in one table update
529 {
530 _PrefixArray pArray, mergedArray;
531 PrefixStringMap pMap, mergedMap;
532 nsCString checksum;
533 nsTArray<uint32_t> removal;
534 nsTArray<TableUpdate*> tableUpdates;
535
536 // Generate first partial update
537 CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
538 CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
539 RemoveIntersection(fArray, pArray);
540 PrefixArrayToPrefixStringMap(pArray, pMap);
541
542 // Remove 1/5 of elements of original prefix set.
543 CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
544 RemoveElements(removal, fArray);
545
546 MergeAndSortArray(fArray, pArray, mergedArray);
547 CalculateCheckSum(mergedArray, checksum);
548
549 GenerateUpdateData(false, pMap, &removal, &checksum, tableUpdates);
550
551 fArray.AppendElements(pArray);
552 fArray.Sort();
553 pArray.Clear();
554 removal.Clear();
555
556 // Generate second partial update.
557 CreateRandomSortedPrefixArray(2000, 4, 4, pArray);
558 CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
559 RemoveIntersection(fArray, pArray);
560 PrefixArrayToPrefixStringMap(pArray, pMap);
561
562 // Remove 1/5 of elements of original prefix set.
563 CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
564 RemoveElements(removal, fArray);
565
566 MergeAndSortArray(fArray, pArray, mergedArray);
567 PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
568 CalculateCheckSum(mergedArray, checksum);
569
570 GenerateUpdateData(false, pMap, &removal, &checksum, tableUpdates);
571
572 testUpdate(tableUpdates, mergedMap);
573 }
574
575 Clear();
576 }
577
578 // Test removal indices are larger than the original prefix set.
TEST(UrlClassifierTableUpdateV4,RemovalIndexTooLarge)579 TEST(UrlClassifierTableUpdateV4, RemovalIndexTooLarge) {
580 _PrefixArray fArray;
581
582 // Apply a full update first
583 {
584 PrefixStringMap fMap;
585 nsCString checksum;
586
587 CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
588 PrefixArrayToPrefixStringMap(fArray, fMap);
589 CalculateCheckSum(fArray, checksum);
590
591 testFullUpdate(fMap, &checksum);
592 }
593
594 // Apply a partial update with removal indice array larger than
595 // old prefix set(fArray). This should cause an error.
596 {
597 _PrefixArray pArray;
598 PrefixStringMap pMap;
599 nsTArray<uint32_t> removal;
600 nsTArray<TableUpdate*> tableUpdates;
601
602 CreateRandomSortedPrefixArray(200, 4, 32, pArray);
603 RemoveIntersection(fArray, pArray);
604 PrefixArrayToPrefixStringMap(pArray, pMap);
605
606 for (uint32_t i = 0; i < fArray.Length() + 1; i++) {
607 removal.AppendElement(i);
608 }
609
610 GenerateUpdateData(false, pMap, &removal, nullptr, tableUpdates);
611 testUpdateFail(tableUpdates);
612 }
613
614 Clear();
615 }
616
TEST(UrlClassifierTableUpdateV4,ChecksumMismatch)617 TEST(UrlClassifierTableUpdateV4, ChecksumMismatch) {
618 // Apply a full update first
619 {
620 _PrefixArray fArray;
621 PrefixStringMap fMap;
622 nsCString checksum;
623
624 CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
625 PrefixArrayToPrefixStringMap(fArray, fMap);
626 CalculateCheckSum(fArray, checksum);
627
628 testFullUpdate(fMap, &checksum);
629 }
630
631 // Apply a partial update with incorrect checksum
632 {
633 _PrefixArray pArray;
634 PrefixStringMap pMap;
635 nsCString checksum;
636 nsTArray<TableUpdate*> tableUpdates;
637
638 CreateRandomSortedPrefixArray(200, 4, 32, pArray);
639 PrefixArrayToPrefixStringMap(pArray, pMap);
640
641 // Checksum should be calculated with both old prefix set and add prefix
642 // set, here we only calculate checksum with add prefix set to check if
643 // applyUpdate will return failure.
644 CalculateCheckSum(pArray, checksum);
645
646 GenerateUpdateData(false, pMap, nullptr, &checksum, tableUpdates);
647 testUpdateFail(tableUpdates);
648 }
649
650 Clear();
651 }
652
TEST(UrlClassifierTableUpdateV4,ApplyUpdateThenLoad)653 TEST(UrlClassifierTableUpdateV4, ApplyUpdateThenLoad) {
654 // Apply update with checksum
655 {
656 _PrefixArray fArray;
657 PrefixStringMap fMap;
658 nsCString checksum;
659
660 CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
661 PrefixArrayToPrefixStringMap(fArray, fMap);
662 CalculateCheckSum(fArray, checksum);
663
664 testFullUpdate(fMap, &checksum);
665
666 // Open lookup cache will load prefix set and verify the checksum
667 testOpenLookupCache();
668 }
669
670 Clear();
671
672 // Apply update without checksum
673 {
674 _PrefixArray fArray;
675 PrefixStringMap fMap;
676
677 CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
678 PrefixArrayToPrefixStringMap(fArray, fMap);
679
680 testFullUpdate(fMap, nullptr);
681
682 testOpenLookupCache();
683 }
684
685 Clear();
686 }
687
688 // This test is used to avoid an eror from nsICryptoHash
TEST(UrlClassifierTableUpdateV4,ApplyUpdateWithFixedChecksum)689 TEST(UrlClassifierTableUpdateV4, ApplyUpdateWithFixedChecksum) {
690 _PrefixArray fArray = {_Prefix("enus"),
691 _Prefix("apollo"),
692 _Prefix("mars"),
693 _Prefix("Hecatonchires cyclopes"),
694 _Prefix("vesta"),
695 _Prefix("neptunus"),
696 _Prefix("jupiter"),
697 _Prefix("diana"),
698 _Prefix("minerva"),
699 _Prefix("ceres"),
700 _Prefix("Aidos,Adephagia,Adikia,Aletheia"),
701 _Prefix("hecatonchires"),
702 _Prefix("alcyoneus"),
703 _Prefix("hades"),
704 _Prefix("vulcanus"),
705 _Prefix("juno"),
706 _Prefix("mercury"),
707 _Prefix("Stheno, Euryale and Medusa")};
708 fArray.Sort();
709
710 PrefixStringMap fMap;
711 PrefixArrayToPrefixStringMap(fArray, fMap);
712
713 nsCString checksum(
714 "\xae\x18\x94\xd7\xd0\x83\x5f\xc1"
715 "\x58\x59\x5c\x2c\x72\xb9\x6e\x5e"
716 "\xf4\xe8\x0a\x6b\xff\x5e\x6b\x81"
717 "\x65\x34\x06\x16\x06\x59\xa0\x67");
718
719 testFullUpdate(fMap, &checksum);
720
721 // Open lookup cache will load prefix set and verify the checksum
722 testOpenLookupCache();
723
724 Clear();
725 }
726
727 // This test ensure that an empty update works correctly. Empty update
728 // should be skipped by CheckValidUpdate in Classifier::UpdateTableV4.
TEST(UrlClassifierTableUpdateV4,EmptyUpdate)729 TEST(UrlClassifierTableUpdateV4, EmptyUpdate) {
730 PrefixStringMap emptyAddition;
731 nsTArray<uint32_t> emptyRemoval;
732
733 _PrefixArray array;
734 PrefixStringMap map;
735 nsCString checksum;
736
737 CalculateCheckSum(array, checksum);
738
739 // Test apply empty full/partial update before we already
740 // have data in DB.
741 testFullUpdate(emptyAddition, &checksum);
742 testPartialUpdate(emptyAddition, &emptyRemoval, &checksum, map);
743
744 // Apply an full update.
745 CreateRandomSortedPrefixArray(100, 4, 4, array);
746 CreateRandomSortedPrefixArray(10, 5, 32, array);
747 PrefixArrayToPrefixStringMap(array, map);
748 CalculateCheckSum(array, checksum);
749
750 testFullUpdate(map, &checksum);
751
752 // Test apply empty full/partial update when we already
753 // have data in DB
754 testPartialUpdate(emptyAddition, &emptyRemoval, &checksum, map);
755 testFullUpdate(emptyAddition, &checksum);
756
757 Clear();
758 }
759
760 // This test ensure applying an empty update directly through update algorithm
761 // should be correct.
TEST(UrlClassifierTableUpdateV4,EmptyUpdate2)762 TEST(UrlClassifierTableUpdateV4, EmptyUpdate2) {
763 // Setup LookupCache with initial data
764 _PrefixArray array;
765 CreateRandomSortedPrefixArray(100, 4, 4, array);
766 CreateRandomSortedPrefixArray(10, 5, 32, array);
767 UniquePtr<LookupCacheV4> cache = SetupLookupCache<LookupCacheV4>(array);
768
769 // Setup TableUpdate object with only checksum from previous update(initial
770 // data).
771 nsCString checksum;
772 CalculateCheckSum(array, checksum);
773 std::string stdChecksum;
774 stdChecksum.assign(const_cast<char*>(checksum.BeginReading()),
775 checksum.Length());
776
777 UniquePtr<TableUpdateV4> tableUpdate = MakeUnique<TableUpdateV4>(GTEST_TABLE);
778 tableUpdate->NewChecksum(stdChecksum);
779
780 // Apply update directly through LookupCache interface
781 PrefixStringMap input, output;
782 PrefixArrayToPrefixStringMap(array, input);
783 nsresult rv = cache->ApplyUpdate(tableUpdate.get(), input, output);
784
785 ASSERT_TRUE(rv == NS_OK);
786
787 Clear();
788 }
789