Bug 1305801 - Part 5: Support SafeBrowsing v4 partial update. r=gcp

MozReview-Commit-ID: 7OEWLaZbotS

--HG--
extra : rebase_source : dea720f726306238f7fe3687d7c2579410923752
This commit is contained in:
Dimi Lee 2016-10-04 09:14:39 +08:00
Родитель f0aab4a2ce
Коммит b4c5400f77
11 изменённых файлов: 780 добавлений и 100 удалений

Просмотреть файл

@ -3821,6 +3821,14 @@
"bug_numbers": [1172688],
"description": "This metric is recorded every time a gethash lookup is performed, `true` is recorded if the lookup times out."
},
"URLCLASSIFIER_UPDATE_ERROR_TYPE": {
"alert_emails": ["safebrowsing-telemetry@mozilla.org"],
"expires_in_version": "58",
"kind": "enumerated",
"n_values": 10,
"bug_numbers": [1305801],
"description": "An error was encountered while parsing a partial update returned by a Safe Browsing V4 server (0 = addition of an already existing prefix, 1 = parser got into an infinite loop, 2 = removal index out of bounds)"
},
"CSP_DOCUMENTS_COUNT": {
"alert_emails": ["seceng@mozilla.com"],
"bug_numbers": [1252829],

Просмотреть файл

@ -4,6 +4,7 @@
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
#include "Classifier.h"
#include "LookupCacheV4.h"
#include "nsIPrefBranch.h"
#include "nsIPrefService.h"
#include "nsISimpleEnumerator.h"
@ -946,7 +947,14 @@ Classifier::UpdateTableV4(nsTArray<TableUpdate*>* aUpdates,
return NS_ERROR_FAILURE;
}
PrefixStringMap prefixes;
nsresult rv = NS_OK;
// prefixes2 is only used in partial update. If there are multiple
// updates for the same table, prefixes1 & prefixes2 will act as
// input and output in turn to reduce memory copy overhead.
PrefixStringMap prefixes1, prefixes2;
PrefixStringMap* output = &prefixes1;
for (uint32_t i = 0; i < aUpdates->Length(); i++) {
TableUpdate *update = aUpdates->ElementAt(i);
if (!update || !update->TableName().Equals(aTable)) {
@ -957,23 +965,44 @@ Classifier::UpdateTableV4(nsTArray<TableUpdate*>* aUpdates,
NS_ENSURE_TRUE(updateV4, NS_ERROR_FAILURE);
if (updateV4->IsFullUpdate()) {
prefixes.Clear();
TableUpdateV4::PrefixesStringMap& map = updateV4->Prefixes();
TableUpdateV4::PrefixStdStringMap& map = updateV4->Prefixes();
output->Clear();
for (auto iter = map.Iter(); !iter.Done(); iter.Next()) {
// prefixes is an nsClassHashtable object stores prefix string.
// It will take the ownership of the put object.
nsCString* prefix = new nsCString(iter.Data()->GetPrefixString());
prefixes.Put(iter.Key(), prefix);
output->Put(iter.Key(), prefix);
}
} else {
// TODO: Bug 1287058, partial update
PrefixStringMap* input = nullptr;
// If both prefix sets are empty, this means we are doing a partial update
// without a prior full/partial update in the loop. In this case we should
// get prefixes from the lookup cache first.
if (prefixes1.IsEmpty() && prefixes2.IsEmpty()) {
lookupCache->GetPrefixes(prefixes1);
input = &prefixes1;
output = &prefixes2;
} else {
MOZ_ASSERT(prefixes1.IsEmpty() ^ prefixes2.IsEmpty());
// When there are multiple partial updates, input should always point
// to the non-empty prefix set(filled by previous full/partial update).
// output should always point to the empty prefix set.
input = prefixes1.IsEmpty() ? &prefixes2 : &prefixes1;
output = prefixes1.IsEmpty() ? &prefixes1 : &prefixes2;
}
rv = lookupCache->ApplyPartialUpdate(updateV4, *input, *output);
NS_ENSURE_SUCCESS(rv, rv);
input->Clear();
}
aUpdates->ElementAt(i) = nullptr;
}
nsresult rv = lookupCache->Build(prefixes);
rv = lookupCache->Build(*output);
NS_ENSURE_SUCCESS(rv, rv);
rv = lookupCache->WriteFile();

Просмотреть файл

@ -168,7 +168,7 @@ TableUpdateV4::NewPrefixes(int32_t aSize, std::string& aPrefixes)
NS_ENSURE_TRUE_VOID(aPrefixes.size() % aSize == 0);
NS_ENSURE_TRUE_VOID(!mPrefixesMap.Get(aSize));
PrefixString* prefix = new PrefixString(aPrefixes);
PrefixStdString* prefix = new PrefixStdString(aPrefixes);
mPrefixesMap.Put(aSize, prefix);
}

Просмотреть файл

@ -132,13 +132,13 @@ private:
// for addition and indices to removal. See Bug 1283009.
class TableUpdateV4 : public TableUpdate {
public:
struct PrefixString {
struct PrefixStdString {
private:
std::string mStorage;
nsDependentCSubstring mString;
public:
explicit PrefixString(std::string& aString)
explicit PrefixStdString(std::string& aString)
{
aString.swap(mStorage);
mString.Rebind(mStorage.data(), mStorage.size());
@ -147,7 +147,7 @@ public:
const nsACString& GetPrefixString() const { return mString; };
};
typedef nsClassHashtable<nsUint32HashKey, PrefixString> PrefixesStringMap;
typedef nsClassHashtable<nsUint32HashKey, PrefixStdString> PrefixStdStringMap;
typedef nsTArray<int32_t> RemovalIndiceArray;
public:
@ -163,7 +163,7 @@ public:
}
bool IsFullUpdate() const { return mFullUpdate; }
PrefixesStringMap& Prefixes() { return mPrefixesMap; }
PrefixStdStringMap& Prefixes() { return mPrefixesMap; }
RemovalIndiceArray& RemovalIndices() { return mRemovalIndiceArray; }
// For downcasting.
@ -177,7 +177,7 @@ private:
virtual int Tag() const override { return TAG; }
bool mFullUpdate;
PrefixesStringMap mPrefixesMap;
PrefixStdStringMap mPrefixesMap;
RemovalIndiceArray mRemovalIndiceArray;
};

Просмотреть файл

@ -41,7 +41,6 @@ namespace mozilla {
namespace safebrowsing {
const int LookupCacheV2::VER = 2;
const int LookupCacheV4::VER = 4;
LookupCache::LookupCache(const nsACString& aTableName, nsIFile* aRootStoreDir)
: mPrimed(false)
@ -588,56 +587,5 @@ LookupCacheV2::DumpCompletions()
}
#endif
nsresult
LookupCacheV4::Init()
{
mVLPrefixSet = new VariableLengthPrefixSet();
nsresult rv = mVLPrefixSet->Init(mTableName);
NS_ENSURE_SUCCESS(rv, rv);
return NS_OK;
}
// TODO : Bug 1298257, Implement url matching for variable-length prefix set
nsresult
LookupCacheV4::Has(const Completion& aCompletion,
bool* aHas, bool* aComplete)
{
*aHas = false;
return NS_OK;
}
nsresult
LookupCacheV4::Build(PrefixStringMap& aPrefixMap)
{
return mVLPrefixSet->SetPrefixes(aPrefixMap);
}
nsresult
LookupCacheV4::ClearPrefixes()
{
// Clear by seting a empty map
PrefixStringMap map;
return mVLPrefixSet->SetPrefixes(map);
}
nsresult
LookupCacheV4::StoreToFile(nsIFile* aFile)
{
return mVLPrefixSet->StoreToFile(aFile);
}
nsresult
LookupCacheV4::LoadFromFile(nsIFile* aFile)
{
return mVLPrefixSet->LoadFromFile(aFile);
}
size_t
LookupCacheV4::SizeOfPrefixSet()
{
return mVLPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
}
} // namespace safebrowsing
} // namespace mozilla

Просмотреть файл

@ -202,33 +202,6 @@ private:
RefPtr<nsUrlClassifierPrefixSet> mPrefixSet;
};
class LookupCacheV4 final : public LookupCache
{
public:
explicit LookupCacheV4(const nsACString& aTableName, nsIFile* aStoreFile)
: LookupCache(aTableName, aStoreFile) {}
~LookupCacheV4() {}
virtual nsresult Init() override;
virtual nsresult Has(const Completion& aCompletion,
bool* aHas, bool* aComplete) override;
nsresult Build(PrefixStringMap& aPrefixMap);
static const int VER;
protected:
virtual nsresult ClearPrefixes() override;
virtual nsresult StoreToFile(nsIFile* aFile) override;
virtual nsresult LoadFromFile(nsIFile* aFile) override;
virtual size_t SizeOfPrefixSet() override;
private:
virtual int Ver() const override { return VER; }
RefPtr<VariableLengthPrefixSet> mVLPrefixSet;
};
} // namespace safebrowsing
} // namespace mozilla

