/* 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 "Common.h"

#include "Classifier.h"

#include "LookupCacheV4.h"

#define GTEST_TABLE_V4 NS_LITERAL_CSTRING("gtest-malware-proto")

#define GTEST_TABLE_V2 NS_LITERAL_CSTRING("gtest-malware-simple")

typedef nsCString _Fragment;

typedef nsTArray<nsCString> _PrefixArray;

static RefPtr<Classifier>

GetClassifier()

{

  nsCOMPtr<nsIFile> file;

  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));

  RefPtr<Classifier> classifier = new Classifier();

  nsresult rv = classifier->Open(*file);

  EXPECT_TRUE(rv == NS_OK);

  return classifier;

}

static nsresult

SetupLookupCacheV4(RefPtr<Classifier> classifier,

                   const _PrefixArray& aPrefixArray,

                   const nsACString& aTable)

{

  RefPtr<LookupCache> lookupCache = classifier->GetLookupCache(aTable, false);

  if (!lookupCache) {

    return NS_ERROR_FAILURE;

  }

  RefPtr<LookupCacheV4> lookupCacheV4 = LookupCache::Cast<LookupCacheV4>(lookupCache);

  if (!lookupCacheV4) {

    return NS_ERROR_FAILURE;

  }

  PrefixStringMap map;

  PrefixArrayToPrefixStringMap(aPrefixArray, map);

  return lookupCacheV4->Build(map);

}

static nsresult

SetupLookupCacheV2(RefPtr<Classifier> classifier,

                   const _PrefixArray& aPrefixArray,

                   const nsACString& aTable)

{

  RefPtr<LookupCache> lookupCache = classifier->GetLookupCache(aTable, false);

  if (!lookupCache) {

    return NS_ERROR_FAILURE;

  }

  RefPtr<LookupCacheV2> lookupCacheV2 =

    LookupCache::Cast<LookupCacheV2>(lookupCache);

  if (!lookupCacheV2) {

    return NS_ERROR_FAILURE;

  }

  AddPrefixArray prefixes;

  AddCompleteArray completions;

  nsresult rv = PrefixArrayToAddPrefixArrayV2(aPrefixArray, prefixes);

  if (NS_FAILED(rv)) {

    return rv;

  }

  EntrySort(prefixes);

  return lookupCacheV2->Build(prefixes, completions);

}

static void

TestReadNoiseEntries(RefPtr<Classifier> classifier,

                     const _PrefixArray& aPrefixArray,

                     const nsCString& aTable,

                     const nsCString& aFragment)

{

  Completion lookupHash;

  lookupHash.FromPlaintext(aFragment);

  RefPtr<LookupResult> result = new LookupResult;

  result->hash.complete = lookupHash;

  PrefixArray noiseEntries;

  uint32_t noiseCount = 3;

  nsresult rv;

  rv = classifier->ReadNoiseEntries(result->hash.fixedLengthPrefix,

                                    aTable, noiseCount,

                                    noiseEntries);

  ASSERT_TRUE(rv == NS_OK);

  EXPECT_TRUE(noiseEntries.Length() > 0);

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

    // Test the noise entry should not equal the "real" hash request

    EXPECT_NE(noiseEntries[i], result->hash.fixedLengthPrefix);

    // Test the noise entry should exist in the cached prefix array

    nsAutoCString partialHash;

    partialHash.Assign(reinterpret_cast<char*>(&noiseEntries[i]), PREFIX_SIZE);

    EXPECT_TRUE(aPrefixArray.Contains(partialHash));

  }

}

TEST(UrlClassifier, ReadNoiseEntriesV4)

{

  RefPtr<Classifier> classifier = GetClassifier();

  _PrefixArray array = { GeneratePrefix(_Fragment("bravo.com/"), 5),

                         GeneratePrefix(_Fragment("browsing.com/"), 9),

                         GeneratePrefix(_Fragment("gound.com/"), 4),

                         GeneratePrefix(_Fragment("small.com/"), 4),

                         GeneratePrefix(_Fragment("gdfad.com/"), 4),

                         GeneratePrefix(_Fragment("afdfound.com/"), 4),

                         GeneratePrefix(_Fragment("dffa.com/"), 4),

                       };

  array.Sort();

  nsresult rv;

  rv = SetupLookupCacheV4(classifier, array, GTEST_TABLE_V4);

  ASSERT_TRUE(rv == NS_OK);

  TestReadNoiseEntries(classifier, array, GTEST_TABLE_V4, _Fragment("gound.com/"));

}

TEST(UrlClassifier, ReadNoiseEntriesV2)

{

  RefPtr<Classifier> classifier = GetClassifier();

  _PrefixArray array = { GeneratePrefix(_Fragment("helloworld.com/"), 4),

                         GeneratePrefix(_Fragment("firefox.com/"), 4),

                         GeneratePrefix(_Fragment("chrome.com/"), 4),

                         GeneratePrefix(_Fragment("safebrowsing.com/"), 4),

                         GeneratePrefix(_Fragment("opera.com/"), 4),

                         GeneratePrefix(_Fragment("torbrowser.com/"), 4),

                         GeneratePrefix(_Fragment("gfaads.com/"), 4),

                         GeneratePrefix(_Fragment("qggdsas.com/"), 4),

                         GeneratePrefix(_Fragment("nqtewq.com/"), 4),

                       };

  nsresult rv;

  rv = SetupLookupCacheV2(classifier, array, GTEST_TABLE_V2);

  ASSERT_TRUE(rv == NS_OK);

  TestReadNoiseEntries(classifier, array, GTEST_TABLE_V2, _Fragment("helloworld.com/"));

}