Bug 1305801 - Part 4: Store variable-length prefix to disk. r=francois, r=gcp

MozReview-Commit-ID: BMTGtgMuQdg

--HG--
extra : rebase_source : dc4975d9ac8d14ecfabf6fc1229fdb8e863cde46
This commit is contained in:
Dimi Lee 2016-09-19 11:51:01 +08:00
Родитель fcaa1bd296
Коммит 3ef7d52f53
9 изменённых файлов: 619 добавлений и 196 удалений

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

@ -517,17 +517,6 @@ Classifier::ApplyUpdates(nsTArray<TableUpdate*>* aUpdates)
{
ScopedUpdatesClearer scopedUpdatesClearer(aUpdates);
// In order to prevent any premature update code from being
// run against V4 updates, we bail out as early as possible
// if aUpdates is using V4.
for (auto update : *aUpdates) {
if (update && TableUpdate::Cast<TableUpdateV4>(update)) {
LOG(("V4 update is not supported yet."));
// TODO: Bug 1283009 - Supports applying table udpate V4.
return NS_ERROR_NOT_IMPLEMENTED;
}
}
LOG(("Backup before update."));
rv = BackupTables();
@ -540,7 +529,13 @@ Classifier::ApplyUpdates(nsTArray<TableUpdate*>* aUpdates)
if ((*aUpdates)[i]) {
// Run all updates for one table
nsCString updateTable(aUpdates->ElementAt(i)->TableName());
rv = UpdateHashStore(aUpdates, updateTable);
if (TableUpdate::Cast<TableUpdateV2>((*aUpdates)[i])) {
rv = UpdateHashStore(aUpdates, updateTable);
} else {
rv = UpdateTableV4(aUpdates, updateTable);
}
if (NS_FAILED(rv)) {
if (rv != NS_ERROR_OUT_OF_MEMORY) {
Reset();
@ -857,7 +852,8 @@ Classifier::UpdateHashStore(nsTArray<TableUpdate*>* aUpdates,
NS_ENSURE_SUCCESS(rv, rv);
// Read the part of the store that is (only) in the cache
LookupCache *lookupCache = GetLookupCache(store.TableName());
LookupCacheV2* lookupCache =
LookupCache::Cast<LookupCacheV2>(GetLookupCache(store.TableName()));
if (!lookupCache) {
return NS_ERROR_FAILURE;
}
@ -922,7 +918,7 @@ Classifier::UpdateHashStore(nsTArray<TableUpdate*>* aUpdates,
NS_ENSURE_SUCCESS(rv, rv);
#if defined(DEBUG)
lookupCache->Dump();
lookupCache->DumpCompletions();
#endif
rv = lookupCache->WriteFile();
NS_ENSURE_SUCCESS(rv, rv);
@ -934,6 +930,62 @@ Classifier::UpdateHashStore(nsTArray<TableUpdate*>* aUpdates,
return NS_OK;
}
nsresult
Classifier::UpdateTableV4(nsTArray<TableUpdate*>* aUpdates,
const nsACString& aTable)
{
LOG(("Classifier::UpdateTableV4(%s)", PromiseFlatCString(aTable).get()));
if (!CheckValidUpdate(aUpdates, aTable)) {
return NS_OK;
}
LookupCacheV4* lookupCache =
LookupCache::Cast<LookupCacheV4>(GetLookupCache(aTable));
if (!lookupCache) {
return NS_ERROR_FAILURE;
}
PrefixStringMap prefixes;
for (uint32_t i = 0; i < aUpdates->Length(); i++) {
TableUpdate *update = aUpdates->ElementAt(i);
if (!update || !update->TableName().Equals(aTable)) {
continue;
}
auto updateV4 = TableUpdate::Cast<TableUpdateV4>(update);
NS_ENSURE_TRUE(updateV4, NS_ERROR_FAILURE);
if (updateV4->IsFullUpdate()) {
prefixes.Clear();
TableUpdateV4::PrefixesStringMap& map = updateV4->Prefixes();
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);
}
} else {
// TODO: Bug 1287058, partial update
}
aUpdates->ElementAt(i) = nullptr;
}
nsresult rv = lookupCache->Build(prefixes);
NS_ENSURE_SUCCESS(rv, rv);
rv = lookupCache->WriteFile();
NS_ENSURE_SUCCESS(rv, rv);
int64_t now = (PR_Now() / PR_USEC_PER_SEC);
LOG(("Successfully updated %s\n", PromiseFlatCString(aTable).get()));
mTableFreshness.Put(aTable, now);
return NS_OK;
}
nsresult
Classifier::UpdateCache(TableUpdate* aUpdate)
{
@ -945,7 +997,9 @@ Classifier::UpdateCache(TableUpdate* aUpdate)
LOG(("Classifier::UpdateCache(%s)", table.get()));
LookupCache *lookupCache = GetLookupCache(table);
NS_ENSURE_TRUE(lookupCache, NS_ERROR_FAILURE);
if (!lookupCache) {
return NS_ERROR_FAILURE;
}
auto updateV2 = TableUpdate::Cast<TableUpdateV2>(aUpdate);
lookupCache->AddCompletionsToCache(updateV2->AddCompletes());
@ -966,7 +1020,16 @@ Classifier::GetLookupCache(const nsACString& aTable)
}
}
UniquePtr<LookupCache> cache(new LookupCache(aTable, mRootStoreDirectory));
// TODO : Bug 1302600, It would be better if we have a more general non-main
// thread method to convert table name to protocol version. Currently
// we can only know this by checking if the table name ends with '-proto'.
UniquePtr<LookupCache> cache;
if (StringEndsWith(aTable, NS_LITERAL_CSTRING("-proto"))) {
cache = MakeUnique<LookupCacheV4>(aTable, mRootStoreDirectory);
} else {
cache = MakeUnique<LookupCacheV2>(aTable, mRootStoreDirectory);
}
nsresult rv = cache->Init();
if (NS_FAILED(rv)) {
return nullptr;
@ -988,7 +1051,8 @@ Classifier::ReadNoiseEntries(const Prefix& aPrefix,
uint32_t aCount,
PrefixArray* aNoiseEntries)
{
LookupCache *cache = GetLookupCache(aTableName);
// TODO : Bug 1297962, support adding noise for v4
LookupCacheV2 *cache = static_cast<LookupCacheV2*>(GetLookupCache(aTableName));
if (!cache) {
return NS_ERROR_FAILURE;
}

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

@ -115,6 +115,9 @@ private:
nsresult UpdateHashStore(nsTArray<TableUpdate*>* aUpdates,
const nsACString& aTable);
nsresult UpdateTableV4(nsTArray<TableUpdate*>* aUpdates,
const nsACString& aTable);
nsresult UpdateCache(TableUpdate* aUpdates);
LookupCache *GetLookupCache(const nsACString& aTable);

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

@ -153,6 +153,7 @@ public:
public:
explicit TableUpdateV4(const nsACString& aTable)
: TableUpdate(aTable)
, mFullUpdate(false)
{
}
@ -161,18 +162,21 @@ public:
return mPrefixesMap.IsEmpty() && mRemovalIndiceArray.IsEmpty();
}
bool IsFullUpdate() const { return mFullUpdate; }
PrefixesStringMap& Prefixes() { return mPrefixesMap; }
RemovalIndiceArray& RemovalIndices() { return mRemovalIndiceArray; }
// For downcasting.
static const int TAG = 4;
void SetFullUpdate(bool aIsFullUpdate) { mFullUpdate = aIsFullUpdate; }
void NewPrefixes(int32_t aSize, std::string& aPrefixes);
void NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices);
private:
virtual int Tag() const override { return TAG; }
bool mFullUpdate;
PrefixesStringMap mPrefixesMap;
RemovalIndiceArray mRemovalIndiceArray;
};

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

@ -40,6 +40,9 @@ extern mozilla::LazyLogModule gUrlClassifierDbServiceLog;
namespace mozilla {
namespace safebrowsing {
const int LookupCacheV2::VER = 2;
const int LookupCacheV4::VER = 4;
LookupCache::LookupCache(const nsACString& aTableName, nsIFile* aRootStoreDir)
: mPrimed(false)
, mTableName(aTableName)
@ -48,29 +51,11 @@ LookupCache::LookupCache(const nsACString& aTableName, nsIFile* aRootStoreDir)
UpdateRootDirHandle(mRootStoreDirectory);
}
nsresult
LookupCache::Init()
{
mPrefixSet = new nsUrlClassifierPrefixSet();
nsresult rv = mPrefixSet->Init(mTableName);
NS_ENSURE_SUCCESS(rv, rv);
return NS_OK;
}
LookupCache::~LookupCache()
{
}
nsresult
LookupCache::Open()
{
LOG(("Reading Completions"));
nsresult rv = ReadCompletions();
NS_ENSURE_SUCCESS(rv, rv);
LOG(("Loading PrefixSet"));
rv = LoadPrefixSet();
nsresult rv = LoadPrefixSet();
NS_ENSURE_SUCCESS(rv, rv);
return NS_OK;
@ -125,32 +110,6 @@ LookupCache::Reset()
return NS_OK;
}
nsresult
LookupCache::Build(AddPrefixArray& aAddPrefixes,
AddCompleteArray& aAddCompletes)
{
Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_COMPLETIONS,
static_cast<uint32_t>(aAddCompletes.Length()));
mUpdateCompletions.Clear();
mUpdateCompletions.SetCapacity(aAddCompletes.Length());
for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
mUpdateCompletions.AppendElement(aAddCompletes[i].CompleteHash());
}
aAddCompletes.Clear();
mUpdateCompletions.Sort();
Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_PREFIXES,
static_cast<uint32_t>(aAddPrefixes.Length()));
nsresult rv = ConstructPrefixSet(aAddPrefixes);
NS_ENSURE_SUCCESS(rv, rv);
mPrimed = true;
return NS_OK;
}
nsresult
LookupCache::AddCompletionsToCache(AddCompleteArray& aAddCompletes)
{
@ -177,50 +136,8 @@ LookupCache::DumpCache()
LOG(("Caches: %s", str.get()));
}
}
void
LookupCache::Dump()
{
if (!LOG_ENABLED())
return;
for (uint32_t i = 0; i < mUpdateCompletions.Length(); i++) {
nsAutoCString str;
mUpdateCompletions[i].ToHexString(str);
LOG(("Update: %s", str.get()));
}
}
#endif
nsresult
LookupCache::Has(const Completion& aCompletion,
bool* aHas, bool* aComplete)
{
*aHas = *aComplete = false;
uint32_t prefix = aCompletion.ToUint32();
bool found;
nsresult rv = mPrefixSet->Contains(prefix, &found);
NS_ENSURE_SUCCESS(rv, rv);
LOG(("Probe in %s: %X, found %d", mTableName.get(), prefix, found));
if (found) {
*aHas = true;
}
// TODO: We may need to distinguish completions found in cache or update in the future
if ((mGetHashCache.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex) ||
(mUpdateCompletions.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex)) {
LOG(("Complete in %s", mTableName.get()));
*aComplete = true;
*aHas = true;
}
return NS_OK;
}
nsresult
LookupCache::WriteFile()
{
@ -231,7 +148,7 @@ LookupCache::WriteFile()
rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
NS_ENSURE_SUCCESS(rv, rv);
rv = mPrefixSet->StoreToFile(psFile);
rv = StoreToFile(psFile);
NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "failed to store the prefixset");
return NS_OK;
@ -241,41 +158,16 @@ void
LookupCache::ClearAll()
{
ClearCache();
ClearUpdatedCompletions();
mPrefixSet->SetPrefixes(nullptr, 0);
ClearPrefixes();
mPrimed = false;
}
void
LookupCache::ClearUpdatedCompletions()
{
mUpdateCompletions.Clear();
}
void
LookupCache::ClearCache()
{
mGetHashCache.Clear();
}
nsresult
LookupCache::ReadCompletions()
{
HashStore store(mTableName, mRootStoreDirectory);
nsresult rv = store.Open();
NS_ENSURE_SUCCESS(rv, rv);
mUpdateCompletions.Clear();
const AddCompleteArray& addComplete = store.AddCompletes();
for (uint32_t i = 0; i < addComplete.Length(); i++) {
mUpdateCompletions.AppendElement(addComplete[i].complete);
}
return NS_OK;
}
/* static */ bool
LookupCache::IsCanonicalizedIP(const nsACString& aHost)
{
@ -452,9 +344,178 @@ LookupCache::GetHostKeys(const nsACString& aSpec,
return NS_OK;
}
bool LookupCache::IsPrimed()
nsresult
LookupCache::LoadPrefixSet()
{
return mPrimed;
nsCOMPtr<nsIFile> psFile;
nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
NS_ENSURE_SUCCESS(rv, rv);
rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
NS_ENSURE_SUCCESS(rv, rv);
bool exists;
rv = psFile->Exists(&exists);
NS_ENSURE_SUCCESS(rv, rv);
if (exists) {
LOG(("stored PrefixSet exists, loading from disk"));
rv = LoadFromFile(psFile);
if (NS_FAILED(rv)) {
if (rv == NS_ERROR_FILE_CORRUPTED) {
Reset();
}
return rv;
}
mPrimed = true;
} else {
LOG(("no (usable) stored PrefixSet found"));
}
#ifdef DEBUG
if (mPrimed) {
uint32_t size = SizeOfPrefixSet();
LOG(("SB tree done, size = %d bytes\n", size));
}
#endif
return NS_OK;
}
nsresult
LookupCacheV2::Init()
{
mPrefixSet = new nsUrlClassifierPrefixSet();
nsresult rv = mPrefixSet->Init(mTableName);
NS_ENSURE_SUCCESS(rv, rv);
return NS_OK;
}
nsresult
LookupCacheV2::Open()
{
nsresult rv = LookupCache::Open();
NS_ENSURE_SUCCESS(rv, rv);
LOG(("Reading Completions"));
rv = ReadCompletions();
NS_ENSURE_SUCCESS(rv, rv);
return NS_OK;
}
void
LookupCacheV2::ClearAll()
{
LookupCache::ClearAll();
mUpdateCompletions.Clear();
}
nsresult
LookupCacheV2::Has(const Completion& aCompletion,
bool* aHas, bool* aComplete)
{
*aHas = *aComplete = false;
uint32_t prefix = aCompletion.ToUint32();
bool found;
nsresult rv = mPrefixSet->Contains(prefix, &found);
NS_ENSURE_SUCCESS(rv, rv);
LOG(("Probe in %s: %X, found %d", mTableName.get(), prefix, found));
if (found) {
*aHas = true;
}
if ((mGetHashCache.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex) ||
(mUpdateCompletions.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex)) {
LOG(("Complete in %s", mTableName.get()));
*aComplete = true;
*aHas = true;
}
return NS_OK;
}
nsresult
LookupCacheV2::Build(AddPrefixArray& aAddPrefixes,
AddCompleteArray& aAddCompletes)
{
Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_COMPLETIONS,
static_cast<uint32_t>(aAddCompletes.Length()));
mUpdateCompletions.Clear();
mUpdateCompletions.SetCapacity(aAddCompletes.Length());
for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
mUpdateCompletions.AppendElement(aAddCompletes[i].CompleteHash());
}
aAddCompletes.Clear();
mUpdateCompletions.Sort();
Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_PREFIXES,
static_cast<uint32_t>(aAddPrefixes.Length()));
nsresult rv = ConstructPrefixSet(aAddPrefixes);
NS_ENSURE_SUCCESS(rv, rv);
mPrimed = true;
return NS_OK;
}
nsresult
LookupCacheV2::GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes)
{
if (!mPrimed) {
// This can happen if its a new table, so no error.
LOG(("GetPrefixes from empty LookupCache"));
return NS_OK;
}
return mPrefixSet->GetPrefixesNative(aAddPrefixes);
}
nsresult
LookupCacheV2::ReadCompletions()
{
HashStore store(mTableName, mRootStoreDirectory);
nsresult rv = store.Open();
NS_ENSURE_SUCCESS(rv, rv);
mUpdateCompletions.Clear();
const AddCompleteArray& addComplete = store.AddCompletes();
for (uint32_t i = 0; i < addComplete.Length(); i++) {
mUpdateCompletions.AppendElement(addComplete[i].complete);
}
return NS_OK;
}
nsresult
LookupCacheV2::ClearPrefixes()
{
return mPrefixSet->SetPrefixes(nullptr, 0);
}
nsresult
LookupCacheV2::StoreToFile(nsIFile* aFile)
{
return mPrefixSet->StoreToFile(aFile);
}
nsresult
LookupCacheV2::LoadFromFile(nsIFile* aFile)
{
return mPrefixSet->LoadFromFile(aFile);
}
size_t
LookupCacheV2::SizeOfPrefixSet()
{
return mPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
}
#ifdef DEBUG
@ -478,7 +539,7 @@ static void EnsureSorted(T* aArray)
#endif
nsresult
LookupCache::ConstructPrefixSet(AddPrefixArray& aAddPrefixes)
LookupCacheV2::ConstructPrefixSet(AddPrefixArray& aAddPrefixes)
{
Telemetry::AutoTimer<Telemetry::URLCLASSIFIER_PS_CONSTRUCT_TIME> timer;
@ -512,55 +573,71 @@ LookupCache::ConstructPrefixSet(AddPrefixArray& aAddPrefixes)
return NS_OK;
}
nsresult
LookupCache::LoadPrefixSet()
#if defined(DEBUG)
void
LookupCacheV2::DumpCompletions()
{
nsCOMPtr<nsIFile> psFile;
nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
NS_ENSURE_SUCCESS(rv, rv);
if (!LOG_ENABLED())
return;
rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
NS_ENSURE_SUCCESS(rv, rv);
bool exists;
rv = psFile->Exists(&exists);
NS_ENSURE_SUCCESS(rv, rv);
if (exists) {
LOG(("stored PrefixSet exists, loading from disk"));
rv = mPrefixSet->LoadFromFile(psFile);
if (NS_FAILED(rv)) {
if (rv == NS_ERROR_FILE_CORRUPTED) {
Reset();
}
return rv;
}
mPrimed = true;
} else {
LOG(("no (usable) stored PrefixSet found"));
}
#ifdef DEBUG
if (mPrimed) {
uint32_t size = mPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
LOG(("SB tree done, size = %d bytes\n", size));
for (uint32_t i = 0; i < mUpdateCompletions.Length(); i++) {
nsAutoCString str;
mUpdateCompletions[i].ToHexString(str);
LOG(("Update: %s", str.get()));
}
}
#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
LookupCache::GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes)
LookupCacheV4::Has(const Completion& aCompletion,
bool* aHas, bool* aComplete)
{
if (!mPrimed) {
// This can happen if its a new table, so no error.
LOG(("GetPrefixes from empty LookupCache"));
return NS_OK;
}
return mPrefixSet->GetPrefixesNative(aAddPrefixes);
*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

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

@ -14,6 +14,7 @@
#include "nsIFileStreams.h"
#include "mozilla/RefPtr.h"
#include "nsUrlClassifierPrefixSet.h"
#include "VariableLengthPrefixSet.h"
#include "mozilla/Logging.h"
namespace mozilla {
@ -96,50 +97,58 @@ public:
nsTArray<nsCString>* aHostKeys);
LookupCache(const nsACString& aTableName, nsIFile* aStoreFile);
~LookupCache();
virtual ~LookupCache() {}
const nsCString &TableName() const { return mTableName; }
nsresult Init();
nsresult Open();
// The directory handle where we operate will
// be moved away when a backup is made.
nsresult UpdateRootDirHandle(nsIFile* aRootStoreDirectory);
// This will Clear() the passed arrays when done.
nsresult Build(AddPrefixArray& aAddPrefixes,
AddCompleteArray& aAddCompletes);
nsresult AddCompletionsToCache(AddCompleteArray& aAddCompletes);
nsresult GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes);
void ClearUpdatedCompletions();
// Write data stored in lookup cache to disk.
nsresult WriteFile();
// Clear completions retrieved from gethash request.
void ClearCache();
bool IsPrimed() const { return mPrimed; };
#if DEBUG
void DumpCache();
void Dump();
#endif
nsresult WriteFile();
nsresult Has(const Completion& aCompletion,
bool* aHas, bool* aComplete);
bool IsPrimed();
virtual nsresult Open();
virtual nsresult Init() = 0;
virtual nsresult ClearPrefixes() = 0;
virtual nsresult Has(const Completion& aCompletion,
bool* aHas, bool* aComplete) = 0;
template<typename T>
static T* Cast(LookupCache* aThat) {
return (T::VER == aThat->Ver() ? reinterpret_cast<T*>(aThat) : nullptr);
}
private:
void ClearAll();
nsresult Reset();
nsresult ReadCompletions();
nsresult LoadPrefixSet();
nsresult LoadCompletions();
// Construct a Prefix Set with known prefixes.
// This will Clear() aAddPrefixes when done.
nsresult ConstructPrefixSet(AddPrefixArray& aAddPrefixes);
virtual nsresult StoreToFile(nsIFile* aFile) = 0;
virtual nsresult LoadFromFile(nsIFile* aFile) = 0;
virtual size_t SizeOfPrefixSet() = 0;
virtual int Ver() const = 0;
protected:
virtual void ClearAll();
bool mPrimed;
nsCString mTableName;
nsCOMPtr<nsIFile> mRootStoreDirectory;
nsCOMPtr<nsIFile> mStoreDirectory;
// Set of prefixes known to be in the database
RefPtr<nsUrlClassifierPrefixSet> mPrefixSet;
// Full length hashes obtained in update request
CompletionArray mUpdateCompletions;
// Full length hashes obtained in gethash request
CompletionArray mGetHashCache;
@ -147,6 +156,79 @@ private:
friend class PerProviderDirectoryTestUtils;
};
class LookupCacheV2 final : public LookupCache
{
public:
explicit LookupCacheV2(const nsACString& aTableName, nsIFile* aStoreFile)
: LookupCache(aTableName, aStoreFile) {}
~LookupCacheV2() {}
virtual nsresult Init() override;
virtual nsresult Open() override;
virtual void ClearAll() override;
virtual nsresult Has(const Completion& aCompletion,
bool* aHas, bool* aComplete) override;
nsresult Build(AddPrefixArray& aAddPrefixes,
AddCompleteArray& aAddCompletes);
nsresult GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes);
#if DEBUG
void DumpCompletions();
#endif
static const int VER;
protected:
nsresult ReadCompletions();
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; }
// Construct a Prefix Set with known prefixes.
// This will Clear() aAddPrefixes when done.
nsresult ConstructPrefixSet(AddPrefixArray& aAddPrefixes);
// Full length hashes obtained in update request
CompletionArray mUpdateCompletions;
// Set of prefixes known to be in the database
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

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

@ -873,6 +873,8 @@ ProtocolParserProtobuf::ProcessOneResponse(const ListUpdateResponse& aResponse)
PARSER_LOG(("* listName: %s\n", listName.get()));
PARSER_LOG(("* newState: %s\n", aResponse.new_client_state().c_str()));
PARSER_LOG(("* isFullUpdate: %s\n", (isFullUpdate ? "yes" : "no")));
tuV4->SetFullUpdate(isFullUpdate);
ProcessAdditionOrRemoval(*tuV4, aResponse.additions(), true /*aIsAddition*/);
ProcessAdditionOrRemoval(*tuV4, aResponse.removals(), false);
PARSER_LOG(("\n\n"));

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

@ -81,12 +81,12 @@ TEST(PerProviderDirectory, LookupCache)
// For V2 tables (NOT ending with '-proto'), root directory should be
// used as the private store.
VerifyPrivateStorePath<LookupCache>("goog-phish-shavar", "google", rootDir, false);
VerifyPrivateStorePath<LookupCacheV2>("goog-phish-shavar", "google", rootDir, false);
// For V4 tables, if provider is found, use per-provider subdirectory;
// If not found, use root directory.
VerifyPrivateStorePath<LookupCache>("goog-noprovider-proto", "", rootDir, false);
VerifyPrivateStorePath<LookupCache>("goog-phish-proto", "google4", rootDir, true);
VerifyPrivateStorePath<LookupCacheV4>("goog-noprovider-proto", "", rootDir, false);
VerifyPrivateStorePath<LookupCacheV4>("goog-phish-proto", "google4", rootDir, true);
});
}

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

