1 files changed, 486 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/gtest/TestVariableLengthPrefixSet.cpp b/toolkit/components/url-classifier/tests/gtest/TestVariableLengthPrefixSet.cpp
new file mode 100644
index 0000000000..6eb36a12d2
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/gtest/TestVariableLengthPrefixSet.cpp
@@ -0,0 +1,486 @@
+/* -*- 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 "LookupCacheV4.h"
+#include "mozilla/Preferences.h"
+#include <mozilla/RefPtr.h>
+#include "nsAppDirectoryServiceDefs.h"
+#include "nsClassHashtable.h"
+#include "nsString.h"
+#include "VariableLengthPrefixSet.h"
+
+#include "Common.h"
+
+// Create fullhash by appending random characters.
+static nsCString CreateFullHash(const nsACString& in) {
+  nsCString out(in);
+  out.SetLength(32);
+  for (size_t i = in.Length(); i < 32; i++) {
+    out.SetCharAt(char(rand() % 256), i);
+  }
+
+  return out;
+}
+
+// This function generate N prefixes with size between MIN and MAX.
+// The output array will not be cleared, random result will append to it
+static void RandomPrefixes(uint32_t N, uint32_t MIN, uint32_t MAX,
+                           _PrefixArray& array) {
+  array.SetCapacity(array.Length() + N);
+
+  uint32_t range = (MAX - MIN + 1);
+
+  for (uint32_t i = 0; i < N; i++) {
+    uint32_t prefixSize = (rand() % range) + MIN;
+    _Prefix prefix;
+    prefix.SetLength(prefixSize);
+
+    bool added = false;
+    while (!added) {
+      char* dst = prefix.BeginWriting();
+      for (uint32_t j = 0; j < prefixSize; j++) {
+        dst[j] = rand() % 256;
+      }
+
+      if (!array.Contains(prefix)) {
+        array.AppendElement(prefix);
+        added = true;
+      }
+    }
+  }
+}
+
+// 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.
+static void DoExpectedLookup(LookupCacheV4* cache, _PrefixArray& array) {
+  uint32_t matchLength = 0;
+  for (uint32_t i = 0; i < array.Length(); i++) {
+    const nsCString& prefix = array[i];
+    Completion complete;
+    complete.Assign(CreateFullHash(prefix));
+
+    // Find match for prefix-generated full hash
+    bool has, confirmed;
+    cache->Has(complete, &has, &matchLength, &confirmed);
+    MOZ_ASSERT(matchLength != 0);
+
+    if (matchLength != prefix.Length()) {
+      // Return match size is not the same as prefix size.
+      // In this case it could be because the generated fullhash match other
+      // prefixes, check if this prefix exist.
+      bool found = false;
+
+      for (uint32_t j = 0; j < array.Length(); j++) {
+        if (array[j].Length() != matchLength) {
+          continue;
+        }
+
+        if (0 == memcmp(complete.buf, array[j].BeginReading(), matchLength)) {
+          found = true;
+          break;
+        }
+      }
+      ASSERT_TRUE(found);
+    }
+  }
+}
+
+static void DoRandomLookup(LookupCacheV4* cache, uint32_t N,
+                           _PrefixArray& array) {
+  for (uint32_t i = 0; i < N; i++) {
+    // Random 32-bytes test fullhash
+    char buf[32];
+    for (uint32_t j = 0; j < 32; j++) {
+      buf[j] = (char)(rand() % 256);
+    }
+
+    // Get the expected result.
+    nsTArray<uint32_t> expected;
+    for (uint32_t j = 0; j < array.Length(); j++) {
+      const nsACString& str = array[j];
+      if (0 == memcmp(buf, str.BeginReading(), str.Length())) {
+        expected.AppendElement(str.Length());
+      }
+    }
+
+    Completion complete;
+    complete.Assign(nsDependentCSubstring(buf, 32));
+    bool has, confirmed;
+    uint32_t matchLength = 0;
+    cache->Has(complete, &has, &matchLength, &confirmed);
+
+    ASSERT_TRUE(expected.IsEmpty() ? !matchLength
+                                   : expected.Contains(matchLength));
+  }
+}
+
+static already_AddRefed<LookupCacheV4> SetupLookupCache(
+    const nsACString& aName) {
+  nsCOMPtr<nsIFile> rootDir;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(rootDir));
+
+  nsAutoCString provider("test");
+  RefPtr<LookupCacheV4> lookup = new LookupCacheV4(aName, provider, rootDir);
+  lookup->Init();
+
+  return lookup.forget();
+}
+
+class UrlClassifierPrefixSetTest : public ::testing::TestWithParam<uint32_t> {
+ protected:
+  void SetUp() override {
+    // max_array_size to 0 means we are testing delta algorithm here.
+    static const char prefKey[] =
+        "browser.safebrowsing.prefixset.max_array_size";
+    mozilla::Preferences::SetUint(prefKey, GetParam());
+
+    mCache = SetupLookupCache("test"_ns);
+  }
+
+  void TearDown() override {
+    mCache = nullptr;
+    mArray.Clear();
+    mMap.Clear();
+  }
+
+  nsresult SetupPrefixes(_PrefixArray&& aArray) {
+    mArray = std::move(aArray);
+    PrefixArrayToPrefixStringMap(mArray, mMap);
+    return mCache->Build(mMap);
+  }
+
+  void SetupPrefixesAndVerify(_PrefixArray& aArray) {
+    mArray = aArray.Clone();
+    PrefixArrayToPrefixStringMap(mArray, mMap);
+
+    ASSERT_NS_SUCCEEDED(mCache->Build(mMap));
+    Verify();
+  }
+
+  void SetupPrefixesAndVerify(_PrefixArray&& aArray) {
+    nsresult rv = SetupPrefixes(std::move(aArray));
+    ASSERT_NS_SUCCEEDED(rv);
+    Verify();
+  }
+
+  void SetupRandomPrefixesAndVerify(uint32_t N, uint32_t MIN, uint32_t MAX) {
+    srand(time(nullptr));
+    RandomPrefixes(N, MIN, MAX, mArray);
+    PrefixArrayToPrefixStringMap(mArray, mMap);
+
+    ASSERT_NS_SUCCEEDED(mCache->Build(mMap));
+    Verify();
+  }
+
+  void Verify() {
+    DoExpectedLookup(mCache, mArray);
+    DoRandomLookup(mCache, 1000, mArray);
+    CheckContent(mCache, mArray);
+  }
+
+  RefPtr<LookupCacheV4> mCache;
+  _PrefixArray mArray;
+  PrefixStringMap mMap;
+};
+
+// Test setting prefix set with only 4-bytes prefixes
+TEST_P(UrlClassifierPrefixSetTest, FixedLengthSet) {
+  SetupPrefixesAndVerify({
+      _Prefix("alph"),
+      _Prefix("brav"),
+      _Prefix("char"),
+      _Prefix("delt"),
+      _Prefix("echo"),
+      _Prefix("foxt"),
+  });
+}
+
+TEST_P(UrlClassifierPrefixSetTest, FixedLengthRandomSet) {
+  SetupRandomPrefixesAndVerify(1500, 4, 4);
+}
+
+TEST_P(UrlClassifierPrefixSetTest, FixedLengthRandomLargeSet) {
+  SetupRandomPrefixesAndVerify(15000, 4, 4);
+}
+
+TEST_P(UrlClassifierPrefixSetTest, FixedLengthTinySet) {
+  SetupPrefixesAndVerify({
+      _Prefix("tiny"),
+  });
+}
+
+// Test setting prefix set with only 5~32 bytes prefixes
+TEST_P(UrlClassifierPrefixSetTest, VariableLengthSet) {
+  SetupPrefixesAndVerify(
+      {_Prefix("bravo"), _Prefix("charlie"), _Prefix("delta"),
+       _Prefix("EchoEchoEchoEchoEcho"), _Prefix("foxtrot"),
+       _Prefix("GolfGolfGolfGolfGolfGolfGolfGolf"), _Prefix("hotel"),
+       _Prefix("november"), _Prefix("oscar"), _Prefix("quebec"),
+       _Prefix("romeo"), _Prefix("sierrasierrasierrasierrasierra"),
+       _Prefix("Tango"), _Prefix("whiskey"), _Prefix("yankee"),
+       _Prefix("ZuluZuluZuluZulu")});
+}
+
+TEST_P(UrlClassifierPrefixSetTest, VariableLengthRandomSet) {
+  SetupRandomPrefixesAndVerify(1500, 5, 32);
+}
+
+// Test setting prefix set with both 4-bytes prefixes and 5~32 bytes prefixes
+TEST_P(UrlClassifierPrefixSetTest, MixedPrefixSet) {
+  SetupPrefixesAndVerify(
+      {_Prefix("enus"), _Prefix("apollo"), _Prefix("mars"),
+       _Prefix("Hecatonchires cyclopes"), _Prefix("vesta"), _Prefix("neptunus"),
+       _Prefix("jupiter"), _Prefix("diana"), _Prefix("minerva"),
+       _Prefix("ceres"), _Prefix("Aidos,Adephagia,Adikia,Aletheia"),
+       _Prefix("hecatonchires"), _Prefix("alcyoneus"), _Prefix("hades"),
+       _Prefix("vulcanus"), _Prefix("juno"), _Prefix("mercury"),
+       _Prefix("Stheno, Euryale and Medusa")});
+}
+
+TEST_P(UrlClassifierPrefixSetTest, MixedRandomPrefixSet) {
+  SetupRandomPrefixesAndVerify(1500, 4, 32);
+}
+
+// Test resetting prefix set
+TEST_P(UrlClassifierPrefixSetTest, ResetPrefix) {
+  // Base prefix set
+  _PrefixArray oldArray = {
+      _Prefix("Iceland"),   _Prefix("Peru"),    _Prefix("Mexico"),
+      _Prefix("Australia"), _Prefix("Japan"),   _Prefix("Egypt"),
+      _Prefix("America"),   _Prefix("Finland"), _Prefix("Germany"),
+      _Prefix("Italy"),     _Prefix("France"),  _Prefix("Taiwan"),
+  };
+  SetupPrefixesAndVerify(oldArray);
+
+  // New prefix set
+  _PrefixArray newArray = {
+      _Prefix("Pikachu"),    _Prefix("Bulbasaur"), _Prefix("Charmander"),
+      _Prefix("Blastoise"),  _Prefix("Pidgey"),    _Prefix("Mewtwo"),
+      _Prefix("Jigglypuff"), _Prefix("Persian"),   _Prefix("Tentacool"),
+      _Prefix("Onix"),       _Prefix("Eevee"),     _Prefix("Jynx"),
+  };
+  SetupPrefixesAndVerify(newArray);
+
+  // Should not match any of the first prefix set
+  uint32_t matchLength = 0;
+  for (uint32_t i = 0; i < oldArray.Length(); i++) {
+    Completion complete;
+    complete.Assign(CreateFullHash(oldArray[i]));
+
+    // Find match for prefix-generated full hash
+    bool has, confirmed;
+    mCache->Has(complete, &has, &matchLength, &confirmed);
+
+    ASSERT_TRUE(matchLength == 0);
+  }
+}
+
+// Test only set one 4-bytes prefix and one full-length prefix
+TEST_P(UrlClassifierPrefixSetTest, TinyPrefixSet) {
+  SetupPrefixesAndVerify({
+      _Prefix("AAAA"),
+      _Prefix("11112222333344445555666677778888"),
+  });
+}
+
+// Test empty prefix set and IsEmpty function
+TEST_P(UrlClassifierPrefixSetTest, EmptyFixedPrefixSet) {
+  ASSERT_TRUE(mCache->IsEmpty());
+
+  SetupPrefixesAndVerify({});
+
+  // Insert an 4-bytes prefix, then IsEmpty should return false
+  SetupPrefixesAndVerify({_Prefix("test")});
+
+  ASSERT_TRUE(!mCache->IsEmpty());
+}
+
+TEST_P(UrlClassifierPrefixSetTest, EmptyVariableLengthPrefixSet) {
+  ASSERT_TRUE(mCache->IsEmpty());
+
+  SetupPrefixesAndVerify({});
+
+  // Insert an 5~32 bytes prefix, then IsEmpty should return false
+  SetupPrefixesAndVerify({_Prefix("test variable length")});
+
+  ASSERT_TRUE(!mCache->IsEmpty());
+}
+
+// Test prefix size should only between 4~32 bytes
+TEST_P(UrlClassifierPrefixSetTest, MinMaxPrefixSet) {
+  // Test prefix set between 4-32 bytes, should success
+  SetupPrefixesAndVerify({_Prefix("1234"), _Prefix("ABCDEFGHIJKKMNOP"),
+                          _Prefix("1aaa2bbb3ccc4ddd5eee6fff7ggg8hhh")});
+
+  // Prefix size less than 4-bytes should fail
+  nsresult rv = SetupPrefixes({_Prefix("123")});
+  ASSERT_NS_FAILED(rv);
+
+  // Prefix size greater than 32-bytes should fail
+  rv = SetupPrefixes({_Prefix("1aaa2bbb3ccc4ddd5eee6fff7ggg8hhh9")});
+  ASSERT_NS_FAILED(rv);
+}
+
+// Test save then load prefix set with only 4-bytes prefixes
+TEST_P(UrlClassifierPrefixSetTest, LoadSaveFixedLengthPrefixSet) {
+  nsCOMPtr<nsIFile> file;
+  _PrefixArray array;
+  PrefixStringMap map;
+
+  // Save
+  {
+    RefPtr<LookupCacheV4> save = SetupLookupCache("test-save"_ns);
+
+    RandomPrefixes(10000, 4, 4, array);
+
+    PrefixArrayToPrefixStringMap(array, map);
+    save->Build(map);
+
+    DoExpectedLookup(save, array);
+    DoRandomLookup(save, 1000, array);
+    CheckContent(save, array);
+
+    NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
+    file->Append(u"test.vlpset"_ns);
+    save->StoreToFile(file);
+  }
+
+  // Load
+  {
+    RefPtr<LookupCacheV4> load = SetupLookupCache("test-load"_ns);
+    load->LoadFromFile(file);
+
+    DoExpectedLookup(load, array);
+    DoRandomLookup(load, 1000, array);
+    CheckContent(load, array);
+  }
+
+  file->Remove(false);
+}
+
+// Test save then load prefix set with only 5~32 bytes prefixes
+TEST_P(UrlClassifierPrefixSetTest, LoadSaveVariableLengthPrefixSet) {
+  nsCOMPtr<nsIFile> file;
+  _PrefixArray array;
+  PrefixStringMap map;
+
+  // Save
+  {
+    RefPtr<LookupCacheV4> save = SetupLookupCache("test-save"_ns);
+
+    RandomPrefixes(10000, 5, 32, array);
+
+    PrefixArrayToPrefixStringMap(array, map);
+    save->Build(map);
+
+    DoExpectedLookup(save, array);
+    DoRandomLookup(save, 1000, array);
+    CheckContent(save, array);
+
+    NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
+    file->Append(u"test.vlpset"_ns);
+    save->StoreToFile(file);
+  }
+
+  // Load
+  {
+    RefPtr<LookupCacheV4> load = SetupLookupCache("test-load"_ns);
+    load->LoadFromFile(file);
+
+    DoExpectedLookup(load, array);
+    DoRandomLookup(load, 1000, array);
+    CheckContent(load, array);
+  }
+
+  file->Remove(false);
+}
+
+// Test save then load prefix with both 4 bytes prefixes and 5~32 bytes prefixes
+TEST_P(UrlClassifierPrefixSetTest, LoadSavePrefixSet) {
+  nsCOMPtr<nsIFile> file;
+  _PrefixArray array;
+  PrefixStringMap map;
+
+  // Save
+  {
+    RefPtr<LookupCacheV4> save = SetupLookupCache("test-save"_ns);
+
+    // Try to simulate the real case that most prefixes are 4bytes
+    RandomPrefixes(20000, 4, 4, array);
+    RandomPrefixes(1000, 5, 32, array);
+
+    PrefixArrayToPrefixStringMap(array, map);
+    save->Build(map);
+
+    DoExpectedLookup(save, array);
+    DoRandomLookup(save, 1000, array);
+    CheckContent(save, array);
+
+    NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
+    file->Append(u"test.vlpset"_ns);
+    save->StoreToFile(file);
+  }
+
+  // Load
+  {
+    RefPtr<LookupCacheV4> load = SetupLookupCache("test-load"_ns);
+    load->LoadFromFile(file);
+
+    DoExpectedLookup(load, array);
+    DoRandomLookup(load, 1000, array);
+    CheckContent(load, array);
+  }
+
+  file->Remove(false);
+}
+
+// This is for fixed-length prefixset
+TEST_P(UrlClassifierPrefixSetTest, LoadSaveNoDelta) {
+  nsCOMPtr<nsIFile> file;
+  _PrefixArray array;
+  PrefixStringMap map;
+
+  for (uint32_t i = 0; i < 100; i++) {
+    // construct a tree without deltas by making the distance
+    // between entries larger than 16 bits
+    uint32_t v = ((1 << 16) + 1) * i;
+    nsCString* ele = array.AppendElement();
+    ele->AppendASCII(reinterpret_cast<const char*>(&v), 4);
+  }
+
+  // Save
+  {
+    RefPtr<LookupCacheV4> save = SetupLookupCache("test-save"_ns);
+
+    PrefixArrayToPrefixStringMap(array, map);
+    save->Build(map);
+
+    DoExpectedLookup(save, array);
+    DoRandomLookup(save, 1000, array);
+    CheckContent(save, array);
+
+    NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
+    file->Append(u"test.vlpset"_ns);
+    save->StoreToFile(file);
+  }
+
+  // Load
+  {
+    RefPtr<LookupCacheV4> load = SetupLookupCache("test-load"_ns);
+    load->LoadFromFile(file);
+
+    DoExpectedLookup(load, array);
+    DoRandomLookup(load, 1000, array);
+    CheckContent(load, array);
+  }
+
+  file->Remove(false);
+}
+
+// To run the same test for different configurations of
+// "browser_safebrowsing_prefixset_max_array_size"
+INSTANTIATE_TEST_SUITE_P(UrlClassifierPrefixSetTest, UrlClassifierPrefixSetTest,
+                         ::testing::Values(0, UINT32_MAX));