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