/* 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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(VariableLengthPrefixSet, 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);
}