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 "nsAutoPtr.h"
michael@0: #include "nsCOMPtr.h"
michael@0: #include "nsAppDirectoryServiceDefs.h"
michael@0: #include "nsCRT.h"
michael@0: #include "nsICryptoHash.h"
michael@0: #include "nsICryptoHMAC.h"
michael@0: #include "nsIDirectoryService.h"
michael@0: #include "nsIKeyModule.h"
michael@0: #include "nsIObserverService.h"
michael@0: #include "nsIPermissionManager.h"
michael@0: #include "nsIPrefBranch.h"
michael@0: #include "nsIPrefService.h"
michael@0: #include "nsIProperties.h"
michael@0: #include "nsToolkitCompsCID.h"
michael@0: #include "nsIUrlClassifierUtils.h"
michael@0: #include "nsUrlClassifierDBService.h"
michael@0: #include "nsUrlClassifierUtils.h"
michael@0: #include "nsUrlClassifierProxies.h"
michael@0: #include "nsURILoader.h"
michael@0: #include "nsString.h"
michael@0: #include "nsReadableUtils.h"
michael@0: #include "nsTArray.h"
michael@0: #include "nsNetUtil.h"
michael@0: #include "nsNetCID.h"
michael@0: #include "nsThreadUtils.h"
michael@0: #include "nsXPCOMStrings.h"
michael@0: #include "nsProxyRelease.h"
michael@0: #include "nsString.h"
michael@0: #include "mozilla/Atomics.h"
michael@0: #include "mozilla/DebugOnly.h"
michael@0: #include "mozilla/Mutex.h"
michael@0: #include "mozilla/Preferences.h"
michael@0: #include "mozilla/TimeStamp.h"
michael@0: #include "mozilla/Telemetry.h"
michael@0: #include "prlog.h"
michael@0: #include "prprf.h"
michael@0: #include "prnetdb.h"
michael@0: #include "Entries.h"
michael@0: #include "mozilla/Attributes.h"
michael@0: #include "nsIPrincipal.h"
michael@0: #include "Classifier.h"
michael@0: #include "ProtocolParser.h"
michael@0: #include "nsContentUtils.h"
michael@0: 
michael@0: using namespace mozilla;
michael@0: using namespace mozilla::safebrowsing;
michael@0: 
michael@0: // NSPR_LOG_MODULES=UrlClassifierDbService:5
michael@0: #if defined(PR_LOGGING)
michael@0: PRLogModuleInfo *gUrlClassifierDbServiceLog = nullptr;
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: // Prefs for implementing nsIURIClassifier to block page loads
michael@0: #define CHECK_MALWARE_PREF      "browser.safebrowsing.malware.enabled"
michael@0: #define CHECK_MALWARE_DEFAULT   false
michael@0: 
michael@0: #define CHECK_PHISHING_PREF     "browser.safebrowsing.enabled"
michael@0: #define CHECK_PHISHING_DEFAULT  false
michael@0: 
michael@0: #define GETHASH_NOISE_PREF      "urlclassifier.gethashnoise"
michael@0: #define GETHASH_NOISE_DEFAULT   4
michael@0: 
michael@0: // Comma-separated lists
michael@0: #define MALWARE_TABLE_PREF      "urlclassifier.malware_table"
michael@0: #define PHISH_TABLE_PREF        "urlclassifier.phish_table"
michael@0: #define DOWNLOAD_BLOCK_TABLE_PREF "urlclassifier.downloadBlockTable"
michael@0: #define DOWNLOAD_ALLOW_TABLE_PREF "urlclassifier.downloadAllowTable"
michael@0: #define DISALLOW_COMPLETION_TABLE_PREF "urlclassifier.disallow_completions"
michael@0: 
michael@0: #define CONFIRM_AGE_PREF        "urlclassifier.max-complete-age"
michael@0: #define CONFIRM_AGE_DEFAULT_SEC (45 * 60)
michael@0: 
michael@0: class nsUrlClassifierDBServiceWorker;
michael@0: 
michael@0: // Singleton instance.
michael@0: static nsUrlClassifierDBService* sUrlClassifierDBService;
michael@0: 
michael@0: nsIThread* nsUrlClassifierDBService::gDbBackgroundThread = nullptr;
michael@0: 
michael@0: // Once we've committed to shutting down, don't do work in the background
michael@0: // thread.
michael@0: static bool gShuttingDownThread = false;
michael@0: 
michael@0: static mozilla::Atomic<int32_t> gFreshnessGuarantee(CONFIRM_AGE_DEFAULT_SEC);
michael@0: 
michael@0: // -------------------------------------------------------------------------
michael@0: // Actual worker implemenatation
michael@0: class nsUrlClassifierDBServiceWorker MOZ_FINAL :
michael@0:   public nsIUrlClassifierDBServiceWorker
michael@0: {
michael@0: public:
michael@0:   nsUrlClassifierDBServiceWorker();
michael@0: 
michael@0:   NS_DECL_THREADSAFE_ISUPPORTS
michael@0:   NS_DECL_NSIURLCLASSIFIERDBSERVICE
michael@0:   NS_DECL_NSIURLCLASSIFIERDBSERVICEWORKER
michael@0: 
michael@0:   nsresult Init(uint32_t aGethashNoise, nsCOMPtr<nsIFile> aCacheDir);
michael@0: 
michael@0:   // Queue a lookup for the worker to perform, called in the main thread.
michael@0:   // tables is a comma-separated list of tables to query
michael@0:   nsresult QueueLookup(const nsACString& lookupKey,
michael@0:                        const nsACString& tables,
michael@0:                        nsIUrlClassifierLookupCallback* callback);
michael@0: 
michael@0:   // Handle any queued-up lookups.  We call this function during long-running
michael@0:   // update operations to prevent lookups from blocking for too long.
michael@0:   nsresult HandlePendingLookups();
michael@0: 
michael@0: private:
michael@0:   // No subclassing
michael@0:   ~nsUrlClassifierDBServiceWorker();
michael@0: 
michael@0:   // Disallow copy constructor
michael@0:   nsUrlClassifierDBServiceWorker(nsUrlClassifierDBServiceWorker&);
michael@0: 
michael@0:   nsresult OpenDb();
michael@0: 
michael@0:   // Applies the current transaction and resets the update/working times.
michael@0:   nsresult ApplyUpdate();
michael@0: 
michael@0:   // Reset the in-progress update stream
michael@0:   void ResetStream();
michael@0: 
michael@0:   // Reset the in-progress update
michael@0:   void ResetUpdate();
michael@0: 
michael@0:   // Perform a classifier lookup for a given url.
michael@0:   nsresult DoLookup(const nsACString& spec,
michael@0:                     const nsACString& tables,
michael@0:                     nsIUrlClassifierLookupCallback* c);
michael@0: 
michael@0:   nsresult AddNoise(const Prefix aPrefix,
michael@0:                     const nsCString tableName,
michael@0:                     uint32_t aCount,
michael@0:                     LookupResultArray& results);
michael@0: 
michael@0:   nsCOMPtr<nsICryptoHash> mCryptoHash;
michael@0: 
michael@0:   nsAutoPtr<Classifier> mClassifier;
michael@0:   // The class that actually parses the update chunks.
michael@0:   nsAutoPtr<ProtocolParser> mProtocolParser;
michael@0: 
michael@0:   // Directory where to store the SB databases.
michael@0:   nsCOMPtr<nsIFile> mCacheDir;
michael@0: 
michael@0:   // XXX: maybe an array of autoptrs.  Or maybe a class specifically
michael@0:   // storing a series of updates.
michael@0:   nsTArray<TableUpdate*> mTableUpdates;
michael@0: 
michael@0:   int32_t mUpdateWait;
michael@0: 
michael@0:   // Entries that cannot be completed. We expect them to die at
michael@0:   // the next update
michael@0:   PrefixArray mMissCache;
michael@0: 
michael@0:   nsresult mUpdateStatus;
michael@0:   nsTArray<nsCString> mUpdateTables;
michael@0: 
michael@0:   nsCOMPtr<nsIUrlClassifierUpdateObserver> mUpdateObserver;
michael@0:   bool mInStream;
michael@0: 
michael@0:   // The number of noise entries to add to the set of lookup results.
michael@0:   uint32_t mGethashNoise;
michael@0: 
michael@0:   // Pending lookups are stored in a queue for processing.  The queue
michael@0:   // is protected by mPendingLookupLock.
michael@0:   Mutex mPendingLookupLock;
michael@0: 
michael@0:   class PendingLookup {
michael@0:   public:
michael@0:     TimeStamp mStartTime;
michael@0:     nsCString mKey;
michael@0:     nsCString mTables;
michael@0:     nsCOMPtr<nsIUrlClassifierLookupCallback> mCallback;
michael@0:   };
michael@0: 
michael@0:   // list of pending lookups
michael@0:   nsTArray<PendingLookup> mPendingLookups;
michael@0: };
michael@0: 
michael@0: NS_IMPL_ISUPPORTS(nsUrlClassifierDBServiceWorker,
michael@0:                   nsIUrlClassifierDBServiceWorker,
michael@0:                   nsIUrlClassifierDBService)
michael@0: 
michael@0: nsUrlClassifierDBServiceWorker::nsUrlClassifierDBServiceWorker()
michael@0:   : mInStream(false)
michael@0:   , mGethashNoise(0)
michael@0:   , mPendingLookupLock("nsUrlClassifierDBServerWorker.mPendingLookupLock")
michael@0: {
michael@0: }
michael@0: 
michael@0: nsUrlClassifierDBServiceWorker::~nsUrlClassifierDBServiceWorker()
michael@0: {
michael@0:   NS_ASSERTION(!mClassifier,
michael@0:                "Db connection not closed, leaking memory!  Call CloseDb "
michael@0:                "to close the connection.");
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBServiceWorker::Init(uint32_t aGethashNoise,
michael@0:                                      nsCOMPtr<nsIFile> aCacheDir)
michael@0: {
michael@0:   mGethashNoise = aGethashNoise;
michael@0:   mCacheDir = aCacheDir;
michael@0: 
michael@0:   ResetUpdate();
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBServiceWorker::QueueLookup(const nsACString& spec,
michael@0:                                             const nsACString& tables,
michael@0:                                             nsIUrlClassifierLookupCallback* callback)
michael@0: {
michael@0:   MutexAutoLock lock(mPendingLookupLock);
michael@0: 
michael@0:   PendingLookup* lookup = mPendingLookups.AppendElement();
michael@0:   if (!lookup) return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:   lookup->mStartTime = TimeStamp::Now();
michael@0:   lookup->mKey = spec;
michael@0:   lookup->mCallback = callback;
michael@0:   lookup->mTables = tables;
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Lookup up a key in the database is a two step process:
michael@0:  *
michael@0:  * a) First we look for any Entries in the database that might apply to this
michael@0:  *    url.  For each URL there are one or two possible domain names to check:
michael@0:  *    the two-part domain name (example.com) and the three-part name
michael@0:  *    (www.example.com).  We check the database for both of these.
michael@0:  * b) If we find any entries, we check the list of fragments for that entry
michael@0:  *    against the possible subfragments of the URL as described in the
michael@0:  *    "Simplified Regular Expression Lookup" section of the protocol doc.
michael@0:  */
michael@0: nsresult
michael@0: nsUrlClassifierDBServiceWorker::DoLookup(const nsACString& spec,
michael@0:                                          const nsACString& tables,
michael@0:                                          nsIUrlClassifierLookupCallback* c)
michael@0: {
michael@0:   if (gShuttingDownThread) {
michael@0:     c->LookupComplete(nullptr);
michael@0:     return NS_ERROR_NOT_INITIALIZED;
michael@0:   }
michael@0: 
michael@0:   nsresult rv = OpenDb();
michael@0:   if (NS_FAILED(rv)) {
michael@0:     c->LookupComplete(nullptr);
michael@0:     NS_ERROR("Unable to open SafeBrowsing database.");
michael@0:     return NS_ERROR_FAILURE;
michael@0:   }
michael@0: 
michael@0: #if defined(PR_LOGGING)
michael@0:   PRIntervalTime clockStart = 0;
michael@0:   if (LOG_ENABLED()) {
michael@0:     clockStart = PR_IntervalNow();
michael@0:   }
michael@0: #endif
michael@0: 
michael@0:   nsAutoPtr<LookupResultArray> results(new LookupResultArray());
michael@0:   if (!results) {
michael@0:     c->LookupComplete(nullptr);
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0:   }
michael@0: 
michael@0:   // we ignore failures from Check because we'd rather return the
michael@0:   // results that were found than fail.
michael@0:   mClassifier->SetFreshTime(gFreshnessGuarantee);
michael@0:   mClassifier->Check(spec, tables, *results);
michael@0: 
michael@0:   LOG(("Found %d results.", results->Length()));
michael@0: 
michael@0: 
michael@0: #if defined(PR_LOGGING)
michael@0:   if (LOG_ENABLED()) {
michael@0:     PRIntervalTime clockEnd = PR_IntervalNow();
michael@0:     LOG(("query took %dms\n",
michael@0:          PR_IntervalToMilliseconds(clockEnd - clockStart)));
michael@0:   }
michael@0: #endif
michael@0: 
michael@0:   nsAutoPtr<LookupResultArray> completes(new LookupResultArray());
michael@0: 
michael@0:   for (uint32_t i = 0; i < results->Length(); i++) {
michael@0:     if (!mMissCache.Contains(results->ElementAt(i).hash.prefix)) {
michael@0:       completes->AppendElement(results->ElementAt(i));
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   for (uint32_t i = 0; i < completes->Length(); i++) {
michael@0:     if (!completes->ElementAt(i).Confirmed()) {
michael@0:       // We're going to be doing a gethash request, add some extra entries.
michael@0:       // Note that we cannot pass the first two by reference, because we
michael@0:       // add to completes, whicah can cause completes to reallocate and move.
michael@0:       AddNoise(completes->ElementAt(i).hash.prefix,
michael@0:                completes->ElementAt(i).mTableName,
michael@0:                mGethashNoise, *completes);
michael@0:       break;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   // At this point ownership of 'results' is handed to the callback.
michael@0:   c->LookupComplete(completes.forget());
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBServiceWorker::HandlePendingLookups()
michael@0: {
michael@0:   MutexAutoLock lock(mPendingLookupLock);
michael@0:   while (mPendingLookups.Length() > 0) {
michael@0:     PendingLookup lookup = mPendingLookups[0];
michael@0:     mPendingLookups.RemoveElementAt(0);
michael@0:     {
michael@0:       MutexAutoUnlock unlock(mPendingLookupLock);
michael@0:       DoLookup(lookup.mKey, lookup.mTables, lookup.mCallback);
michael@0:     }
michael@0:     double lookupTime = (TimeStamp::Now() - lookup.mStartTime).ToMilliseconds();
michael@0:     Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LOOKUP_TIME,
michael@0:                           static_cast<uint32_t>(lookupTime));
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBServiceWorker::AddNoise(const Prefix aPrefix,
michael@0:                                          const nsCString tableName,
michael@0:                                          uint32_t aCount,
michael@0:                                          LookupResultArray& results)
michael@0: {
michael@0:   if (aCount < 1) {
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   PrefixArray noiseEntries;
michael@0:   nsresult rv = mClassifier->ReadNoiseEntries(aPrefix, tableName,
michael@0:                                               aCount, &noiseEntries);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   for (uint32_t i = 0; i < noiseEntries.Length(); i++) {
michael@0:     LookupResult *result = results.AppendElement();
michael@0:     if (!result)
michael@0:       return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:     result->hash.prefix = noiseEntries[i];
michael@0:     result->mNoise = true;
michael@0: 
michael@0:     result->mTableName.Assign(tableName);
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: // Lookup a key in the db.
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::Lookup(nsIPrincipal* aPrincipal,
michael@0:                                        const nsACString& aTables,
michael@0:                                        nsIUrlClassifierCallback* c)
michael@0: {
michael@0:   return HandlePendingLookups();
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::GetTables(nsIUrlClassifierCallback* c)
michael@0: {
michael@0:   if (gShuttingDownThread)
michael@0:     return NS_ERROR_NOT_INITIALIZED;
michael@0: 
michael@0:   nsresult rv = OpenDb();
michael@0:   if (NS_FAILED(rv)) {
michael@0:     NS_ERROR("Unable to open SafeBrowsing database");
michael@0:     return NS_ERROR_FAILURE;
michael@0:   }
michael@0: 
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   nsAutoCString response;
michael@0:   mClassifier->TableRequest(response);
michael@0:   c->HandleEvent(response);
michael@0: 
michael@0:   return rv;
michael@0: }
michael@0: 
michael@0: void
michael@0: nsUrlClassifierDBServiceWorker::ResetStream()
michael@0: {
michael@0:   LOG(("ResetStream"));
michael@0:   mInStream = false;
michael@0:   mProtocolParser = nullptr;
michael@0: }
michael@0: 
michael@0: void
michael@0: nsUrlClassifierDBServiceWorker::ResetUpdate()
michael@0: {
michael@0:   LOG(("ResetUpdate"));
michael@0:   mUpdateWait = 0;
michael@0:   mUpdateStatus = NS_OK;
michael@0:   mUpdateObserver = nullptr;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::SetHashCompleter(const nsACString &tableName,
michael@0:                                                  nsIUrlClassifierHashCompleter *completer)
michael@0: {
michael@0:   return NS_ERROR_NOT_IMPLEMENTED;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::BeginUpdate(nsIUrlClassifierUpdateObserver *observer,
michael@0:                                             const nsACString &tables)
michael@0: {
michael@0:   LOG(("nsUrlClassifierDBServiceWorker::BeginUpdate [%s]", PromiseFlatCString(tables).get()));
michael@0: 
michael@0:   if (gShuttingDownThread)
michael@0:     return NS_ERROR_NOT_INITIALIZED;
michael@0: 
michael@0:   NS_ENSURE_STATE(!mUpdateObserver);
michael@0: 
michael@0:   nsresult rv = OpenDb();
michael@0:   if (NS_FAILED(rv)) {
michael@0:     NS_ERROR("Unable to open SafeBrowsing database");
michael@0:     return NS_ERROR_FAILURE;
michael@0:   }
michael@0: 
michael@0:   mUpdateStatus = NS_OK;
michael@0:   mUpdateObserver = observer;
michael@0:   Classifier::SplitTables(tables, mUpdateTables);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: // Called from the stream updater.
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::BeginStream(const nsACString &table)
michael@0: {
michael@0:   LOG(("nsUrlClassifierDBServiceWorker::BeginStream"));
michael@0: 
michael@0:   if (gShuttingDownThread)
michael@0:     return NS_ERROR_NOT_INITIALIZED;
michael@0: 
michael@0:   NS_ENSURE_STATE(mUpdateObserver);
michael@0:   NS_ENSURE_STATE(!mInStream);
michael@0: 
michael@0:   mInStream = true;
michael@0: 
michael@0:   NS_ASSERTION(!mProtocolParser, "Should not have a protocol parser.");
michael@0: 
michael@0:   mProtocolParser = new ProtocolParser();
michael@0:   if (!mProtocolParser)
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:   mProtocolParser->Init(mCryptoHash);
michael@0: 
michael@0:   if (!table.IsEmpty()) {
michael@0:     mProtocolParser->SetCurrentTable(table);
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Updating the database:
michael@0:  *
michael@0:  * The Update() method takes a series of chunks separated with control data,
michael@0:  * as described in
michael@0:  * http://code.google.com/p/google-safe-browsing/wiki/Protocolv2Spec
michael@0:  *
michael@0:  * It will iterate through the control data until it reaches a chunk.  By
michael@0:  * the time it reaches a chunk, it should have received
michael@0:  * a) the table to which this chunk applies
michael@0:  * b) the type of chunk (add, delete, expire add, expire delete).
michael@0:  * c) the chunk ID
michael@0:  * d) the length of the chunk.
michael@0:  *
michael@0:  * For add and subtract chunks, it needs to read the chunk data (expires
michael@0:  * don't have any data).  Chunk data is a list of URI fragments whose
michael@0:  * encoding depends on the type of table (which is indicated by the end
michael@0:  * of the table name):
michael@0:  * a) tables ending with -exp are a zlib-compressed list of URI fragments
michael@0:  *    separated by newlines.
michael@0:  * b) tables ending with -sha128 have the form
michael@0:  *    [domain][N][frag0]...[fragN]
michael@0:  *       16    1   16        16
michael@0:  *    If N is 0, the domain is reused as a fragment.
michael@0:  * c) any other tables are assumed to be a plaintext list of URI fragments
michael@0:  *    separated by newlines.
michael@0:  *
michael@0:  * Update() can be fed partial data;  It will accumulate data until there is
michael@0:  * enough to act on.  Finish() should be called when there will be no more
michael@0:  * data.
michael@0:  */
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::UpdateStream(const nsACString& chunk)
michael@0: {
michael@0:   if (gShuttingDownThread)
michael@0:     return NS_ERROR_NOT_INITIALIZED;
michael@0: 
michael@0:   NS_ENSURE_STATE(mInStream);
michael@0: 
michael@0:   HandlePendingLookups();
michael@0: 
michael@0:   // Feed the chunk to the parser.
michael@0:   return mProtocolParser->AppendStream(chunk);
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::FinishStream()
michael@0: {
michael@0:   if (gShuttingDownThread)
michael@0:     return NS_ERROR_NOT_INITIALIZED;
michael@0: 
michael@0:   NS_ENSURE_STATE(mInStream);
michael@0:   NS_ENSURE_STATE(mUpdateObserver);
michael@0: 
michael@0:   mInStream = false;
michael@0: 
michael@0:   if (NS_SUCCEEDED(mProtocolParser->Status())) {
michael@0:     if (mProtocolParser->UpdateWait()) {
michael@0:       mUpdateWait = mProtocolParser->UpdateWait();
michael@0:     }
michael@0:     // XXX: Only allow forwards from the initial update?
michael@0:     const nsTArray<ProtocolParser::ForwardedUpdate> &forwards =
michael@0:       mProtocolParser->Forwards();
michael@0:     for (uint32_t i = 0; i < forwards.Length(); i++) {
michael@0:       const ProtocolParser::ForwardedUpdate &forward = forwards[i];
michael@0:       mUpdateObserver->UpdateUrlRequested(forward.url, forward.table);
michael@0:     }
michael@0:     // Hold on to any TableUpdate objects that were created by the
michael@0:     // parser.
michael@0:     mTableUpdates.AppendElements(mProtocolParser->GetTableUpdates());
michael@0:     mProtocolParser->ForgetTableUpdates();
michael@0:   } else {
michael@0:     mUpdateStatus = mProtocolParser->Status();
michael@0:   }
michael@0:   mUpdateObserver->StreamFinished(mProtocolParser->Status(), 0);
michael@0: 
michael@0:   if (NS_SUCCEEDED(mUpdateStatus)) {
michael@0:     if (mProtocolParser->ResetRequested()) {
michael@0:       mClassifier->Reset();
michael@0:    }
michael@0:   }
michael@0: 
michael@0:   mProtocolParser = nullptr;
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::FinishUpdate()
michael@0: {
michael@0:   if (gShuttingDownThread)
michael@0:     return NS_ERROR_NOT_INITIALIZED;
michael@0:   NS_ENSURE_STATE(mUpdateObserver);
michael@0: 
michael@0:   if (NS_SUCCEEDED(mUpdateStatus)) {
michael@0:     mUpdateStatus = ApplyUpdate();
michael@0:   }
michael@0: 
michael@0:   mMissCache.Clear();
michael@0: 
michael@0:   if (NS_SUCCEEDED(mUpdateStatus)) {
michael@0:     LOG(("Notifying success: %d", mUpdateWait));
michael@0:     mUpdateObserver->UpdateSuccess(mUpdateWait);
michael@0:   } else {
michael@0:     LOG(("Notifying error: %d", mUpdateStatus));
michael@0:     mUpdateObserver->UpdateError(mUpdateStatus);
michael@0:     /*
michael@0:      * mark the tables as spoiled, we don't want to block hosts
michael@0:      * longer than normal because our update failed
michael@0:     */
michael@0:     mClassifier->MarkSpoiled(mUpdateTables);
michael@0:   }
michael@0:   mUpdateObserver = nullptr;
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBServiceWorker::ApplyUpdate()
michael@0: {
michael@0:   LOG(("nsUrlClassifierDBServiceWorker::ApplyUpdate()"));
michael@0:   return mClassifier->ApplyUpdates(&mTableUpdates);
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::ResetDatabase()
michael@0: {
michael@0:   nsresult rv = OpenDb();
michael@0: 
michael@0:   if (NS_SUCCEEDED(rv)) {
michael@0:     mClassifier->Reset();
michael@0:   }
michael@0: 
michael@0:   rv = CloseDb();
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::CancelUpdate()
michael@0: {
michael@0:   LOG(("nsUrlClassifierDBServiceWorker::CancelUpdate"));
michael@0: 
michael@0:   if (mUpdateObserver) {
michael@0:     LOG(("UpdateObserver exists, cancelling"));
michael@0: 
michael@0:     mUpdateStatus = NS_BINDING_ABORTED;
michael@0: 
michael@0:     mUpdateObserver->UpdateError(mUpdateStatus);
michael@0: 
michael@0:     /*
michael@0:      * mark the tables as spoiled, we don't want to block hosts
michael@0:      * longer than normal because our update failed
michael@0:     */
michael@0:     mClassifier->MarkSpoiled(mUpdateTables);
michael@0: 
michael@0:     ResetStream();
michael@0:     ResetUpdate();
michael@0:   } else {
michael@0:     LOG(("No UpdateObserver, nothing to cancel"));
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: // Allows the main thread to delete the connection which may be in
michael@0: // a background thread.
michael@0: // XXX This could be turned into a single shutdown event so the logic
michael@0: // is simpler in nsUrlClassifierDBService::Shutdown.
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::CloseDb()
michael@0: {
michael@0:   if (mClassifier) {
michael@0:     mClassifier->Close();
michael@0:     mClassifier = nullptr;
michael@0:   }
michael@0: 
michael@0:   mCryptoHash = nullptr;
michael@0:   LOG(("urlclassifier db closed\n"));
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::CacheCompletions(CacheResultArray *results)
michael@0: {
michael@0:   LOG(("nsUrlClassifierDBServiceWorker::CacheCompletions [%p]", this));
michael@0:   if (!mClassifier)
michael@0:     return NS_OK;
michael@0: 
michael@0:   // Ownership is transferred in to us
michael@0:   nsAutoPtr<CacheResultArray> resultsPtr(results);
michael@0: 
michael@0:   nsAutoPtr<ProtocolParser> pParse(new ProtocolParser());
michael@0:   nsTArray<TableUpdate*> updates;
michael@0: 
michael@0:   // Only cache results for tables that we have, don't take
michael@0:   // in tables we might accidentally have hit during a completion.
michael@0:   // This happens due to goog vs googpub lists existing.
michael@0:   nsTArray<nsCString> tables;
michael@0:   nsresult rv = mClassifier->ActiveTables(tables);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   for (uint32_t i = 0; i < resultsPtr->Length(); i++) {
michael@0:     bool activeTable = false;
michael@0:     for (uint32_t table = 0; table < tables.Length(); table++) {
michael@0:       if (tables[table].Equals(resultsPtr->ElementAt(i).table)) {
michael@0:         activeTable = true;
michael@0:         break;
michael@0:       }
michael@0:     }
michael@0:     if (activeTable) {
michael@0:       TableUpdate * tu = pParse->GetTableUpdate(resultsPtr->ElementAt(i).table);
michael@0:       LOG(("CacheCompletion Addchunk %d hash %X", resultsPtr->ElementAt(i).entry.addChunk,
michael@0:            resultsPtr->ElementAt(i).entry.ToUint32()));
michael@0:       tu->NewAddComplete(resultsPtr->ElementAt(i).entry.addChunk,
michael@0:                          resultsPtr->ElementAt(i).entry.complete);
michael@0:       tu->NewAddChunk(resultsPtr->ElementAt(i).entry.addChunk);
michael@0:       tu->SetLocalUpdate();
michael@0:       updates.AppendElement(tu);
michael@0:       pParse->ForgetTableUpdates();
michael@0:     } else {
michael@0:       LOG(("Completion received, but table is not active, so not caching."));
michael@0:     }
michael@0:    }
michael@0: 
michael@0:   mClassifier->ApplyUpdates(&updates);
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBServiceWorker::CacheMisses(PrefixArray *results)
michael@0: {
michael@0:   LOG(("nsUrlClassifierDBServiceWorker::CacheMisses [%p] %d",
michael@0:        this, results->Length()));
michael@0: 
michael@0:   // Ownership is transferred in to us
michael@0:   nsAutoPtr<PrefixArray> resultsPtr(results);
michael@0: 
michael@0:   for (uint32_t i = 0; i < resultsPtr->Length(); i++) {
michael@0:     mMissCache.AppendElement(resultsPtr->ElementAt(i));
michael@0:    }
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBServiceWorker::OpenDb()
michael@0: {
michael@0:   // Connection already open, don't do anything.
michael@0:   if (mClassifier) {
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   LOG(("Opening db"));
michael@0: 
michael@0:   nsresult rv;
michael@0:   mCryptoHash = do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID, &rv);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   nsAutoPtr<Classifier> classifier(new Classifier());
michael@0:   if (!classifier) {
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0:   }
michael@0: 
michael@0:   classifier->SetFreshTime(gFreshnessGuarantee);
michael@0: 
michael@0:   rv = classifier->Open(*mCacheDir);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   mClassifier = classifier;
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: // -------------------------------------------------------------------------
michael@0: // nsUrlClassifierLookupCallback
michael@0: //
michael@0: // This class takes the results of a lookup found on the worker thread
michael@0: // and handles any necessary partial hash expansions before calling
michael@0: // the client callback.
michael@0: 
michael@0: class nsUrlClassifierLookupCallback MOZ_FINAL : public nsIUrlClassifierLookupCallback
michael@0:                                               , public nsIUrlClassifierHashCompleterCallback
michael@0: {
michael@0: public:
michael@0:   NS_DECL_THREADSAFE_ISUPPORTS
michael@0:   NS_DECL_NSIURLCLASSIFIERLOOKUPCALLBACK
michael@0:   NS_DECL_NSIURLCLASSIFIERHASHCOMPLETERCALLBACK
michael@0: 
michael@0:   nsUrlClassifierLookupCallback(nsUrlClassifierDBService *dbservice,
michael@0:                                 nsIUrlClassifierCallback *c)
michael@0:     : mDBService(dbservice)
michael@0:     , mResults(nullptr)
michael@0:     , mPendingCompletions(0)
michael@0:     , mCallback(c)
michael@0:     {}
michael@0: 
michael@0:   ~nsUrlClassifierLookupCallback();
michael@0: 
michael@0: private:
michael@0:   nsresult HandleResults();
michael@0: 
michael@0:   nsRefPtr<nsUrlClassifierDBService> mDBService;
michael@0:   nsAutoPtr<LookupResultArray> mResults;
michael@0: 
michael@0:   // Completed results to send back to the worker for caching.
michael@0:   nsAutoPtr<CacheResultArray> mCacheResults;
michael@0: 
michael@0:   uint32_t mPendingCompletions;
michael@0:   nsCOMPtr<nsIUrlClassifierCallback> mCallback;
michael@0: };
michael@0: 
michael@0: NS_IMPL_ISUPPORTS(nsUrlClassifierLookupCallback,
michael@0:                   nsIUrlClassifierLookupCallback,
michael@0:                   nsIUrlClassifierHashCompleterCallback)
michael@0: 
michael@0: nsUrlClassifierLookupCallback::~nsUrlClassifierLookupCallback()
michael@0: {
michael@0:   nsCOMPtr<nsIThread> thread;
michael@0:   (void)NS_GetMainThread(getter_AddRefs(thread));
michael@0: 
michael@0:   if (mCallback) {
michael@0:     (void)NS_ProxyRelease(thread, mCallback, false);
michael@0:   }
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierLookupCallback::LookupComplete(nsTArray<LookupResult>* results)
michael@0: {
michael@0:   NS_ASSERTION(mResults == nullptr,
michael@0:                "Should only get one set of results per nsUrlClassifierLookupCallback!");
michael@0: 
michael@0:   if (!results) {
michael@0:     HandleResults();
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   mResults = results;
michael@0: 
michael@0:   // Check the results entries that need to be completed.
michael@0:   for (uint32_t i = 0; i < results->Length(); i++) {
michael@0:     LookupResult& result = results->ElementAt(i);
michael@0: 
michael@0:     // We will complete partial matches and matches that are stale.
michael@0:     if (!result.Confirmed()) {
michael@0:       nsCOMPtr<nsIUrlClassifierHashCompleter> completer;
michael@0:       if (mDBService->GetCompleter(result.mTableName,
michael@0:                                    getter_AddRefs(completer))) {
michael@0:         nsAutoCString partialHash;
michael@0:         partialHash.Assign(reinterpret_cast<char*>(&result.hash.prefix),
michael@0:                            PREFIX_SIZE);
michael@0: 
michael@0:         nsresult rv = completer->Complete(partialHash, this);
michael@0:         if (NS_SUCCEEDED(rv)) {
michael@0:           mPendingCompletions++;
michael@0:         }
michael@0:       } else {
michael@0:         // For tables with no hash completer, a complete hash match is
michael@0:         // good enough, we'll consider it fresh.
michael@0:         if (result.Complete()) {
michael@0:           result.mFresh = true;
michael@0:         } else {
michael@0:           NS_WARNING("Partial match in a table without a valid completer, ignoring partial match.");
michael@0:         }
michael@0:       }
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   if (mPendingCompletions == 0) {
michael@0:     // All results were complete, we're ready!
michael@0:     HandleResults();
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierLookupCallback::CompletionFinished(nsresult status)
michael@0: {
michael@0:   LOG(("nsUrlClassifierLookupCallback::CompletionFinished [%p, %08x]",
michael@0:        this, status));
michael@0:   if (NS_FAILED(status)) {
michael@0:     NS_WARNING("gethash response failed.");
michael@0:   }
michael@0: 
michael@0:   mPendingCompletions--;
michael@0:   if (mPendingCompletions == 0) {
michael@0:     HandleResults();
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierLookupCallback::Completion(const nsACString& completeHash,
michael@0:                                           const nsACString& tableName,
michael@0:                                           uint32_t chunkId)
michael@0: {
michael@0:   LOG(("nsUrlClassifierLookupCallback::Completion [%p, %s, %d]",
michael@0:        this, PromiseFlatCString(tableName).get(), chunkId));
michael@0:   mozilla::safebrowsing::Completion hash;
michael@0:   hash.Assign(completeHash);
michael@0: 
michael@0:   // Send this completion to the store for caching.
michael@0:   if (!mCacheResults) {
michael@0:     mCacheResults = new CacheResultArray();
michael@0:     if (!mCacheResults)
michael@0:       return NS_ERROR_OUT_OF_MEMORY;
michael@0:   }
michael@0: 
michael@0:   CacheResult result;
michael@0:   result.entry.addChunk = chunkId;
michael@0:   result.entry.complete = hash;
michael@0:   result.table = tableName;
michael@0: 
michael@0:   // OK if this fails, we just won't cache the item.
michael@0:   mCacheResults->AppendElement(result);
michael@0: 
michael@0:   // Check if this matched any of our results.
michael@0:   for (uint32_t i = 0; i < mResults->Length(); i++) {
michael@0:     LookupResult& result = mResults->ElementAt(i);
michael@0: 
michael@0:     // Now, see if it verifies a lookup
michael@0:     if (result.CompleteHash() == hash && result.mTableName.Equals(tableName)) {
michael@0:       result.mProtocolConfirmed = true;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierLookupCallback::HandleResults()
michael@0: {
michael@0:   if (!mResults) {
michael@0:     // No results, this URI is clean.
michael@0:     return mCallback->HandleEvent(NS_LITERAL_CSTRING(""));
michael@0:   }
michael@0: 
michael@0:   nsTArray<nsCString> tables;
michael@0:   // Build a stringified list of result tables.
michael@0:   for (uint32_t i = 0; i < mResults->Length(); i++) {
michael@0:     LookupResult& result = mResults->ElementAt(i);
michael@0: 
michael@0:     // Leave out results that weren't confirmed, as their existence on
michael@0:     // the list can't be verified.  Also leave out randomly-generated
michael@0:     // noise.
michael@0:     if (!result.Confirmed() || result.mNoise) {
michael@0:       LOG(("Skipping result from table %s", result.mTableName.get()));
michael@0:       continue;
michael@0:     }
michael@0: 
michael@0:     LOG(("Confirmed result from table %s", result.mTableName.get()));
michael@0: 
michael@0:     if (tables.IndexOf(result.mTableName) == nsTArray<nsCString>::NoIndex) {
michael@0:       tables.AppendElement(result.mTableName);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   // Some parts of this gethash request generated no hits at all.
michael@0:   // Prefixes must have been removed from the database since our last update.
michael@0:   // Save the prefixes we checked to prevent repeated requests
michael@0:   // until the next update.
michael@0:   nsAutoPtr<PrefixArray> cacheMisses(new PrefixArray());
michael@0:   if (cacheMisses) {
michael@0:     for (uint32_t i = 0; i < mResults->Length(); i++) {
michael@0:       LookupResult &result = mResults->ElementAt(i);
michael@0:       if (!result.Confirmed() && !result.mNoise) {
michael@0:         cacheMisses->AppendElement(result.PrefixHash());
michael@0:       }
michael@0:     }
michael@0:     // Hands ownership of the miss array back to the worker thread.
michael@0:     mDBService->CacheMisses(cacheMisses.forget());
michael@0:   }
michael@0: 
michael@0:   if (mCacheResults) {
michael@0:     // This hands ownership of the cache results array back to the worker
michael@0:     // thread.
michael@0:     mDBService->CacheCompletions(mCacheResults.forget());
michael@0:   }
michael@0: 
michael@0:   nsAutoCString tableStr;
michael@0:   for (uint32_t i = 0; i < tables.Length(); i++) {
michael@0:     if (i != 0)
michael@0:       tableStr.Append(',');
michael@0:     tableStr.Append(tables[i]);
michael@0:   }
michael@0: 
michael@0:   return mCallback->HandleEvent(tableStr);
michael@0: }
michael@0: 
michael@0: 
michael@0: // -------------------------------------------------------------------------
michael@0: // Helper class for nsIURIClassifier implementation, translates table names
michael@0: // to nsIURIClassifier enums.
michael@0: 
michael@0: class nsUrlClassifierClassifyCallback MOZ_FINAL : public nsIUrlClassifierCallback
michael@0: {
michael@0: public:
michael@0:   NS_DECL_THREADSAFE_ISUPPORTS
michael@0:   NS_DECL_NSIURLCLASSIFIERCALLBACK
michael@0: 
michael@0:   nsUrlClassifierClassifyCallback(nsIURIClassifierCallback *c,
michael@0:                                   bool checkMalware,
michael@0:                                   bool checkPhishing)
michael@0:     : mCallback(c)
michael@0:     , mCheckMalware(checkMalware)
michael@0:     , mCheckPhishing(checkPhishing)
michael@0:     {}
michael@0: 
michael@0: private:
michael@0:   nsCOMPtr<nsIURIClassifierCallback> mCallback;
michael@0:   bool mCheckMalware;
michael@0:   bool mCheckPhishing;
michael@0: };
michael@0: 
michael@0: NS_IMPL_ISUPPORTS(nsUrlClassifierClassifyCallback,
michael@0:                   nsIUrlClassifierCallback)
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierClassifyCallback::HandleEvent(const nsACString& tables)
michael@0: {
michael@0:   // XXX: we should probably have the wardens tell the service which table
michael@0:   // names match with which classification.  For now the table names give
michael@0:   // enough information.
michael@0:   nsresult response = NS_OK;
michael@0: 
michael@0:   nsACString::const_iterator begin, end;
michael@0: 
michael@0:   tables.BeginReading(begin);
michael@0:   tables.EndReading(end);
michael@0:   if (mCheckMalware &&
michael@0:       FindInReadable(NS_LITERAL_CSTRING("-malware-"), begin, end)) {
michael@0:     response = NS_ERROR_MALWARE_URI;
michael@0:   } else {
michael@0:     // Reset begin before checking phishing table
michael@0:     tables.BeginReading(begin);
michael@0: 
michael@0:     if (mCheckPhishing &&
michael@0:         FindInReadable(NS_LITERAL_CSTRING("-phish-"), begin, end)) {
michael@0:       response = NS_ERROR_PHISHING_URI;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   mCallback->OnClassifyComplete(response);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: 
michael@0: // -------------------------------------------------------------------------
michael@0: // Proxy class implementation
michael@0: 
michael@0: NS_IMPL_ISUPPORTS(nsUrlClassifierDBService,
michael@0:                   nsIUrlClassifierDBService,
michael@0:                   nsIURIClassifier,
michael@0:                   nsIObserver)
michael@0: 
michael@0: /* static */ nsUrlClassifierDBService*
michael@0: nsUrlClassifierDBService::GetInstance(nsresult *result)
michael@0: {
michael@0:   *result = NS_OK;
michael@0:   if (!sUrlClassifierDBService) {
michael@0:     sUrlClassifierDBService = new nsUrlClassifierDBService();
michael@0:     if (!sUrlClassifierDBService) {
michael@0:       *result = NS_ERROR_OUT_OF_MEMORY;
michael@0:       return nullptr;
michael@0:     }
michael@0: 
michael@0:     NS_ADDREF(sUrlClassifierDBService);   // addref the global
michael@0: 
michael@0:     *result = sUrlClassifierDBService->Init();
michael@0:     if (NS_FAILED(*result)) {
michael@0:       NS_RELEASE(sUrlClassifierDBService);
michael@0:       return nullptr;
michael@0:     }
michael@0:   } else {
michael@0:     // Already exists, just add a ref
michael@0:     NS_ADDREF(sUrlClassifierDBService);   // addref the return result
michael@0:   }
michael@0:   return sUrlClassifierDBService;
michael@0: }
michael@0: 
michael@0: 
michael@0: nsUrlClassifierDBService::nsUrlClassifierDBService()
michael@0:  : mCheckMalware(CHECK_MALWARE_DEFAULT)
michael@0:  , mCheckPhishing(CHECK_PHISHING_DEFAULT)
michael@0:  , mInUpdate(false)
michael@0: {
michael@0: }
michael@0: 
michael@0: nsUrlClassifierDBService::~nsUrlClassifierDBService()
michael@0: {
michael@0:   sUrlClassifierDBService = nullptr;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBService::ReadTablesFromPrefs()
michael@0: {
michael@0:   nsCString allTables;
michael@0:   nsCString tables;
michael@0:   Preferences::GetCString(PHISH_TABLE_PREF, &allTables);
michael@0: 
michael@0:   Preferences::GetCString(MALWARE_TABLE_PREF, &tables);
michael@0:   if (!tables.IsEmpty()) {
michael@0:     allTables.Append(',');
michael@0:     allTables.Append(tables);
michael@0:   }
michael@0: 
michael@0:   Preferences::GetCString(DOWNLOAD_BLOCK_TABLE_PREF, &tables);
michael@0:   if (!tables.IsEmpty()) {
michael@0:     allTables.Append(',');
michael@0:     allTables.Append(tables);
michael@0:   }
michael@0: 
michael@0:   Preferences::GetCString(DOWNLOAD_ALLOW_TABLE_PREF, &tables);
michael@0:   if (!tables.IsEmpty()) {
michael@0:     allTables.Append(',');
michael@0:     allTables.Append(tables);
michael@0:   }
michael@0: 
michael@0:   Classifier::SplitTables(allTables, mGethashTables);
michael@0: 
michael@0:   Preferences::GetCString(DISALLOW_COMPLETION_TABLE_PREF, &tables);
michael@0:   Classifier::SplitTables(tables, mDisallowCompletionsTables);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBService::Init()
michael@0: {
michael@0: #if defined(PR_LOGGING)
michael@0:   if (!gUrlClassifierDbServiceLog)
michael@0:     gUrlClassifierDbServiceLog = PR_NewLogModule("UrlClassifierDbService");
michael@0: #endif
michael@0: 
michael@0:   // Retrieve all the preferences.
michael@0:   mCheckMalware = Preferences::GetBool(CHECK_MALWARE_PREF,
michael@0:     CHECK_MALWARE_DEFAULT);
michael@0:   mCheckPhishing = Preferences::GetBool(CHECK_PHISHING_PREF,
michael@0:     CHECK_PHISHING_DEFAULT);
michael@0:   uint32_t gethashNoise = Preferences::GetUint(GETHASH_NOISE_PREF,
michael@0:     GETHASH_NOISE_DEFAULT);
michael@0:   gFreshnessGuarantee = Preferences::GetInt(CONFIRM_AGE_PREF,
michael@0:     CONFIRM_AGE_DEFAULT_SEC);
michael@0:   ReadTablesFromPrefs();
michael@0: 
michael@0:   // Do we *really* need to be able to change all of these at runtime?
michael@0:   Preferences::AddStrongObserver(this, CHECK_MALWARE_PREF);
michael@0:   Preferences::AddStrongObserver(this, CHECK_PHISHING_PREF);
michael@0:   Preferences::AddStrongObserver(this, GETHASH_NOISE_PREF);
michael@0:   Preferences::AddStrongObserver(this, CONFIRM_AGE_PREF);
michael@0:   Preferences::AddStrongObserver(this, PHISH_TABLE_PREF);
michael@0:   Preferences::AddStrongObserver(this, MALWARE_TABLE_PREF);
michael@0:   Preferences::AddStrongObserver(this, DOWNLOAD_BLOCK_TABLE_PREF);
michael@0:   Preferences::AddStrongObserver(this, DOWNLOAD_ALLOW_TABLE_PREF);
michael@0:   Preferences::AddStrongObserver(this, DISALLOW_COMPLETION_TABLE_PREF);
michael@0: 
michael@0:   // Force PSM loading on main thread
michael@0:   nsresult rv;
michael@0:   nsCOMPtr<nsICryptoHash> acryptoHash = do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID, &rv);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   // Directory providers must also be accessed on the main thread.
michael@0:   nsCOMPtr<nsIFile> cacheDir;
michael@0:   rv = NS_GetSpecialDirectory(NS_APP_USER_PROFILE_LOCAL_50_DIR,
michael@0:                               getter_AddRefs(cacheDir));
michael@0:   if (NS_FAILED(rv)) {
michael@0:     rv = NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
michael@0:                                 getter_AddRefs(cacheDir));
michael@0:   }
michael@0: 
michael@0:   // Start the background thread.
michael@0:   rv = NS_NewNamedThread("URL Classifier", &gDbBackgroundThread);
michael@0:   if (NS_FAILED(rv))
michael@0:     return rv;
michael@0: 
michael@0:   mWorker = new nsUrlClassifierDBServiceWorker();
michael@0:   if (!mWorker)
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:   rv = mWorker->Init(gethashNoise, cacheDir);
michael@0:   if (NS_FAILED(rv)) {
michael@0:     mWorker = nullptr;
michael@0:     return rv;
michael@0:   }
michael@0: 
michael@0:   // Proxy for calling the worker on the background thread
michael@0:   mWorkerProxy = new UrlClassifierDBServiceWorkerProxy(mWorker);
michael@0: 
michael@0:   // Add an observer for shutdown
michael@0:   nsCOMPtr<nsIObserverService> observerService =
michael@0:       mozilla::services::GetObserverService();
michael@0:   if (!observerService)
michael@0:     return NS_ERROR_FAILURE;
michael@0: 
michael@0:   observerService->AddObserver(this, "profile-before-change", false);
michael@0:   observerService->AddObserver(this, "xpcom-shutdown-threads", false);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: // nsChannelClassifier is the only consumer of this interface.
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::Classify(nsIPrincipal* aPrincipal,
michael@0:                                    nsIURIClassifierCallback* c,
michael@0:                                    bool* result)
michael@0: {
michael@0:   NS_ENSURE_ARG(aPrincipal);
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   if (!(mCheckMalware || mCheckPhishing)) {
michael@0:     *result = false;
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   nsRefPtr<nsUrlClassifierClassifyCallback> callback =
michael@0:     new nsUrlClassifierClassifyCallback(c, mCheckMalware, mCheckPhishing);
michael@0:   if (!callback) return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:   nsAutoCString tables;
michael@0:   nsAutoCString malware;
michael@0:   // LookupURI takes a comma-separated list already.
michael@0:   Preferences::GetCString(MALWARE_TABLE_PREF, &malware);
michael@0:   if (!malware.IsEmpty()) {
michael@0:     tables.Append(malware);
michael@0:   }
michael@0:   nsAutoCString phishing;
michael@0:   Preferences::GetCString(PHISH_TABLE_PREF, &phishing);
michael@0:   if (!phishing.IsEmpty()) {
michael@0:     tables.Append(",");
michael@0:     tables.Append(phishing);
michael@0:   }
michael@0:   nsresult rv = LookupURI(aPrincipal, tables, callback, false, result);
michael@0:   if (rv == NS_ERROR_MALFORMED_URI) {
michael@0:     *result = false;
michael@0:     // The URI had no hostname, don't try to classify it.
michael@0:     return NS_OK;
michael@0:   }
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::Lookup(nsIPrincipal* aPrincipal,
michael@0:                                  const nsACString& tables,
michael@0:                                  nsIUrlClassifierCallback* c)
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   bool dummy;
michael@0:   return LookupURI(aPrincipal, tables, c, true, &dummy);
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBService::LookupURI(nsIPrincipal* aPrincipal,
michael@0:                                     const nsACString& tables,
michael@0:                                     nsIUrlClassifierCallback* c,
michael@0:                                     bool forceLookup,
michael@0:                                     bool *didLookup)
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0:   NS_ENSURE_ARG(aPrincipal);
michael@0: 
michael@0:   if (nsContentUtils::IsSystemPrincipal(aPrincipal)) {
michael@0:     *didLookup = false;
michael@0:     return NS_OK;
michael@0:   }
michael@0: 
michael@0:   nsCOMPtr<nsIURI> uri;
michael@0:   nsresult rv = aPrincipal->GetURI(getter_AddRefs(uri));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   NS_ENSURE_TRUE(uri, NS_ERROR_FAILURE);
michael@0: 
michael@0:   uri = NS_GetInnermostURI(uri);
michael@0:   NS_ENSURE_TRUE(uri, NS_ERROR_FAILURE);
michael@0: 
michael@0:   nsAutoCString key;
michael@0:   // Canonicalize the url
michael@0:   nsCOMPtr<nsIUrlClassifierUtils> utilsService =
michael@0:     do_GetService(NS_URLCLASSIFIERUTILS_CONTRACTID);
michael@0:   rv = utilsService->GetKeyForURI(uri, key);
michael@0:   if (NS_FAILED(rv))
michael@0:     return rv;
michael@0: 
michael@0:   if (forceLookup) {
michael@0:     *didLookup = true;
michael@0:   } else {
michael@0:     bool clean = false;
michael@0: 
michael@0:     if (!clean) {
michael@0:       nsCOMPtr<nsIPermissionManager> permissionManager =
michael@0:         do_GetService(NS_PERMISSIONMANAGER_CONTRACTID);
michael@0: 
michael@0:       if (permissionManager) {
michael@0:         uint32_t perm;
michael@0:         rv = permissionManager->TestPermissionFromPrincipal(aPrincipal,
michael@0:                                                            "safe-browsing", &perm);
michael@0:         NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:         clean |= (perm == nsIPermissionManager::ALLOW_ACTION);
michael@0:       }
michael@0:     }
michael@0: 
michael@0:     *didLookup = !clean;
michael@0:     if (clean) {
michael@0:       return NS_OK;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   // Create an nsUrlClassifierLookupCallback object.  This object will
michael@0:   // take care of confirming partial hash matches if necessary before
michael@0:   // calling the client's callback.
michael@0:   nsCOMPtr<nsIUrlClassifierLookupCallback> callback =
michael@0:     new nsUrlClassifierLookupCallback(this, c);
michael@0:   if (!callback)
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:   nsCOMPtr<nsIUrlClassifierLookupCallback> proxyCallback =
michael@0:     new UrlClassifierLookupCallbackProxy(callback);
michael@0: 
michael@0:   // Queue this lookup and call the lookup function to flush the queue if
michael@0:   // necessary.
michael@0:   rv = mWorker->QueueLookup(key, tables, proxyCallback);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   // This seems to just call HandlePendingLookups.
michael@0:   nsAutoCString dummy;
michael@0:   return mWorkerProxy->Lookup(nullptr, dummy, nullptr);
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::GetTables(nsIUrlClassifierCallback* c)
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   // The proxy callback uses the current thread.
michael@0:   nsCOMPtr<nsIUrlClassifierCallback> proxyCallback =
michael@0:     new UrlClassifierCallbackProxy(c);
michael@0: 
michael@0:   return mWorkerProxy->GetTables(proxyCallback);
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::SetHashCompleter(const nsACString &tableName,
michael@0:                                            nsIUrlClassifierHashCompleter *completer)
michael@0: {
michael@0:   if (completer) {
michael@0:     mCompleters.Put(tableName, completer);
michael@0:   } else {
michael@0:     mCompleters.Remove(tableName);
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::BeginUpdate(nsIUrlClassifierUpdateObserver *observer,
michael@0:                                       const nsACString &updateTables)
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   if (mInUpdate)
michael@0:     return NS_ERROR_NOT_AVAILABLE;
michael@0: 
michael@0:   mInUpdate = true;
michael@0: 
michael@0:   // The proxy observer uses the current thread
michael@0:   nsCOMPtr<nsIUrlClassifierUpdateObserver> proxyObserver =
michael@0:     new UrlClassifierUpdateObserverProxy(observer);
michael@0: 
michael@0:   return mWorkerProxy->BeginUpdate(proxyObserver, updateTables);
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::BeginStream(const nsACString &table)
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   return mWorkerProxy->BeginStream(table);
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::UpdateStream(const nsACString& aUpdateChunk)
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   return mWorkerProxy->UpdateStream(aUpdateChunk);
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::FinishStream()
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   return mWorkerProxy->FinishStream();
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::FinishUpdate()
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   mInUpdate = false;
michael@0: 
michael@0:   return mWorkerProxy->FinishUpdate();
michael@0: }
michael@0: 
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::CancelUpdate()
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   mInUpdate = false;
michael@0: 
michael@0:   return mWorkerProxy->CancelUpdate();
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::ResetDatabase()
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   return mWorkerProxy->ResetDatabase();
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBService::CacheCompletions(CacheResultArray *results)
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   return mWorkerProxy->CacheCompletions(results);
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsUrlClassifierDBService::CacheMisses(PrefixArray *results)
michael@0: {
michael@0:   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
michael@0: 
michael@0:   return mWorkerProxy->CacheMisses(results);
michael@0: }
michael@0: 
michael@0: bool
michael@0: nsUrlClassifierDBService::GetCompleter(const nsACString &tableName,
michael@0:                                        nsIUrlClassifierHashCompleter **completer)
michael@0: {
michael@0:   // If we have specified a completer, go ahead and query it. This is only
michael@0:   // used by tests.
michael@0:   if (mCompleters.Get(tableName, completer)) {
michael@0:     return true;
michael@0:   }
michael@0: 
michael@0:   // If we don't know about this table at all, or are disallowing completions
michael@0:   // for it, skip completion checks.
michael@0:   if (!mGethashTables.Contains(tableName) ||
michael@0:       mDisallowCompletionsTables.Contains(tableName)) {
michael@0:     return false;
michael@0:   }
michael@0: 
michael@0:   MOZ_ASSERT(!StringBeginsWith(tableName, NS_LITERAL_CSTRING("test-")),
michael@0:              "We should never fetch hash completions for test tables");
michael@0: 
michael@0:   // Otherwise, call gethash to find the hash completions.
michael@0:   return NS_SUCCEEDED(CallGetService(NS_URLCLASSIFIERHASHCOMPLETER_CONTRACTID,
michael@0:                                      completer));
michael@0: }
michael@0: 
michael@0: NS_IMETHODIMP
michael@0: nsUrlClassifierDBService::Observe(nsISupports *aSubject, const char *aTopic,
michael@0:                                   const char16_t *aData)
michael@0: {
michael@0:   if (!strcmp(aTopic, NS_PREFBRANCH_PREFCHANGE_TOPIC_ID)) {
michael@0:     nsresult rv;
michael@0:     nsCOMPtr<nsIPrefBranch> prefs(do_QueryInterface(aSubject, &rv));
michael@0:     NS_ENSURE_SUCCESS(rv, rv);
michael@0:     if (NS_LITERAL_STRING(CHECK_MALWARE_PREF).Equals(aData)) {
michael@0:       mCheckMalware = Preferences::GetBool(CHECK_MALWARE_PREF,
michael@0:         CHECK_MALWARE_DEFAULT);
michael@0:     } else if (NS_LITERAL_STRING(CHECK_PHISHING_PREF).Equals(aData)) {
michael@0:       mCheckPhishing = Preferences::GetBool(CHECK_PHISHING_PREF,
michael@0:         CHECK_PHISHING_DEFAULT);
michael@0:     } else if (
michael@0:       NS_LITERAL_STRING(PHISH_TABLE_PREF).Equals(aData) ||
michael@0:       NS_LITERAL_STRING(MALWARE_TABLE_PREF).Equals(aData) ||
michael@0:       NS_LITERAL_STRING(DOWNLOAD_BLOCK_TABLE_PREF).Equals(aData) ||
michael@0:       NS_LITERAL_STRING(DOWNLOAD_ALLOW_TABLE_PREF).Equals(aData) ||
michael@0:       NS_LITERAL_STRING(DISALLOW_COMPLETION_TABLE_PREF).Equals(aData)) {
michael@0:       // Just read everything again.
michael@0:       ReadTablesFromPrefs();
michael@0:     } else if (NS_LITERAL_STRING(CONFIRM_AGE_PREF).Equals(aData)) {
michael@0:       gFreshnessGuarantee = Preferences::GetInt(CONFIRM_AGE_PREF,
michael@0:         CONFIRM_AGE_DEFAULT_SEC);
michael@0:     }
michael@0:   } else if (!strcmp(aTopic, "profile-before-change") ||
michael@0:              !strcmp(aTopic, "xpcom-shutdown-threads")) {
michael@0:     Shutdown();
michael@0:   } else {
michael@0:     return NS_ERROR_UNEXPECTED;
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: // Join the background thread if it exists.
michael@0: nsresult
michael@0: nsUrlClassifierDBService::Shutdown()
michael@0: {
michael@0:   LOG(("shutting down db service\n"));
michael@0: 
michael@0:   if (!gDbBackgroundThread)
michael@0:     return NS_OK;
michael@0: 
michael@0:   mCompleters.Clear();
michael@0: 
michael@0:   nsCOMPtr<nsIPrefBranch> prefs = do_GetService(NS_PREFSERVICE_CONTRACTID);
michael@0:   if (prefs) {
michael@0:     prefs->RemoveObserver(CHECK_MALWARE_PREF, this);
michael@0:     prefs->RemoveObserver(CHECK_PHISHING_PREF, this);
michael@0:     prefs->RemoveObserver(PHISH_TABLE_PREF, this);
michael@0:     prefs->RemoveObserver(MALWARE_TABLE_PREF, this);
michael@0:     prefs->RemoveObserver(DOWNLOAD_BLOCK_TABLE_PREF, this);
michael@0:     prefs->RemoveObserver(DOWNLOAD_ALLOW_TABLE_PREF, this);
michael@0:     prefs->RemoveObserver(DISALLOW_COMPLETION_TABLE_PREF, this);
michael@0:     prefs->RemoveObserver(CONFIRM_AGE_PREF, this);
michael@0:   }
michael@0: 
michael@0:   DebugOnly<nsresult> rv;
michael@0:   // First close the db connection.
michael@0:   if (mWorker) {
michael@0:     rv = mWorkerProxy->CancelUpdate();
michael@0:     NS_ASSERTION(NS_SUCCEEDED(rv), "failed to post cancel update event");
michael@0: 
michael@0:     rv = mWorkerProxy->CloseDb();
michael@0:     NS_ASSERTION(NS_SUCCEEDED(rv), "failed to post close db event");
michael@0:   }
michael@0: 
michael@0:   mWorkerProxy = nullptr;
michael@0: 
michael@0:   LOG(("joining background thread"));
michael@0: 
michael@0:   gShuttingDownThread = true;
michael@0: 
michael@0:   nsIThread *backgroundThread = gDbBackgroundThread;
michael@0:   gDbBackgroundThread = nullptr;
michael@0:   backgroundThread->Shutdown();
michael@0:   NS_RELEASE(backgroundThread);
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsIThread*
michael@0: nsUrlClassifierDBService::BackgroundThread()
michael@0: {
michael@0:   return gDbBackgroundThread;
michael@0: }
michael@0: