/* 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 <mozilla/RefPtr.h>

#include "nsString.h"

#include "nsTArray.h"

#include "nsClassHashtable.h"

#include "VariableLengthPrefixSet.h"

#include "nsAppDirectoryServiceDefs.h"

#include "nsIFile.h"

#include "gtest/gtest.h"

using namespace mozilla::safebrowsing;

typedef nsCString _Prefix;

typedef nsTArray<_Prefix> _PrefixArray;

// Create fullhash by appending random characters.

static nsCString* CreateFullHash(const nsACString& in)

{

  nsCString* out = new nsCString(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;

      }

    }

  }

}

static void CheckContent(VariableLengthPrefixSet* pset,

                         PrefixStringMap& expected)

{

  PrefixStringMap vlPSetMap;

  pset->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.

static void DoExpectedLookup(VariableLengthPrefixSet* pset,

                             _PrefixArray& array)

{

  uint32_t matchLength = 0;

  for (uint32_t i = 0; i < array.Length(); i++) {

    const nsCString& prefix = array[i];

    UniquePtr<nsCString> fullhash(CreateFullHash(prefix));

    // Find match for prefix-generated full hash

    pset->Matches(*fullhash, &matchLength);

    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(fullhash->BeginReading(),

                        array[j].BeginReading(),

                        matchLength)) {

          found = true;

          break;

        }

      }

      ASSERT_TRUE(found);

    }

  }

}

static void DoRandomLookup(VariableLengthPrefixSet* pset,

                           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());

      }

    }

    uint32_t matchLength = 0;

    pset->Matches(nsDependentCSubstring(buf, 32), &matchLength);

    ASSERT_TRUE(expected.IsEmpty() ? !matchLength : expected.Contains(matchLength));

  }

}

static void SetupPrefixMap(const _PrefixArray& array,

                             PrefixStringMap& map)

{

  map.Clear();

  // Buckets are keyed by prefix length and contain an array of

  // all prefixes of that length.

  nsClassHashtable<nsUint32HashKey, _PrefixArray> table;

  for (uint32_t i = 0; i < array.Length(); i++) {

    _PrefixArray* prefixes = table.Get(array[i].Length());

    if (!prefixes) {

      prefixes = new _PrefixArray();

      table.Put(array[i].Length(), prefixes);

    }

    prefixes->AppendElement(array[i]);

  }

  // The resulting map entries will be a concatenation of all

  // prefix data for the prefixes of a given size.

  for (auto iter = table.Iter(); !iter.Done(); iter.Next()) {

    uint32_t size = iter.Key();

    uint32_t count = iter.Data()->Length();

    _Prefix* str = new _Prefix();

    str->SetLength(size * count);

    char* dst = str->BeginWriting();

    iter.Data()->Sort();

    for (uint32_t i = 0; i < count; i++) {

      memcpy(dst, iter.Data()->ElementAt(i).get(), size);

      dst += size;

    }

    map.Put(size, str);

  }

}

// Test setting prefix set with only 4-bytes prefixes

TEST(UrlClassifierVLPrefixSet, FixedLengthSet)

{

  srand(time(nullptr));

  RefPtr<VariableLengthPrefixSet> pset = new VariableLengthPrefixSet;

  pset->Init(NS_LITERAL_CSTRING("test"));

  PrefixStringMap map;

  _PrefixArray array = { _Prefix("alph"), _Prefix("brav"), _Prefix("char"),

                         _Prefix("delt"), _Prefix("echo"), _Prefix("foxt"),

                       };

  SetupPrefixMap(array, map);

  pset->SetPrefixes(map);

  DoExpectedLookup(pset, array);

  DoRandomLookup(pset, 1000, array);

  CheckContent(pset, map);

  // Run random test

  array.Clear();

  map.Clear();

  RandomPrefixes(1500, 4, 4, array);

  SetupPrefixMap(array, map);

  pset->SetPrefixes(map);

  DoExpectedLookup(pset, array);

  DoRandomLookup(pset, 1000, array);

  CheckContent(pset, map);

}

// Test setting prefix set with only 5~32 bytes prefixes

TEST(UrlClassifierVLPrefixSet, VariableLengthSet)

{

  RefPtr<VariableLengthPrefixSet> pset = new VariableLengthPrefixSet;

  pset->Init(NS_LITERAL_CSTRING("test"));

  PrefixStringMap map;

  _PrefixArray array = { _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")

                       };

  SetupPrefixMap(array, map);

  pset->SetPrefixes(map);

  DoExpectedLookup(pset, array);

  DoRandomLookup(pset, 1000, array);

  CheckContent(pset, map);

  // Run random test

  array.Clear();

  map.Clear();

  RandomPrefixes(1500, 5, 32, array);

  SetupPrefixMap(array, map);

  pset->SetPrefixes(map);

  DoExpectedLookup(pset, array);

  DoRandomLookup(pset, 1000, array);

  CheckContent(pset, map);

}

// Test setting prefix set with both 4-bytes prefixes and 5~32 bytes prefixes

TEST(UrlClassifierVLPrefixSet, MixedPrefixSet)

{

  RefPtr<VariableLengthPrefixSet> pset = new VariableLengthPrefixSet;

  pset->Init(NS_LITERAL_CSTRING("test"));

  PrefixStringMap map;

  _PrefixArray array = { _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")

                       };

  SetupPrefixMap(array, map);

  pset->SetPrefixes(map);

  DoExpectedLookup(pset, array);

  DoRandomLookup(pset, 1000, array);

  CheckContent(pset, map);

  // Run random test

  array.Clear();

  map.Clear();

  RandomPrefixes(1500, 4, 32, array);

  SetupPrefixMap(array, map);

  pset->SetPrefixes(map);

  DoExpectedLookup(pset, array);

  DoRandomLookup(pset, 1000, array);

  CheckContent(pset, map);

}

// Test resetting prefix set

TEST(UrlClassifierVLPrefixSet, ResetPrefix)

{

  RefPtr<VariableLengthPrefixSet> pset = new VariableLengthPrefixSet;

  pset->Init(NS_LITERAL_CSTRING("test"));

  // First prefix set

  _PrefixArray array1 = { _Prefix("Iceland"), _Prefix("Peru"), _Prefix("Mexico"),

                          _Prefix("Australia"), _Prefix("Japan"), _Prefix("Egypt"),

                          _Prefix("America"), _Prefix("Finland"), _Prefix("Germany"),

                          _Prefix("Italy"), _Prefix("France"), _Prefix("Taiwan"),

                        };

  {

    PrefixStringMap map;

    SetupPrefixMap(array1, map);

    pset->SetPrefixes(map);

    DoExpectedLookup(pset, array1);

  }

  // Second

  _PrefixArray array2 = { _Prefix("Pikachu"), _Prefix("Bulbasaur"), _Prefix("Charmander"),

                          _Prefix("Blastoise"), _Prefix("Pidgey"), _Prefix("Mewtwo"),

                          _Prefix("Jigglypuff"), _Prefix("Persian"), _Prefix("Tentacool"),

                          _Prefix("Onix"), _Prefix("Eevee"), _Prefix("Jynx"),

                        };

  {

    PrefixStringMap map;

    SetupPrefixMap(array2, map);

    pset->SetPrefixes(map);

    DoExpectedLookup(pset, array2);

  }

  // Should not match any of the first prefix set

  uint32_t matchLength = 0;

  for (uint32_t i = 0; i < array1.Length(); i++) {

    UniquePtr<nsACString> fullhash(CreateFullHash(array1[i]));

    pset->Matches(*fullhash, &matchLength);

    ASSERT_TRUE(matchLength == 0);

  }

}

// Test only set one 4-bytes prefix and one full-length prefix

TEST(UrlClassifierVLPrefixSet, TinyPrefixSet)

{

  RefPtr<VariableLengthPrefixSet> pset = new VariableLengthPrefixSet;

  pset->Init(NS_LITERAL_CSTRING("test"));

  PrefixStringMap map;

  _PrefixArray array = { _Prefix("AAAA"),

                         _Prefix("11112222333344445555666677778888"),

                       };

  SetupPrefixMap(array, map);

  pset->SetPrefixes(map);

  DoExpectedLookup(pset, array);

  DoRandomLookup(pset, 1000, array);

  CheckContent(pset, map);

}

// Test empty prefix set and IsEmpty function

TEST(UrlClassifierVLPrefixSet, EmptyPrefixSet)

{

  RefPtr<VariableLengthPrefixSet> pset = new VariableLengthPrefixSet;

  pset->Init(NS_LITERAL_CSTRING("test"));

  bool empty;

  pset->IsEmpty(&empty);

  ASSERT_TRUE(empty);

  PrefixStringMap map;

  _PrefixArray array1;

  // Lookup an empty array should never match

  DoRandomLookup(pset, 100, array1);

  // Insert an 4-bytes prefix, then IsEmpty should return false

  _PrefixArray array2 = { _Prefix("test") };

  SetupPrefixMap(array2, map);

  pset->SetPrefixes(map);

  pset->IsEmpty(&empty);

  ASSERT_TRUE(!empty);

  _PrefixArray array3 = { _Prefix("test variable length") };

  // Insert an 5~32 bytes prefix, then IsEmpty should return false

  SetupPrefixMap(array3, map);

  pset->SetPrefixes(map);

  pset->IsEmpty(&empty);

  ASSERT_TRUE(!empty);

}

// Test prefix size should only between 4~32 bytes

TEST(UrlClassifierVLPrefixSet, MinMaxPrefixSet)

{

  RefPtr<VariableLengthPrefixSet> pset = new VariableLengthPrefixSet;

  pset->Init(NS_LITERAL_CSTRING("test"));

  PrefixStringMap map;

  {

    _PrefixArray array = { _Prefix("1234"),

                           _Prefix("ABCDEFGHIJKKMNOP"),

                           _Prefix("1aaa2bbb3ccc4ddd5eee6fff7ggg8hhh") };

    SetupPrefixMap(array, map);

    nsresult rv = pset->SetPrefixes(map);

    ASSERT_TRUE(rv == NS_OK);

  }

  // Prefix size less than 4-bytes should fail

  {

    _PrefixArray array = { _Prefix("123") };

    SetupPrefixMap(array, map);

    nsresult rv = pset->SetPrefixes(map);

    ASSERT_TRUE(NS_FAILED(rv));

  }

  // Prefix size greater than 32-bytes should fail

  {

    _PrefixArray array = { _Prefix("1aaa2bbb3ccc4ddd5eee6fff7ggg8hhh9") };

    SetupPrefixMap(array, map);

    nsresult rv = pset->SetPrefixes(map);

    ASSERT_TRUE(NS_FAILED(rv));

  }

}

// Test save then load prefix set with only 4-bytes prefixes

TEST(UrlClassifierVLPrefixSet, LoadSaveFixedLengthPrefixSet)

{

  RefPtr<VariableLengthPrefixSet> save = new VariableLengthPrefixSet;

  save->Init(NS_LITERAL_CSTRING("test-save"));

  _PrefixArray array;

  RandomPrefixes(10000, 4, 4, array);

  PrefixStringMap map;

  SetupPrefixMap(array, map);

  save->SetPrefixes(map);

  DoExpectedLookup(save, array);

  DoRandomLookup(save, 1000, array);

  CheckContent(save, map);

  nsCOMPtr<nsIFile> file;

  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,

                         getter_AddRefs(file));

  file->Append(NS_LITERAL_STRING("test.vlpset"));

  save->StoreToFile(file);

  RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;

  load->Init(NS_LITERAL_CSTRING("test-load"));

  load->LoadFromFile(file);

  DoExpectedLookup(load, array);

  DoRandomLookup(load, 1000, array);

  CheckContent(load, map);

  file->Remove(false);

}

// Test save then load prefix set with only 5~32 bytes prefixes

TEST(UrlClassifierVLPrefixSet, LoadSaveVariableLengthPrefixSet)

{

  RefPtr<VariableLengthPrefixSet> save = new VariableLengthPrefixSet;

  save->Init(NS_LITERAL_CSTRING("test-save"));

  _PrefixArray array;

  RandomPrefixes(10000, 5, 32, array);

  PrefixStringMap map;

  SetupPrefixMap(array, map);

  save->SetPrefixes(map);

  DoExpectedLookup(save, array);

  DoRandomLookup(save, 1000, array);

  CheckContent(save, map);

  nsCOMPtr<nsIFile> file;

  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,

                         getter_AddRefs(file));

  file->Append(NS_LITERAL_STRING("test.vlpset"));

  save->StoreToFile(file);

  RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;

  load->Init(NS_LITERAL_CSTRING("test-load"));

  load->LoadFromFile(file);

  DoExpectedLookup(load, array);

  DoRandomLookup(load, 1000, array);

  CheckContent(load, map);

  file->Remove(false);

}

// Test save then load prefix with both 4 bytes prefixes and 5~32 bytes prefixes

TEST(UrlClassifierVLPrefixSet, LoadSavePrefixSet)

{

  RefPtr<VariableLengthPrefixSet> save = new VariableLengthPrefixSet;

  save->Init(NS_LITERAL_CSTRING("test-save"));

  // Try to simulate the real case that most prefixes are 4bytes

  _PrefixArray array;

  RandomPrefixes(20000, 4, 4, array);

  RandomPrefixes(1000, 5, 32, array);

  PrefixStringMap map;

  SetupPrefixMap(array, map);

  save->SetPrefixes(map);

  DoExpectedLookup(save, array);

  DoRandomLookup(save, 1000, array);

  CheckContent(save, map);

  nsCOMPtr<nsIFile> file;

  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,

                         getter_AddRefs(file));

  file->Append(NS_LITERAL_STRING("test.vlpset"));

  save->StoreToFile(file);

  RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;

  load->Init(NS_LITERAL_CSTRING("test-load"));

  load->LoadFromFile(file);

  DoExpectedLookup(load, array);

  DoRandomLookup(load, 1000, array);

  CheckContent(load, map);

  file->Remove(false);

}