Bug 1353956 - P7. Add gtest to ensure .pset is correctly loaded and removed. r=gcp

Differential Revision: https://phabricator.services.mozilla.com/D22490

--HG--
extra : moz-landing-system : lando

toolkit/components/url-classifier/tests/gtest/Common.cpp

@@ -147,6 +147,18 @@ nsCString GeneratePrefix(const nsCString& aFragment, uint8_t aLength) {
   return hash;
 }
 
+void CheckContent(LookupCacheV4* cache, PrefixStringMap& expected) {
+  PrefixStringMap vlPSetMap;
+  cache->GetPrefixes(vlPSetMap);
+
+  for (auto iter = vlPSetMap.Iter(); !iter.Done(); iter.Next()) {
+    nsCString* expectedPrefix = expected.Get(iter.Key());
+    nsCString* resultPrefix = iter.Data();
+
+    ASSERT_TRUE(resultPrefix->Equals(*expectedPrefix));
+  }
+}
+
 static nsresult BuildCache(LookupCacheV2* cache,
                            const _PrefixArray& prefixArray) {
   AddPrefixArray prefixes;

toolkit/components/url-classifier/tests/gtest/Common.h

@@ -10,6 +10,7 @@ using namespace mozilla::safebrowsing;
 namespace mozilla {
 namespace safebrowsing {
 class Classifier;
+class LookupCacheV4;
 }
 }  // namespace mozilla
 
@@ -43,6 +44,9 @@ nsresult PrefixArrayToAddPrefixArrayV2(const nsTArray<nsCString>& prefixArray,
 // Generate a hash prefix from string
 nsCString GeneratePrefix(const nsCString& aFragment, uint8_t aLength);
 
+// To test if the content is equal
+void CheckContent(LookupCacheV4* cache, PrefixStringMap& expected);
+
 // Create a LookupCacheV4 object with sepecified prefix array.
 template <typename T>
 RefPtr<T> SetupLookupCache(const _PrefixArray& prefixArray);

toolkit/components/url-classifier/tests/gtest/TestLookupCacheV4.cpp

@@ -43,3 +43,85 @@ TEST(UrlClassifierLookupCacheV4, HasPrefix) {
 TEST(UrlClassifierLookupCacheV4, Nomatch) {
   TestHasPrefix(_Fragment("nomatch.com/"), false, false);
 }
+
+// Test an existing .pset should be removed after .vlpset is written
+TEST(UrlClassifierLookupCacheV4, RemoveOldPset) {
+  nsCOMPtr<nsIFile> oldPsetFile;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                         getter_AddRefs(oldPsetFile));
+  oldPsetFile->AppendNative(NS_LITERAL_CSTRING("safebrowsing"));
+  oldPsetFile->AppendNative(GTEST_TABLE + NS_LITERAL_CSTRING(".pset"));
+
+  nsCOMPtr<nsIFile> newPsetFile;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                         getter_AddRefs(newPsetFile));
+  newPsetFile->AppendNative(NS_LITERAL_CSTRING("safebrowsing"));
+  newPsetFile->AppendNative(GTEST_TABLE + NS_LITERAL_CSTRING(".vlpset"));
+
+  // Create the legacy .pset file
+  nsresult rv = oldPsetFile->Create(nsIFile::NORMAL_FILE_TYPE, 0666);
+  EXPECT_EQ(rv, NS_OK);
+
+  bool exists;
+  rv = oldPsetFile->Exists(&exists);
+  EXPECT_EQ(rv, NS_OK);
+  EXPECT_EQ(exists, true);
+
+  // Setup the data in lookup cache and write its data to disk
+  RefPtr<Classifier> classifier = GetClassifier();
+  _PrefixArray array = {GeneratePrefix(_Fragment("entry.com/"), 4)};
+  rv = SetupLookupCacheV4(classifier, array, GTEST_TABLE);
+  EXPECT_EQ(rv, NS_OK);
+
+  RefPtr<LookupCache> cache = classifier->GetLookupCache(GTEST_TABLE, false);
+  rv = cache->WriteFile();
+  EXPECT_EQ(rv, NS_OK);
+
+  // .vlpset should exist while .pset should be removed
+  rv = newPsetFile->Exists(&exists);
+  EXPECT_EQ(rv, NS_OK);
+  EXPECT_EQ(exists, true);
+
+  rv = oldPsetFile->Exists(&exists);
+  EXPECT_EQ(rv, NS_OK);
+  EXPECT_EQ(exists, false);
+
+  newPsetFile->Remove(false);
+}
+
+// Test the legacy load
+TEST(UrlClassifierLookupCacheV4, LoadOldPset) {
+  nsCOMPtr<nsIFile> oldPsetFile;
+
+  _PrefixArray array = {GeneratePrefix(_Fragment("entry.com/"), 4)};
+  PrefixStringMap map;
+  PrefixArrayToPrefixStringMap(array, map);
+
+  // Prepare .pset file on disk
+  {
+    NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                           getter_AddRefs(oldPsetFile));
+    oldPsetFile->AppendNative(NS_LITERAL_CSTRING("safebrowsing"));
+    oldPsetFile->AppendNative(GTEST_TABLE + NS_LITERAL_CSTRING(".pset"));
+
+    RefPtr<VariableLengthPrefixSet> pset = new VariableLengthPrefixSet;
+    pset->SetPrefixes(map);
+
+    nsCOMPtr<nsIOutputStream> stream;
+    nsresult rv =
+        NS_NewLocalFileOutputStream(getter_AddRefs(stream), oldPsetFile);
+    EXPECT_EQ(rv, NS_OK);
+
+    rv = pset->WritePrefixes(stream);
+    EXPECT_EQ(rv, NS_OK);
+  }
+
+  // Load data from disk
+  RefPtr<Classifier> classifier = GetClassifier();
+  RefPtr<LookupCache> cache = classifier->GetLookupCache(GTEST_TABLE, false);
+
+  RefPtr<LookupCacheV4> cacheV4 = LookupCache::Cast<LookupCacheV4>(cache);
+  CheckContent(cacheV4, map);
+
+  oldPsetFile->Remove(false);
+}

toolkit/components/url-classifier/tests/gtest/TestVariableLengthPrefixSet.cpp

@@ -55,18 +55,6 @@ static void RandomPrefixes(uint32_t N, uint32_t MIN, uint32_t MAX,
   }
 }
 
-static void CheckContent(LookupCacheV4* cache, PrefixStringMap& expected) {
-  PrefixStringMap vlPSetMap;
-  cache->GetPrefixes(vlPSetMap);
-
-  for (auto iter = vlPSetMap.Iter(); !iter.Done(); iter.Next()) {
-    nsCString* expectedPrefix = expected.Get(iter.Key());
-    nsCString* resultPrefix = iter.Data();
-
-    ASSERT_TRUE(resultPrefix->Equals(*expectedPrefix));
-  }
-}
-
 // This test loops through all the prefixes and converts each prefix to
 // fullhash by appending random characters, each converted fullhash
 // should at least match its original length in the prefixSet.