Просмотреть файл

@ -0,0 +1,299 @@
//* -*- 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 "HashStore.h"
// MOZ_LOG=UrlClassifierDbService:5
extern mozilla::LazyLogModule gUrlClassifierDbServiceLog;
#define LOG(args) MOZ_LOG(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug, args)
#define LOG_ENABLED() MOZ_LOG_TEST(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug)
namespace mozilla {
namespace safebrowsing {
const int LookupCacheV4::VER = 4;
// Prefixes coming from updates and VLPrefixSet are both stored in the HashTable
// where the (key, value) pair is a prefix size and a lexicographic-sorted string.
// The difference is prefixes from updates use std:string(to avoid additional copies)
// and prefixes from VLPrefixSet use nsCString.
// This class provides a common interface for the partial update algorithm to make it
// easier to operate on two different kind prefix string map..
class VLPrefixSet
{
public:
explicit VLPrefixSet(const PrefixStringMap& aMap);
explicit VLPrefixSet(const TableUpdateV4::PrefixStdStringMap& aMap);
// This function will merge the prefix map in VLPrefixSet to aPrefixMap.
void Merge(PrefixStringMap& aPrefixMap);
// Find the smallest string from the map in VLPrefixSet.
bool GetSmallestPrefix(nsDependentCSubstring& aOutString);
// Return the number of prefixes in the map
uint32_t Count() const { return mCount; }
private:
// PrefixString structure contains a lexicographic-sorted string with
// a |pos| variable to indicate which substring we are pointing to right now.
// |pos| increases each time GetSmallestPrefix finds the smallest string.
struct PrefixString {
PrefixString(const nsACString& aStr, uint32_t aSize)
: pos(0)
, size(aSize)
{
data.Rebind(aStr.BeginReading(), aStr.Length());
}
const char* get() {
return pos < data.Length() ? data.BeginReading() + pos : nullptr;
}
void next() { pos += size; }
uint32_t remaining() { return data.Length() - pos; }
nsDependentCSubstring data;
uint32_t pos;
uint32_t size;
};
nsClassHashtable<nsUint32HashKey, PrefixString> mMap;
uint32_t mCount;
};
nsresult
LookupCacheV4::Init()
{
mVLPrefixSet = new VariableLengthPrefixSet();
nsresult rv = mVLPrefixSet->Init(mTableName);
NS_ENSURE_SUCCESS(rv, rv);
return NS_OK;
}
// TODO : Bug 1298257, Implement url matching for variable-length prefix set
nsresult
LookupCacheV4::Has(const Completion& aCompletion,
bool* aHas, bool* aComplete)
{
*aHas = false;
return NS_OK;
}
nsresult
LookupCacheV4::Build(PrefixStringMap& aPrefixMap)
{
return mVLPrefixSet->SetPrefixes(aPrefixMap);
}
nsresult
LookupCacheV4::GetPrefixes(PrefixStringMap& aPrefixMap)
{
return mVLPrefixSet->GetPrefixes(aPrefixMap);
}
nsresult
LookupCacheV4::ClearPrefixes()
{
// Clear by seting a empty map
PrefixStringMap map;
return mVLPrefixSet->SetPrefixes(map);
}
nsresult
LookupCacheV4::StoreToFile(nsIFile* aFile)
{
return mVLPrefixSet->StoreToFile(aFile);
}
nsresult
LookupCacheV4::LoadFromFile(nsIFile* aFile)
{
return mVLPrefixSet->LoadFromFile(aFile);
}
size_t
LookupCacheV4::SizeOfPrefixSet()
{
return mVLPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
}
static void
AppendPrefixToMap(PrefixStringMap& prefixes, nsDependentCSubstring& prefix)
{
if (!prefix.Length()) {
return;
}
nsCString* prefixString = prefixes.LookupOrAdd(prefix.Length());
prefixString->Append(prefix.BeginReading(), prefix.Length());
}
// Please see https://bug1287058.bmoattachments.org/attachment.cgi?id=8795366
// for detail about partial update algorithm.
nsresult
LookupCacheV4::ApplyPartialUpdate(TableUpdateV4* aTableUpdate,
PrefixStringMap& aInputMap,
PrefixStringMap& aOutputMap)
{
MOZ_ASSERT(aOutputMap.IsEmpty());
// oldPSet contains prefixes we already have or we just merged last round.
// addPSet contains prefixes stored in tableUpdate which should be merged with oldPSet.
VLPrefixSet oldPSet(aInputMap);
VLPrefixSet addPSet(aTableUpdate->Prefixes());
// RemovalIndiceArray is a sorted integer array indicating the index of prefix we should
// remove from the old prefix set(according to lexigraphic order).
// |removalIndex| is the current index of RemovalIndiceArray.
// |numOldPrefixPicked| is used to record how many prefixes we picked from the old map.
TableUpdateV4::RemovalIndiceArray& removalArray = aTableUpdate->RemovalIndices();
uint32_t removalIndex = 0;
int32_t numOldPrefixPicked = -1;
nsDependentCSubstring smallestOldPrefix;
nsDependentCSubstring smallestAddPrefix;
// This is used to avoid infinite loop for partial update algorithm.
// The maximum loops will be the number of old prefixes plus the number of add prefixes.
uint32_t index = oldPSet.Count() + addPSet.Count() + 1;
for(;index > 0; index--) {
// Get smallest prefix from the old prefix set if we don't have one
if (smallestOldPrefix.IsEmpty()) {
// If prefixes from the old prefix set are all merged,
// then we can merge the entire add prefix set directly.
if (!oldPSet.GetSmallestPrefix(smallestOldPrefix)) {
AppendPrefixToMap(aOutputMap, smallestAddPrefix);
addPSet.Merge(aOutputMap);
break;
}
}
// Get smallest prefix from add prefix set if we don't have one
if (smallestAddPrefix.IsEmpty()) {
// If add prefixes are all merged and there is no removalIndices left,
// then merge the entire old prefix set directly. If there are still
// removalIndices left, we should still merge prefixes one by one
// to know which prefix from old prefix set should be removed.
if (!addPSet.GetSmallestPrefix(smallestAddPrefix) &&
removalIndex >= removalArray.Length()) {
AppendPrefixToMap(aOutputMap, smallestOldPrefix);
oldPSet.Merge(aOutputMap);
break;
}
}
// Compare the smallest string in old prefix set and add prefix set, merge the
// smaller one into new map to ensure merged string still follows
// lexigraphic order.
if (smallestOldPrefix < smallestAddPrefix ||
smallestAddPrefix.IsEmpty()) {
numOldPrefixPicked++;
// If the number of picks from old map matches the removalIndex, then this prefix
// will be removed by not merging it to new map.
if (removalIndex < removalArray.Length() &&
numOldPrefixPicked == removalArray[removalIndex]) {
removalIndex++;
} else {
AppendPrefixToMap(aOutputMap, smallestOldPrefix);
}
smallestOldPrefix.SetLength(0);
} else if (smallestOldPrefix > smallestAddPrefix ||
smallestOldPrefix.IsEmpty()){
AppendPrefixToMap(aOutputMap, smallestAddPrefix);
smallestAddPrefix.SetLength(0);
} else {
NS_WARNING("Add prefix should not exist in the original prefix set.");
Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
DUPLICATE_PREFIX);
return NS_ERROR_FAILURE;
}
}
// We expect index will be greater to 0 because max number of runs will be
// the number of original prefix plus add prefix.
if (index <= 0) {
NS_WARNING("There are still prefixes remaining after reaching maximum runs.");
Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
INFINITE_LOOP);
return NS_ERROR_FAILURE;
}
if (removalIndex < removalArray.Length()) {
NS_WARNING("There are still prefixes to remove after exhausting the old PrefixSet.");
Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
WRONG_REMOVAL_INDICES);
return NS_ERROR_FAILURE;
}
return NS_OK;
}
VLPrefixSet::VLPrefixSet(const PrefixStringMap& aMap)
: mCount(0)
{
for (auto iter = aMap.ConstIter(); !iter.Done(); iter.Next()) {
uint32_t size = iter.Key();
mMap.Put(size, new PrefixString(*iter.Data(), size));
mCount += iter.Data()->Length() / size;
}
}
VLPrefixSet::VLPrefixSet(const TableUpdateV4::PrefixStdStringMap& aMap)
: mCount(0)
{
for (auto iter = aMap.ConstIter(); !iter.Done(); iter.Next()) {
uint32_t size = iter.Key();
mMap.Put(size, new PrefixString(iter.Data()->GetPrefixString(), size));
mCount += iter.Data()->GetPrefixString().Length() / size;
}
}
void
VLPrefixSet::Merge(PrefixStringMap& aPrefixMap) {
for (auto iter = mMap.ConstIter(); !iter.Done(); iter.Next()) {
nsCString* prefixString = aPrefixMap.LookupOrAdd(iter.Key());
PrefixString* str = iter.Data();
if (str->get()) {
prefixString->Append(str->get(), str->remaining());
}
}
}
bool
VLPrefixSet::GetSmallestPrefix(nsDependentCSubstring& aOutString) {
PrefixString* pick = nullptr;
for (auto iter = mMap.ConstIter(); !iter.Done(); iter.Next()) {
PrefixString* str = iter.Data();
if (!str->get()) {
continue;
}
if (aOutString.IsEmpty()) {
aOutString.Rebind(str->get(), iter.Key());
pick = str;
continue;
}
nsDependentCSubstring cur(str->get(), iter.Key());
if (cur < aOutString) {
aOutString.Rebind(str->get(), iter.Key());
pick = str;
}
}
if (pick) {
pick->next();
}
return pick != nullptr;
}
} // namespace safebrowsing
} // namespace mozilla

