1 files changed, 152 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/gtest/TestLookupCacheV4.cpp b/toolkit/components/url-classifier/tests/gtest/TestLookupCacheV4.cpp
new file mode 100644
index 0000000000..e1482c1416
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/gtest/TestLookupCacheV4.cpp
@@ -0,0 +1,152 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "Classifier.h"
+#include "LookupCacheV4.h"
+#include "nsAppDirectoryServiceDefs.h"
+
+#include "Common.h"
+
+#define GTEST_SAFEBROWSING_DIR "safebrowsing"_ns
+
+static void TestHasPrefix(const nsCString& aURL, bool aExpectedHas,
+                          bool aExpectedComplete) {
+  _PrefixArray array = {CreatePrefixFromURL("bravo.com/", 32),
+                        CreatePrefixFromURL("browsing.com/", 8),
+                        CreatePrefixFromURL("gound.com/", 5),
+                        CreatePrefixFromURL("small.com/", 4)};
+
+  nsCOMPtr<nsIFile> file;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
+  file->AppendNative(GTEST_SAFEBROWSING_DIR);
+
+  RunTestInNewThread([&]() -> void {
+    RefPtr<LookupCache> cache = SetupLookupCache<LookupCacheV4>(array, file);
+
+    Completion lookupHash;
+    lookupHash.FromPlaintext(aURL);
+
+    bool has, confirmed;
+    uint32_t matchLength;
+    // Freshness is not used in V4 so we just put dummy values here.
+    TableFreshnessMap dummy;
+    nsresult rv = cache->Has(lookupHash, &has, &matchLength, &confirmed);
+
+    EXPECT_EQ(rv, NS_OK);
+    EXPECT_EQ(has, aExpectedHas);
+    EXPECT_EQ(matchLength == COMPLETE_SIZE, aExpectedComplete);
+    EXPECT_EQ(confirmed, false);
+
+    cache->ClearAll();
+  });
+}
+
+TEST(UrlClassifierLookupCacheV4, HasComplete)
+{ TestHasPrefix("bravo.com/"_ns, true, true); }
+
+TEST(UrlClassifierLookupCacheV4, HasPrefix)
+{ TestHasPrefix("browsing.com/"_ns, true, false); }
+
+TEST(UrlClassifierLookupCacheV4, Nomatch)
+{ TestHasPrefix("nomatch.com/"_ns, 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("safebrowsing"_ns);
+  oldPsetFile->AppendNative(GTEST_TABLE_V4 + ".pset"_ns);
+
+  nsCOMPtr<nsIFile> newPsetFile;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                         getter_AddRefs(newPsetFile));
+  newPsetFile->AppendNative("safebrowsing"_ns);
+  newPsetFile->AppendNative(GTEST_TABLE_V4 + ".vlpset"_ns);
+
+  // 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 = {CreatePrefixFromURL("entry.com/", 4)};
+  rv = BuildLookupCache(classifier, GTEST_TABLE_V4, array);
+  EXPECT_EQ(rv, NS_OK);
+
+  RefPtr<LookupCache> cache = classifier->GetLookupCache(GTEST_TABLE_V4, 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 = {CreatePrefixFromURL("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("safebrowsing"_ns);
+    oldPsetFile->AppendNative(GTEST_TABLE_V4 + ".pset"_ns);
+
+    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_V4, false);
+
+  RefPtr<LookupCacheV4> cacheV4 = LookupCache::Cast<LookupCacheV4>(cache);
+  CheckContent(cacheV4, array);
+
+  oldPsetFile->Remove(false);
+}
+
+TEST(UrlClassifierLookupCacheV4, BuildAPI)
+{
+  _PrefixArray init = {_Prefix("alph")};
+  RefPtr<LookupCacheV4> cache = SetupLookupCache<LookupCacheV4>(init);
+
+  _PrefixArray update = {_Prefix("beta")};
+  PrefixStringMap map;
+  PrefixArrayToPrefixStringMap(update, map);
+
+  cache->Build(map);
+  EXPECT_TRUE(map.IsEmpty());
+
+  CheckContent(cache, update);
+}