@ -0,0 +1,190 @@
#include "Classifier.h"
#include "HashStore.h"
#include "nsAppDirectoryServiceDefs.h"
#include "nsIFile.h"
#include "nsIThread.h"
#include "string.h"
#include "gtest/gtest.h"
using namespace mozilla;
using namespace mozilla::safebrowsing;
typedef nsCString _Prefix;
typedef nsTArray<_Prefix> _PrefixArray;
static void
PrefixArrayToPrefixStringMap(const _PrefixArray& prefixArray,
PrefixStringMap& outMap)
{
outMap.Clear();
for (uint32_t i = 0; i < prefixArray.Length(); i++) {
const _Prefix& prefix = prefixArray[i];
nsCString* prefixString = outMap.LookupOrAdd(prefix.Length());
prefixString->Append(prefix.BeginReading(), prefix.Length());
}
}
// 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,
uint32_t MAX,
_PrefixArray& outArray)
{
outArray.SetCapacity(outArray.Length() + N);
const 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);
while (true) {
char* dst = prefix.BeginWriting();
for (uint32_t j = 0; j < prefixSize; j++) {
dst[j] = rand() % 256;
}
if (!outArray.Contains(prefix)) {
outArray.AppendElement(prefix);
break;
}
}
}
outArray.Sort();
}
// Function to generate TableUpdateV4.
static void
GenerateUpdateData(bool fullUpdate,
PrefixStringMap& add,
nsTArray<uint32_t>& removal,
nsTArray<TableUpdate*>& tableUpdates)
{
TableUpdateV4* tableUpdate = new TableUpdateV4(NS_LITERAL_CSTRING("gtest-malware-proto"));
tableUpdate->SetFullUpdate(fullUpdate);
for (auto iter = add.ConstIter(); !iter.Done(); iter.Next()) {
nsCString* pstring = iter.Data();
std::string str(pstring->BeginReading(), pstring->Length());
tableUpdate->NewPrefixes(iter.Key(), str);
}
tableUpdate->NewRemovalIndices(removal.Elements(), removal.Length());
tableUpdates.AppendElement(tableUpdate);
}
static void
VerifyPrefixSet(PrefixStringMap& expected)
{
// Verify the prefix set is written to disk.
nsCOMPtr<nsIFile> file;
NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
getter_AddRefs(file));
file->AppendRelativeNativePath(NS_LITERAL_CSTRING("safebrowsing/gtest-malware-proto.pset"));
RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;
load->Init(NS_LITERAL_CSTRING("gtest-malware-proto"));
PrefixStringMap out;
load->LoadFromFile(file);
load->GetPrefixes(out);
for (auto iter = expected.ConstIter(); !iter.Done(); iter.Next()) {
nsCString* str1 = iter.Data();
nsCString* str2 = out.Get(iter.Key());
ASSERT_TRUE(*str1 == *str2);
}
}
static void
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);
}
static void
testUpdate(nsTArray<TableUpdate*>& tableUpdates,
PrefixStringMap& expected)
{
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(rv == NS_OK);
VerifyPrefixSet(expected);
});
}
static void
testFullUpdate(PrefixStringMap& add)
{
nsTArray<uint32_t> empty;
nsTArray<TableUpdate*> tableUpdates;
GenerateUpdateData(true, add, empty, tableUpdates);
testUpdate(tableUpdates, add);
}
TEST(UrlClassifierTableUpdateV4, FixLenghtPSetFullUpdate)
{
srand(time(NULL));
_PrefixArray array;
PrefixStringMap map;
CreateRandomSortedPrefixArray(5000, 4, 4, array);
PrefixArrayToPrefixStringMap(array, map);
testFullUpdate(map);
Clear();
}
TEST(UrlClassifierTableUpdateV4, VariableLenghtPSetFullUpdate)
{
_PrefixArray array;
PrefixStringMap map;
CreateRandomSortedPrefixArray(5000, 5, 32, array);
PrefixArrayToPrefixStringMap(array, map);
testFullUpdate(map);
Clear();
}
// This test contain both variable length prefix set and fixed-length prefix set
TEST(UrlClassifierTableUpdateV4, MixedPSetFullUpdate)
{
_PrefixArray array;
PrefixStringMap map;
CreateRandomSortedPrefixArray(5000, 4, 4, array);
CreateRandomSortedPrefixArray(1000, 5, 32, array);
PrefixArrayToPrefixStringMap(array, map);
testFullUpdate(map);
Clear();
}

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

@ -14,6 +14,7 @@ UNIFIED_SOURCES += [
'TestSafebrowsingHash.cpp',
'TestSafeBrowsingProtobuf.cpp',
'TestTable.cpp',
'TestUrlClassifierTableUpdateV4.cpp',
'TestUrlClassifierUtils.cpp',
'TestVariableLengthPrefixSet.cpp',
]