Просмотреть файл

@ -0,0 +1,61 @@
//* -*- 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/. */
#ifndef LookupCacheV4_h__
#define LookupCacheV4_h__
#include "LookupCache.h"
namespace mozilla {
namespace safebrowsing {
// Forward declaration.
class TableUpdateV4;
class LookupCacheV4 final : public LookupCache
{
public:
explicit LookupCacheV4(const nsACString& aTableName, nsIFile* aStoreFile)
: LookupCache(aTableName, aStoreFile) {}
~LookupCacheV4() {}
virtual nsresult Init() override;
virtual nsresult Has(const Completion& aCompletion,
bool* aHas, bool* aComplete) override;
nsresult Build(PrefixStringMap& aPrefixMap);
nsresult GetPrefixes(PrefixStringMap& aPrefixMap);
// ApplyPartialUpdate will merge partial update data stored in aTableUpdate
// with prefixes in aInputMap.
nsresult ApplyPartialUpdate(TableUpdateV4* aTableUpdate,
PrefixStringMap& aInputMap,
PrefixStringMap& aOutputMap);
static const int VER;
protected:
virtual nsresult ClearPrefixes() override;
virtual nsresult StoreToFile(nsIFile* aFile) override;
virtual nsresult LoadFromFile(nsIFile* aFile) override;
virtual size_t SizeOfPrefixSet() override;
private:
virtual int Ver() const override { return VER; }
enum UPDATE_ERROR_TYPES {
DUPLICATE_PREFIX = 0,
INFINITE_LOOP = 1,
WRONG_REMOVAL_INDICES = 2,
};
RefPtr<VariableLengthPrefixSet> mVLPrefixSet;
};
} // namespace safebrowsing
} // namespace mozilla
#endif

