michael@0: //* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: #include "LookupCache.h"
michael@0: #include "HashStore.h"
michael@0: #include "nsISeekableStream.h"
michael@0: #include "mozilla/Telemetry.h"
michael@0: #include "prlog.h"
michael@0: #include "prprf.h"
michael@0: 
michael@0: // We act as the main entry point for all the real lookups,
michael@0: // so note that those are not done to the actual HashStore.
michael@0: // The latter solely exists to store the data needed to handle
michael@0: // the updates from the protocol.
michael@0: 
michael@0: // This module has its own store, which stores the Completions,
michael@0: // mostly caching lookups that have happened over the net.
michael@0: // The prefixes are cached/checked by looking them up in the
michael@0: // PrefixSet.
michael@0: 
michael@0: // Data format for the ".cache" files:
michael@0: //    uint32_t magic           Identify the file type
michael@0: //    uint32_t version         Version identifier for file format
michael@0: //    uint32_t numCompletions  Amount of completions stored
michael@0: //    0...numCompletions     256-bit Completions
michael@0: 
michael@0: // Name of the lookupcomplete cache
michael@0: #define CACHE_SUFFIX ".cache"
michael@0: 
michael@0: // Name of the persistent PrefixSet storage
michael@0: #define PREFIXSET_SUFFIX  ".pset"
michael@0: 
michael@0: // NSPR_LOG_MODULES=UrlClassifierDbService:5
michael@0: extern PRLogModuleInfo *gUrlClassifierDbServiceLog;
michael@0: #if defined(PR_LOGGING)
michael@0: #define LOG(args) PR_LOG(gUrlClassifierDbServiceLog, PR_LOG_DEBUG, args)
michael@0: #define LOG_ENABLED() PR_LOG_TEST(gUrlClassifierDbServiceLog, 4)
michael@0: #else
michael@0: #define LOG(args)
michael@0: #define LOG_ENABLED() (false)
michael@0: #endif
michael@0: 
michael@0: namespace mozilla {
michael@0: namespace safebrowsing {
michael@0: 
michael@0: const uint32_t LOOKUPCACHE_MAGIC = 0x1231af3e;
michael@0: const uint32_t CURRENT_VERSION = 2;
michael@0: 
michael@0: LookupCache::LookupCache(const nsACString& aTableName, nsIFile* aStoreDir)
michael@0:   : mPrimed(false)
michael@0:   , mTableName(aTableName)
michael@0:   , mStoreDirectory(aStoreDir)
michael@0: {
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::Init()
michael@0: {
michael@0:   mPrefixSet = new nsUrlClassifierPrefixSet();
michael@0:   nsresult rv = mPrefixSet->Init(mTableName);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: LookupCache::~LookupCache()
michael@0: {
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::Open()
michael@0: {
michael@0:   nsCOMPtr<nsIFile> storeFile;
michael@0: 
michael@0:   nsresult rv = mStoreDirectory->Clone(getter_AddRefs(storeFile));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = storeFile->AppendNative(mTableName + NS_LITERAL_CSTRING(CACHE_SUFFIX));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   nsCOMPtr<nsIInputStream> inputStream;
michael@0:   rv = NS_NewLocalFileInputStream(getter_AddRefs(inputStream), storeFile,
michael@0:                                   PR_RDONLY | nsIFile::OS_READAHEAD);
michael@0: 
michael@0:   if (NS_FAILED(rv) && rv != NS_ERROR_FILE_NOT_FOUND) {
michael@0:     Reset();
michael@0:     return rv;
michael@0:   }
michael@0: 
michael@0:   if (rv == NS_ERROR_FILE_NOT_FOUND) {
michael@0:     // Simply lacking a .cache file is a recoverable error,
michael@0:     // as unlike the .pset/.sbstore files it is a pure cache.
michael@0:     // Just create a new empty one.
michael@0:     ClearCompleteCache();
michael@0:   } else {
michael@0:     // Read in the .cache file
michael@0:     rv = ReadHeader(inputStream);
michael@0:     NS_ENSURE_SUCCESS(rv, rv);
michael@0:     LOG(("ReadCompletions"));
michael@0:     rv = ReadCompletions(inputStream);
michael@0:     NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:     rv = inputStream->Close();
michael@0:     NS_ENSURE_SUCCESS(rv, rv);
michael@0:   }
michael@0: 
michael@0:   LOG(("Loading PrefixSet"));
michael@0:   rv = LoadPrefixSet();
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::UpdateDirHandle(nsIFile* aStoreDirectory)
michael@0: {
michael@0:   return aStoreDirectory->Clone(getter_AddRefs(mStoreDirectory));
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::Reset()
michael@0: {
michael@0:   LOG(("LookupCache resetting"));
michael@0: 
michael@0:   nsCOMPtr<nsIFile> storeFile;
michael@0:   nsCOMPtr<nsIFile> prefixsetFile;
michael@0:   nsresult rv = mStoreDirectory->Clone(getter_AddRefs(storeFile));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   rv = mStoreDirectory->Clone(getter_AddRefs(prefixsetFile));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = storeFile->AppendNative(mTableName + NS_LITERAL_CSTRING(CACHE_SUFFIX));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   rv = prefixsetFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = storeFile->Remove(false);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   rv = prefixsetFile->Remove(false);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   ClearAll();
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: 
michael@0: nsresult
michael@0: LookupCache::Build(AddPrefixArray& aAddPrefixes,
michael@0:                    AddCompleteArray& aAddCompletes)
michael@0: {
michael@0:   Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_COMPLETIONS,
michael@0:                         static_cast<uint32_t>(aAddCompletes.Length()));
michael@0: 
michael@0:   mCompletions.Clear();
michael@0:   mCompletions.SetCapacity(aAddCompletes.Length());
michael@0:   for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
michael@0:     mCompletions.AppendElement(aAddCompletes[i].CompleteHash());
michael@0:   }
michael@0:   aAddCompletes.Clear();
michael@0:   mCompletions.Sort();
michael@0: 
michael@0:   Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_PREFIXES,
michael@0:                         static_cast<uint32_t>(aAddPrefixes.Length()));
michael@0: 
michael@0:   nsresult rv = ConstructPrefixSet(aAddPrefixes);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   mPrimed = true;
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: #if defined(DEBUG) && defined(PR_LOGGING)
michael@0: void
michael@0: LookupCache::Dump()
michael@0: {
michael@0:   if (!LOG_ENABLED())
michael@0:     return;
michael@0: 
michael@0:   for (uint32_t i = 0; i < mCompletions.Length(); i++) {
michael@0:     nsAutoCString str;
michael@0:     mCompletions[i].ToString(str);
michael@0:     LOG(("Completion: %s", str.get()));
michael@0:   }
michael@0: }
michael@0: #endif
michael@0: 
michael@0: nsresult
michael@0: LookupCache::Has(const Completion& aCompletion,
michael@0:                  bool* aHas, bool* aComplete)
michael@0: {
michael@0:   *aHas = *aComplete = false;
michael@0: 
michael@0:   uint32_t prefix = aCompletion.ToUint32();
michael@0: 
michael@0:   bool found;
michael@0:   nsresult rv = mPrefixSet->Contains(prefix, &found);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   LOG(("Probe in %s: %X, found %d", mTableName.get(), prefix, found));
michael@0: 
michael@0:   if (found) {
michael@0:     *aHas = true;
michael@0:   }
michael@0: 
michael@0:   if (mCompletions.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex) {
michael@0:     LOG(("Complete in %s", mTableName.get()));
michael@0:     *aComplete = true;
michael@0:     *aHas = true;
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::WriteFile()
michael@0: {
michael@0:   nsCOMPtr<nsIFile> storeFile;
michael@0:   nsresult rv = mStoreDirectory->Clone(getter_AddRefs(storeFile));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   rv = storeFile->AppendNative(mTableName + NS_LITERAL_CSTRING(CACHE_SUFFIX));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   nsCOMPtr<nsIOutputStream> out;
michael@0:   rv = NS_NewSafeLocalFileOutputStream(getter_AddRefs(out), storeFile,
michael@0:                                        PR_WRONLY | PR_TRUNCATE | PR_CREATE_FILE);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   UpdateHeader();
michael@0:   LOG(("Writing %d completions", mHeader.numCompletions));
michael@0: 
michael@0:   uint32_t written;
michael@0:   rv = out->Write(reinterpret_cast<char*>(&mHeader), sizeof(mHeader), &written);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = WriteTArray(out, mCompletions);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   nsCOMPtr<nsISafeOutputStream> safeOut = do_QueryInterface(out);
michael@0:   rv = safeOut->Finish();
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = EnsureSizeConsistent();
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   nsCOMPtr<nsIFile> psFile;
michael@0:   rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = mPrefixSet->StoreToFile(psFile);
michael@0:   NS_WARN_IF_FALSE(NS_SUCCEEDED(rv), "failed to store the prefixset");
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: void
michael@0: LookupCache::ClearAll()
michael@0: {
michael@0:   ClearCompleteCache();
michael@0:   mPrefixSet->SetPrefixes(nullptr, 0);
michael@0:   mPrimed = false;
michael@0: }
michael@0: 
michael@0: void
michael@0: LookupCache::ClearCompleteCache()
michael@0: {
michael@0:   mCompletions.Clear();
michael@0:   UpdateHeader();
michael@0: }
michael@0: 
michael@0: void
michael@0: LookupCache::UpdateHeader()
michael@0: {
michael@0:   mHeader.magic = LOOKUPCACHE_MAGIC;
michael@0:   mHeader.version = CURRENT_VERSION;
michael@0:   mHeader.numCompletions = mCompletions.Length();
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::EnsureSizeConsistent()
michael@0: {
michael@0:   nsCOMPtr<nsIFile> storeFile;
michael@0:   nsresult rv = mStoreDirectory->Clone(getter_AddRefs(storeFile));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   rv = storeFile->AppendNative(mTableName + NS_LITERAL_CSTRING(CACHE_SUFFIX));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   int64_t fileSize;
michael@0:   rv = storeFile->GetFileSize(&fileSize);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   if (fileSize < 0) {
michael@0:     return NS_ERROR_FAILURE;
michael@0:   }
michael@0: 
michael@0:   int64_t expectedSize = sizeof(mHeader)
michael@0:                         + mHeader.numCompletions*sizeof(Completion);
michael@0:   if (expectedSize != fileSize) {
michael@0:     NS_WARNING("File length does not match. Probably corrupted.");
michael@0:     Reset();
michael@0:     return NS_ERROR_FILE_CORRUPTED;
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::ReadHeader(nsIInputStream* aInputStream)
michael@0: {
michael@0:   if (!aInputStream) {
michael@0:     ClearCompleteCache();
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   nsCOMPtr<nsISeekableStream> seekable = do_QueryInterface(aInputStream);
michael@0:   nsresult rv = seekable->Seek(nsISeekableStream::NS_SEEK_SET, 0);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   void *buffer = &mHeader;
michael@0:   rv = NS_ReadInputStreamToBuffer(aInputStream,
michael@0:                                   &buffer,
michael@0:                                   sizeof(Header));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   if (mHeader.magic != LOOKUPCACHE_MAGIC || mHeader.version != CURRENT_VERSION) {
michael@0:     NS_WARNING("Unexpected header data in the store.");
michael@0:     Reset();
michael@0:     return NS_ERROR_FILE_CORRUPTED;
michael@0:   }
michael@0:   LOG(("%d completions present", mHeader.numCompletions));
michael@0: 
michael@0:   rv = EnsureSizeConsistent();
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::ReadCompletions(nsIInputStream* aInputStream)
michael@0: {
michael@0:   if (!mHeader.numCompletions) {
michael@0:     mCompletions.Clear();
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   nsCOMPtr<nsISeekableStream> seekable = do_QueryInterface(aInputStream);
michael@0:   nsresult rv = seekable->Seek(nsISeekableStream::NS_SEEK_SET, sizeof(Header));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = ReadTArray(aInputStream, &mCompletions, mHeader.numCompletions);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   LOG(("Read %d completions", mCompletions.Length()));
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: /* static */ bool
michael@0: LookupCache::IsCanonicalizedIP(const nsACString& aHost)
michael@0: {
michael@0:   // The canonicalization process will have left IP addresses in dotted
michael@0:   // decimal with no surprises.
michael@0:   uint32_t i1, i2, i3, i4;
michael@0:   char c;
michael@0:   if (PR_sscanf(PromiseFlatCString(aHost).get(), "%u.%u.%u.%u%c",
michael@0:                 &i1, &i2, &i3, &i4, &c) == 4) {
michael@0:     return (i1 <= 0xFF && i2 <= 0xFF && i3 <= 0xFF && i4 <= 0xFF);
michael@0:   }
michael@0: 
michael@0:   return false;
michael@0: }
michael@0: 
michael@0: /* static */ nsresult
michael@0: LookupCache::GetKey(const nsACString& aSpec,
michael@0:                     Completion* aHash,
michael@0:                     nsCOMPtr<nsICryptoHash>& aCryptoHash)
michael@0: {
michael@0:   nsACString::const_iterator begin, end, iter;
michael@0:   aSpec.BeginReading(begin);
michael@0:   aSpec.EndReading(end);
michael@0: 
michael@0:   iter = begin;
michael@0:   if (!FindCharInReadable('/', iter, end)) {
michael@0:    return NS_OK;
michael@0:   }
michael@0: 
michael@0:   const nsCSubstring& host = Substring(begin, iter);
michael@0: 
michael@0:   if (IsCanonicalizedIP(host)) {
michael@0:     nsAutoCString key;
michael@0:     key.Assign(host);
michael@0:     key.Append("/");
michael@0:     return aHash->FromPlaintext(key, aCryptoHash);
michael@0:   }
michael@0: 
michael@0:   nsTArray<nsCString> hostComponents;
michael@0:   ParseString(PromiseFlatCString(host), '.', hostComponents);
michael@0: 
michael@0:   if (hostComponents.Length() < 2)
michael@0:     return NS_ERROR_FAILURE;
michael@0: 
michael@0:   int32_t last = int32_t(hostComponents.Length()) - 1;
michael@0:   nsAutoCString lookupHost;
michael@0: 
michael@0:   if (hostComponents.Length() > 2) {
michael@0:     lookupHost.Append(hostComponents[last - 2]);
michael@0:     lookupHost.Append(".");
michael@0:   }
michael@0: 
michael@0:   lookupHost.Append(hostComponents[last - 1]);
michael@0:   lookupHost.Append(".");
michael@0:   lookupHost.Append(hostComponents[last]);
michael@0:   lookupHost.Append("/");
michael@0: 
michael@0:   return aHash->FromPlaintext(lookupHost, aCryptoHash);
michael@0: }
michael@0: 
michael@0: /* static */ nsresult
michael@0: LookupCache::GetLookupFragments(const nsACString& aSpec,
michael@0:                                 nsTArray<nsCString>* aFragments)
michael@0: 
michael@0: {
michael@0:   aFragments->Clear();
michael@0: 
michael@0:   nsACString::const_iterator begin, end, iter;
michael@0:   aSpec.BeginReading(begin);
michael@0:   aSpec.EndReading(end);
michael@0: 
michael@0:   iter = begin;
michael@0:   if (!FindCharInReadable('/', iter, end)) {
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   const nsCSubstring& host = Substring(begin, iter++);
michael@0:   nsAutoCString path;
michael@0:   path.Assign(Substring(iter, end));
michael@0: 
michael@0:   /**
michael@0:    * From the protocol doc:
michael@0:    * For the hostname, the client will try at most 5 different strings.  They
michael@0:    * are:
michael@0:    * a) The exact hostname of the url
michael@0:    * b) The 4 hostnames formed by starting with the last 5 components and
michael@0:    *    successivly removing the leading component.  The top-level component
michael@0:    *    can be skipped. This is not done if the hostname is a numerical IP.
michael@0:    */
michael@0:   nsTArray<nsCString> hosts;
michael@0:   hosts.AppendElement(host);
michael@0: 
michael@0:   if (!IsCanonicalizedIP(host)) {
michael@0:     host.BeginReading(begin);
michael@0:     host.EndReading(end);
michael@0:     int numHostComponents = 0;
michael@0:     while (RFindInReadable(NS_LITERAL_CSTRING("."), begin, end) &&
michael@0:            numHostComponents < MAX_HOST_COMPONENTS) {
michael@0:       // don't bother checking toplevel domains
michael@0:       if (++numHostComponents >= 2) {
michael@0:         host.EndReading(iter);
michael@0:         hosts.AppendElement(Substring(end, iter));
michael@0:       }
michael@0:       end = begin;
michael@0:       host.BeginReading(begin);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * From the protocol doc:
michael@0:    * For the path, the client will also try at most 6 different strings.
michael@0:    * They are:
michael@0:    * a) the exact path of the url, including query parameters
michael@0:    * b) the exact path of the url, without query parameters
michael@0:    * c) the 4 paths formed by starting at the root (/) and
michael@0:    *    successively appending path components, including a trailing
michael@0:    *    slash.  This behavior should only extend up to the next-to-last
michael@0:    *    path component, that is, a trailing slash should never be
michael@0:    *    appended that was not present in the original url.
michael@0:    */
michael@0:   nsTArray<nsCString> paths;
michael@0:   nsAutoCString pathToAdd;
michael@0: 
michael@0:   path.BeginReading(begin);
michael@0:   path.EndReading(end);
michael@0:   iter = begin;
michael@0:   if (FindCharInReadable('?', iter, end)) {
michael@0:     pathToAdd = Substring(begin, iter);
michael@0:     paths.AppendElement(pathToAdd);
michael@0:     end = iter;
michael@0:   }
michael@0: 
michael@0:   int numPathComponents = 1;
michael@0:   iter = begin;
michael@0:   while (FindCharInReadable('/', iter, end) &&
michael@0:          numPathComponents < MAX_PATH_COMPONENTS) {
michael@0:     iter++;
michael@0:     pathToAdd.Assign(Substring(begin, iter));
michael@0:     paths.AppendElement(pathToAdd);
michael@0:     numPathComponents++;
michael@0:   }
michael@0: 
michael@0:   // If we haven't already done so, add the full path
michael@0:   if (!pathToAdd.Equals(path)) {
michael@0:     paths.AppendElement(path);
michael@0:   }
michael@0:   // Check an empty path (for whole-domain blacklist entries)
michael@0:   paths.AppendElement(EmptyCString());
michael@0: 
michael@0:   for (uint32_t hostIndex = 0; hostIndex < hosts.Length(); hostIndex++) {
michael@0:     for (uint32_t pathIndex = 0; pathIndex < paths.Length(); pathIndex++) {
michael@0:       nsCString key;
michael@0:       key.Assign(hosts[hostIndex]);
michael@0:       key.Append('/');
michael@0:       key.Append(paths[pathIndex]);
michael@0:       LOG(("Checking fragment %s", key.get()));
michael@0: 
michael@0:       aFragments->AppendElement(key);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: /* static */ nsresult
michael@0: LookupCache::GetHostKeys(const nsACString& aSpec,
michael@0:                          nsTArray<nsCString>* aHostKeys)
michael@0: {
michael@0:   nsACString::const_iterator begin, end, iter;
michael@0:   aSpec.BeginReading(begin);
michael@0:   aSpec.EndReading(end);
michael@0: 
michael@0:   iter = begin;
michael@0:   if (!FindCharInReadable('/', iter, end)) {
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   const nsCSubstring& host = Substring(begin, iter);
michael@0: 
michael@0:   if (IsCanonicalizedIP(host)) {
michael@0:     nsCString *key = aHostKeys->AppendElement();
michael@0:     if (!key)
michael@0:       return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:     key->Assign(host);
michael@0:     key->Append("/");
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   nsTArray<nsCString> hostComponents;
michael@0:   ParseString(PromiseFlatCString(host), '.', hostComponents);
michael@0: 
michael@0:   if (hostComponents.Length() < 2) {
michael@0:     // no host or toplevel host, this won't match anything in the db
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   // First check with two domain components
michael@0:   int32_t last = int32_t(hostComponents.Length()) - 1;
michael@0:   nsCString *lookupHost = aHostKeys->AppendElement();
michael@0:   if (!lookupHost)
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:   lookupHost->Assign(hostComponents[last - 1]);
michael@0:   lookupHost->Append(".");
michael@0:   lookupHost->Append(hostComponents[last]);
michael@0:   lookupHost->Append("/");
michael@0: 
michael@0:   // Now check with three domain components
michael@0:   if (hostComponents.Length() > 2) {
michael@0:     nsCString *lookupHost2 = aHostKeys->AppendElement();
michael@0:     if (!lookupHost2)
michael@0:       return NS_ERROR_OUT_OF_MEMORY;
michael@0:     lookupHost2->Assign(hostComponents[last - 2]);
michael@0:     lookupHost2->Append(".");
michael@0:     lookupHost2->Append(*lookupHost);
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: bool LookupCache::IsPrimed()
michael@0: {
michael@0:   return mPrimed;
michael@0: }
michael@0: 
michael@0: #ifdef DEBUG
michael@0: template <class T>
michael@0: static void EnsureSorted(T* aArray)
michael@0: {
michael@0:   typename T::elem_type* start = aArray->Elements();
michael@0:   typename T::elem_type* end = aArray->Elements() + aArray->Length();
michael@0:   typename T::elem_type* iter = start;
michael@0:   typename T::elem_type* previous = start;
michael@0: 
michael@0:   while (iter != end) {
michael@0:     previous = iter;
michael@0:     ++iter;
michael@0:     if (iter != end) {
michael@0:       MOZ_ASSERT(*previous <= *iter);
michael@0:     }
michael@0:   }
michael@0:   return;
michael@0: }
michael@0: #endif
michael@0: 
michael@0: nsresult
michael@0: LookupCache::ConstructPrefixSet(AddPrefixArray& aAddPrefixes)
michael@0: {
michael@0:   Telemetry::AutoTimer<Telemetry::URLCLASSIFIER_PS_CONSTRUCT_TIME> timer;
michael@0: 
michael@0:   nsTArray<uint32_t> array;
michael@0:   array.SetCapacity(aAddPrefixes.Length());
michael@0: 
michael@0:   for (uint32_t i = 0; i < aAddPrefixes.Length(); i++) {
michael@0:     array.AppendElement(aAddPrefixes[i].PrefixHash().ToUint32());
michael@0:   }
michael@0:   aAddPrefixes.Clear();
michael@0: 
michael@0: #ifdef DEBUG
michael@0:   // PrefixSet requires sorted order
michael@0:   EnsureSorted(&array);
michael@0: #endif
michael@0: 
michael@0:   // construct new one, replace old entries
michael@0:   nsresult rv = mPrefixSet->SetPrefixes(array.Elements(), array.Length());
michael@0:   if (NS_FAILED(rv)) {
michael@0:     goto error_bailout;
michael@0:   }
michael@0: 
michael@0: #ifdef DEBUG
michael@0:   uint32_t size;
michael@0:   size = mPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
michael@0:   LOG(("SB tree done, size = %d bytes\n", size));
michael@0: #endif
michael@0: 
michael@0:   mPrimed = true;
michael@0: 
michael@0:   return NS_OK;
michael@0: 
michael@0:  error_bailout:
michael@0:   Telemetry::Accumulate(Telemetry::URLCLASSIFIER_PS_FAILURE, 1);
michael@0:   return rv;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::LoadPrefixSet()
michael@0: {
michael@0:   nsCOMPtr<nsIFile> psFile;
michael@0:   nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   bool exists;
michael@0:   rv = psFile->Exists(&exists);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   if (exists) {
michael@0:     LOG(("stored PrefixSet exists, loading from disk"));
michael@0:     rv = mPrefixSet->LoadFromFile(psFile);
michael@0:     if (NS_FAILED(rv)) {
michael@0:       if (rv == NS_ERROR_FILE_CORRUPTED) {
michael@0:         Reset();
michael@0:       }
michael@0:       return rv;
michael@0:     }
michael@0:     mPrimed = true;
michael@0:   } else {
michael@0:     LOG(("no (usable) stored PrefixSet found"));
michael@0:   }
michael@0: 
michael@0: #ifdef DEBUG
michael@0:   if (mPrimed) {
michael@0:     uint32_t size = mPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
michael@0:     LOG(("SB tree done, size = %d bytes\n", size));
michael@0:   }
michael@0: #endif
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: LookupCache::GetPrefixes(nsTArray<uint32_t>* aAddPrefixes)
michael@0: {
michael@0:   if (!mPrimed) {
michael@0:     // This can happen if its a new table, so no error.
michael@0:     LOG(("GetPrefixes from empty LookupCache"));
michael@0:     return NS_OK;
michael@0:   }
michael@0:   uint32_t cnt;
michael@0:   uint32_t *arr;
michael@0:   nsresult rv = mPrefixSet->GetPrefixes(&cnt, &arr);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   if (!aAddPrefixes->AppendElements(arr, cnt))
michael@0:     return NS_ERROR_FAILURE;
michael@0:   nsMemory::Free(arr);
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: 
michael@0: }
michael@0: }