diff options
Diffstat (limited to '')
-rw-r--r-- | toolkit/components/url-classifier/tests/gtest/Common.cpp | 219 |
1 files changed, 219 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/gtest/Common.cpp b/toolkit/components/url-classifier/tests/gtest/Common.cpp new file mode 100644 index 0000000000..172ded051f --- /dev/null +++ b/toolkit/components/url-classifier/tests/gtest/Common.cpp @@ -0,0 +1,219 @@ +/* -*- 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 "Common.h" + +#include "Classifier.h" +#include "HashStore.h" +#include "mozilla/Components.h" +#include "mozilla/SpinEventLoopUntil.h" +#include "mozilla/Unused.h" +#include "nsAppDirectoryServiceDefs.h" +#include "nsIThread.h" +#include "nsTArray.h" +#include "nsThreadUtils.h" +#include "nsUrlClassifierUtils.h" + +using namespace mozilla; +using namespace mozilla::safebrowsing; + +nsresult SyncApplyUpdates(TableUpdateArray& aUpdates) { + // We need to spin a new thread specifically because the callback + // will be on the caller thread. If we call Classifier::AsyncApplyUpdates + // and wait on the same thread, this function will never return. + + nsresult ret = NS_ERROR_FAILURE; + bool done = false; + auto onUpdateComplete = [&done, &ret](nsresult rv) { + // We are on the "ApplyUpdate" thread. Post an event to main thread + // so that we can avoid busy waiting on the main thread. + nsCOMPtr<nsIRunnable> r = + NS_NewRunnableFunction("SyncApplyUpdates", [&done, &ret, rv] { + ret = rv; + done = true; + }); + NS_DispatchToMainThread(r); + }; + + nsCOMPtr<nsIFile> file; + NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file)); + + nsCOMPtr<nsIRunnable> r = NS_NewRunnableFunction("SyncApplyUpdates", [&]() { + RefPtr<Classifier> classifier = new Classifier(); + classifier->Open(*file); + + nsresult rv = classifier->AsyncApplyUpdates(aUpdates, onUpdateComplete); + if (NS_FAILED(rv)) { + onUpdateComplete(rv); + } + }); + + nsCOMPtr<nsIThread> testingThread; + NS_NewNamedThread("ApplyUpdates", getter_AddRefs(testingThread)); + if (!testingThread) { + return NS_ERROR_FAILURE; + } + + testingThread->Dispatch(r, NS_DISPATCH_NORMAL); + + // NS_NewCheckSummedOutputStream in HashStore::WriteFile + // will synchronously init NS_CRYPTO_HASH_CONTRACTID on + // the main thread. As a result we have to keep processing + // pending event until |done| becomes true. If there's no + // more pending event, what we only can do is wait. + MOZ_ALWAYS_TRUE(SpinEventLoopUntil("url-classifier:SyncApplyUpdates"_ns, + [&]() { return done; })); + + return ret; +} + +already_AddRefed<nsIFile> GetFile(const nsTArray<nsString>& path) { + nsCOMPtr<nsIFile> file; + nsresult rv = + NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file)); + if (NS_WARN_IF(NS_FAILED(rv))) { + return nullptr; + } + + for (uint32_t i = 0; i < path.Length(); i++) { + file->Append(path[i]); + } + return file.forget(); +} + +void ApplyUpdate(TableUpdateArray& updates) { + // Force nsUrlClassifierUtils loading on main thread + // because nsIUrlClassifierDBService will not run in advance + // in gtest. + nsUrlClassifierUtils::GetInstance(); + + SyncApplyUpdates(updates); +} + +void ApplyUpdate(TableUpdate* update) { + TableUpdateArray updates = {update}; + ApplyUpdate(updates); +} + +nsresult PrefixArrayToPrefixStringMap(const _PrefixArray& aPrefixArray, + PrefixStringMap& aOut) { + aOut.Clear(); + + // Buckets are keyed by prefix length and contain an array of + // all prefixes of that length. + nsClassHashtable<nsUint32HashKey, _PrefixArray> table; + for (const auto& prefix : aPrefixArray) { + _PrefixArray* array = table.GetOrInsertNew(prefix.Length()); + array->AppendElement(prefix); + } + + // The resulting map entries will be a concatenation of all + // prefix data for the prefixes of a given size. + for (const auto& entry : table) { + uint32_t size = entry.GetKey(); + uint32_t count = entry.GetData()->Length(); + + auto str = MakeUnique<_Prefix>(); + str->SetLength(size * count); + + char* dst = str->BeginWriting(); + + entry.GetData()->Sort(); + for (uint32_t i = 0; i < count; i++) { + memcpy(dst, entry.GetData()->ElementAt(i).get(), size); + dst += size; + } + + aOut.InsertOrUpdate(size, std::move(str)); + } + + return NS_OK; +} + +nsresult PrefixArrayToAddPrefixArray(const _PrefixArray& aPrefixArray, + AddPrefixArray& aOut) { + aOut.Clear(); + + for (const auto& prefix : aPrefixArray) { + // Create prefix hash from string + AddPrefix* add = aOut.AppendElement(fallible); + if (!add) { + return NS_ERROR_OUT_OF_MEMORY; + } + + add->addChunk = 1; + add->prefix.Assign(prefix); + } + + EntrySort(aOut); + + return NS_OK; +} + +_Prefix CreatePrefixFromURL(const char* aURL, uint8_t aPrefixSize) { + return CreatePrefixFromURL(nsCString(aURL), aPrefixSize); +} + +_Prefix CreatePrefixFromURL(const nsCString& aURL, uint8_t aPrefixSize) { + Completion complete; + complete.FromPlaintext(aURL); + + _Prefix prefix; + prefix.Assign((const char*)complete.buf, aPrefixSize); + return prefix; +} + +void CheckContent(LookupCacheV4* aCache, const _PrefixArray& aPrefixArray) { + PrefixStringMap vlPSetMap; + aCache->GetPrefixes(vlPSetMap); + + PrefixStringMap expected; + PrefixArrayToPrefixStringMap(aPrefixArray, expected); + + for (const auto& entry : vlPSetMap) { + nsCString* expectedPrefix = expected.Get(entry.GetKey()); + nsCString* resultPrefix = entry.GetWeak(); + + ASSERT_TRUE(resultPrefix->Equals(*expectedPrefix)); + } +} + +nsresult BuildLookupCache(const RefPtr<Classifier>& classifier, + const nsACString& aTable, + _PrefixArray& aPrefixArray) { + RefPtr<LookupCache> cache = classifier->GetLookupCache(aTable, false); + if (!cache) { + return NS_ERROR_FAILURE; + } + + if (LookupCache::Cast<LookupCacheV4>(cache)) { + // V4 + RefPtr<LookupCacheV4> cacheV4 = LookupCache::Cast<LookupCacheV4>(cache); + + PrefixStringMap map; + PrefixArrayToPrefixStringMap(aPrefixArray, map); + return cacheV4->Build(map); + } else { + // V2 + RefPtr<LookupCacheV2> cacheV2 = LookupCache::Cast<LookupCacheV2>(cache); + + AddPrefixArray addPrefixes; + AddCompleteArray addComples; + + PrefixArrayToAddPrefixArray(aPrefixArray, addPrefixes); + return cacheV2->Build(addPrefixes, addComples); + } +} + +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; +} |