Просмотреть файл

@ -24,6 +24,7 @@ UNIFIED_SOURCES += [
'ChunkSet.cpp',
'Classifier.cpp',
'LookupCache.cpp',
'LookupCacheV4.cpp',
'nsCheckSummedOutputStream.cpp',
'nsUrlClassifierDBService.cpp',
'nsUrlClassifierProxies.cpp',
@ -62,6 +63,7 @@ EXTRA_JS_MODULES += [
EXPORTS += [
'Entries.h',
'LookupCache.h',
'LookupCacheV4.h',
'nsUrlClassifierPrefixSet.h',
'protobuf/safebrowsing.pb.h',
'VariableLengthPrefixSet.h',

Просмотреть файл

@ -1,4 +1,5 @@
#include "LookupCache.h"
#include "LookupCacheV4.h"
#include "HashStore.h"
#include "gtest/gtest.h"
#include "nsIThread.h"

Просмотреть файл

@ -5,6 +5,7 @@
#include "nsIThread.h"
#include "string.h"
#include "gtest/gtest.h"
#include "nsThreadUtils.h"
using namespace mozilla;
using namespace mozilla::safebrowsing;
@ -12,6 +13,32 @@ using namespace mozilla::safebrowsing;
typedef nsCString _Prefix;
typedef nsTArray<_Prefix> _PrefixArray;
// This function removes common elements of inArray and outArray from
// outArray. This is used by partial update testcase to ensure partial update
// data won't contain prefixes we already have.
static void
RemoveIntersection(const _PrefixArray& inArray, _PrefixArray& outArray)
{
for (uint32_t i = 0; i < inArray.Length(); i++) {
int32_t idx = outArray.BinaryIndexOf(inArray[i]);
if (idx >= 0) {
outArray.RemoveElementAt(idx);
}
}
}
// This fucntion removes elements from outArray by index specified in
// removal array.
static void
RemoveElements(const nsTArray<uint32_t>& removal, _PrefixArray& outArray)
{
for (int32_t i = removal.Length() - 1; i >= 0; i--) {
outArray.RemoveElementAt(removal[i]);
}
}
// This function converts lexigraphic-sorted prefixes to a hashtable
// which key is prefix size and value is concatenated prefix string.
static void
PrefixArrayToPrefixStringMap(const _PrefixArray& prefixArray,
PrefixStringMap& outMap)
@ -27,7 +54,6 @@ PrefixArrayToPrefixStringMap(const _PrefixArray& prefixArray,
// N: Number of prefixes, MIN/MAX: minimum/maximum prefix size
// This function will append generated prefixes to outArray.
// All elements in the generated array will be different.
static void
CreateRandomSortedPrefixArray(uint32_t N,
uint32_t MIN,
@ -59,6 +85,20 @@ CreateRandomSortedPrefixArray(uint32_t N,
outArray.Sort();
}
// N: Number of removal indices, MAX: maximum index
static void
CreateRandomRemovalIndices(uint32_t N,
uint32_t MAX,
nsTArray<uint32_t>& outArray)
{
for (uint32_t i = 0; i < N; i++) {
uint32_t idx = rand() % MAX;
if (!outArray.Contains(idx)) {
outArray.InsertElementSorted(idx);
}
}
}
// Function to generate TableUpdateV4.
static void
GenerateUpdateData(bool fullUpdate,
@ -89,20 +129,21 @@ VerifyPrefixSet(PrefixStringMap& expected)
NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
getter_AddRefs(file));
file->AppendRelativeNativePath(NS_LITERAL_CSTRING("safebrowsing/gtest-malware-proto.pset"));
file->AppendNative(NS_LITERAL_CSTRING("safebrowsing"));
file->AppendNative(NS_LITERAL_CSTRING("gtest-malware-proto.pset"));
RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;
load->Init(NS_LITERAL_CSTRING("gtest-malware-proto"));
PrefixStringMap out;
PrefixStringMap prefixesInFile;
load->LoadFromFile(file);
load->GetPrefixes(out);
load->GetPrefixes(prefixesInFile);
for (auto iter = expected.ConstIter(); !iter.Done(); iter.Next()) {
nsCString* str1 = iter.Data();
nsCString* str2 = out.Get(iter.Key());
nsCString* expectedPrefix = iter.Data();
nsCString* resultPrefix = prefixesInFile.Get(iter.Key());
ASSERT_TRUE(*str1 == *str2);
ASSERT_TRUE(*resultPrefix == *expectedPrefix);
}
}
@ -112,8 +153,26 @@ Clear()
nsCOMPtr<nsIFile> file;
NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
getter_AddRefs(file));
file->AppendRelativeNativePath(NS_LITERAL_CSTRING("safebrowsing/gtest-malware-proto.pset"));
file->Remove(false);
UniquePtr<Classifier> classifier(new Classifier());
classifier->Open(*file);
classifier->Reset();
}
static void
testUpdateFail(nsTArray<TableUpdate*>& tableUpdates)
{
nsCOMPtr<nsIFile> file;
NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
getter_AddRefs(file));
UniquePtr<Classifier> classifier(new Classifier());
classifier->Open(*file);
RunTestInNewThread([&] () -> void {
nsresult rv = classifier->ApplyUpdates(&tableUpdates);
ASSERT_TRUE(NS_FAILED(rv));
});
}
static void
@ -145,6 +204,18 @@ testFullUpdate(PrefixStringMap& add)
testUpdate(tableUpdates, add);
}
static void
testPartialUpdate(PrefixStringMap& add,
nsTArray<uint32_t>& removal,
PrefixStringMap& expected)
{
nsTArray<TableUpdate*> tableUpdates;
GenerateUpdateData(false, add, removal, tableUpdates);
testUpdate(tableUpdates, expected);
}
TEST(UrlClassifierTableUpdateV4, FixLenghtPSetFullUpdate)
{
srand(time(NULL));
@ -188,3 +259,291 @@ TEST(UrlClassifierTableUpdateV4, MixedPSetFullUpdate)
Clear();
}
TEST(UrlClassifierTableUpdateV4, PartialUpdateWithRemoval)
{
_PrefixArray fArray, pArray, mergedArray;
PrefixStringMap fMap, pMap, mergedMap;
{
CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
PrefixArrayToPrefixStringMap(fArray, fMap);
testFullUpdate(fMap);
}
{
CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
RemoveIntersection(fArray, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
// Remove 1/5 of elements of original prefix set.
nsTArray<uint32_t> removal;
CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
RemoveElements(removal, fArray);
// Calculate the expected prefix map.
mergedArray.AppendElements(fArray);
mergedArray.AppendElements(pArray);
mergedArray.Sort();
PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
testPartialUpdate(pMap, removal, mergedMap);
}
Clear();
}
TEST(UrlClassifierTableUpdateV4, PartialUpdateWithoutRemoval)
{
_PrefixArray fArray, pArray, mergedArray;
PrefixStringMap fMap, pMap, mergedMap;
{
CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
PrefixArrayToPrefixStringMap(fArray, fMap);
testFullUpdate(fMap);
}
{
nsTArray<uint32_t> empty;
CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
RemoveIntersection(fArray, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
// Calculate the expected prefix map.
mergedArray.AppendElements(fArray);
mergedArray.AppendElements(pArray);
mergedArray.Sort();
PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
testPartialUpdate(pMap, empty, mergedMap);
}
Clear();
}
// Expect failure because partial update contains prefix already
// in old prefix set.
TEST(UrlClassifierTableUpdateV4, PartialUpdatePrefixAlreadyExist)
{
_PrefixArray fArray, pArray;
PrefixStringMap fMap, pMap;
{
CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
PrefixArrayToPrefixStringMap(fArray, fMap);
testFullUpdate(fMap);
}
{
nsTArray<uint32_t> empty;
nsTArray<TableUpdate*> tableUpdates;
// Pick one prefix from full update prefix and add it to partial update.
// This should result a failure when call ApplyUpdates.
pArray.AppendElement(fArray[rand() % fArray.Length()]);
CreateRandomSortedPrefixArray(200, 4, 32, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
GenerateUpdateData(false, pMap, empty, tableUpdates);
testUpdateFail(tableUpdates);
}
Clear();
}
// Test apply partial update directly without applying an full update first.
TEST(UrlClassifierTableUpdateV4, OnlyPartialUpdate)
{
_PrefixArray pArray;
PrefixStringMap pMap;
nsTArray<uint32_t> empty;
CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
testPartialUpdate(pMap, empty, pMap);
Clear();
}
// Test partial update without any ADD prefixes, only removalIndices.
TEST(UrlClassifierTableUpdateV4, PartialUpdateOnlyRemoval)
{
_PrefixArray fArray, pArray;
PrefixStringMap fMap, pMap, mergedMap;
{
CreateRandomSortedPrefixArray(5000, 4, 4, fArray);
CreateRandomSortedPrefixArray(1000, 5, 32, fArray);
PrefixArrayToPrefixStringMap(fArray, fMap);
testFullUpdate(fMap);
}
{
// Remove 1/5 of elements of original prefix set.
nsTArray<uint32_t> removal;
CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
RemoveElements(removal, fArray);
PrefixArrayToPrefixStringMap(fArray, mergedMap);
testPartialUpdate(pMap, removal, mergedMap);
}
Clear();
}
// Test one tableupdate array contains full update and multiple partial updates.
TEST(UrlClassifierTableUpdateV4, MultipleTableUpdates)
{
_PrefixArray fArray, pArray, mergedArray;
PrefixStringMap fMap, pMap, mergedMap;
{
nsTArray<uint32_t> empty;
nsTArray<TableUpdate*> tableUpdates;
// Generate first full udpate
CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
PrefixArrayToPrefixStringMap(fArray, fMap);
GenerateUpdateData(true, fMap, empty, tableUpdates);
// Generate second partial update
CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
RemoveIntersection(fArray, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
GenerateUpdateData(false, pMap, empty, tableUpdates);
// Generate thrid partial update
fArray.AppendElements(pArray);
fArray.Sort();
pArray.Clear();
CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
RemoveIntersection(fArray, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
// Remove 1/5 of elements of original prefix set.
nsTArray<uint32_t> removal;
CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
RemoveElements(removal, fArray);
GenerateUpdateData(false, pMap, removal, tableUpdates);
mergedArray.AppendElements(fArray);
mergedArray.AppendElements(pArray);
mergedArray.Sort();
PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
testUpdate(tableUpdates, mergedMap);
}
Clear();
}
// Test apply full update first, and then apply multiple partial updates
// in one tableupdate array.
TEST(UrlClassifierTableUpdateV4, MultiplePartialUpdateTableUpdates)
{
_PrefixArray fArray, pArray, mergedArray;
PrefixStringMap fMap, pMap, mergedMap;
{
// Generate first full udpate
CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
CreateRandomSortedPrefixArray(3000, 5, 32, fArray);
PrefixArrayToPrefixStringMap(fArray, fMap);
testFullUpdate(fMap);
}
{
nsTArray<uint32_t> removal;
nsTArray<TableUpdate*> tableUpdates;
// Generate first partial update
CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
RemoveIntersection(fArray, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
// Remove 1/5 of elements of original prefix set.
CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
RemoveElements(removal, fArray);
GenerateUpdateData(false, pMap, removal, tableUpdates);
fArray.AppendElements(pArray);
fArray.Sort();
pArray.Clear();
removal.Clear();
// Generate second partial update.
CreateRandomSortedPrefixArray(2000, 4, 4, pArray);
CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
RemoveIntersection(fArray, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
// Remove 1/5 of elements of original prefix set.
CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
RemoveElements(removal, fArray);
GenerateUpdateData(false, pMap, removal, tableUpdates);
mergedArray.AppendElements(fArray);
mergedArray.AppendElements(pArray);
mergedArray.Sort();
PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
testUpdate(tableUpdates, mergedMap);
}
Clear();
}
// Test removal indices are larger than the original prefix set.
TEST(UrlClassifierTableUpdateV4, RemovalIndexTooLarge)
{
_PrefixArray fArray, pArray;
PrefixStringMap fMap, pMap;
{
CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
PrefixArrayToPrefixStringMap(fArray, fMap);
testFullUpdate(fMap);
}
{
nsTArray<uint32_t> removal;
nsTArray<TableUpdate*> tableUpdates;
CreateRandomSortedPrefixArray(200, 4, 32, pArray);
RemoveIntersection(fArray, pArray);
PrefixArrayToPrefixStringMap(pArray, pMap);
for (uint32_t i = 0; i < fArray.Length() + 1 ;i++) {
removal.AppendElement(i);
}
GenerateUpdateData(false, pMap, removal, tableUpdates);
testUpdateFail(tableUpdates);
}
Clear();
}