1 files changed, 258 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/Classifier.h b/toolkit/components/url-classifier/Classifier.h
new file mode 100644
index 0000000000..a9af2736d2
--- /dev/null
+++ b/toolkit/components/url-classifier/Classifier.h
@@ -0,0 +1,258 @@
+//* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef Classifier_h__
+#define Classifier_h__
+
+#include "Entries.h"
+#include "HashStore.h"
+#include "ProtocolParser.h"
+#include "LookupCache.h"
+#include "mozilla/Atomics.h"
+#include "nsCOMPtr.h"
+#include "nsString.h"
+#include "nsIFile.h"
+
+namespace mozilla {
+
+class LazyIdleThread;
+
+namespace safebrowsing {
+
+/**
+ * Maintains the stores and LookupCaches for the url classifier.
+ */
+class Classifier {
+ public:
+  NS_INLINE_DECL_THREADSAFE_REFCOUNTING(Classifier);
+
+  Classifier();
+
+  nsresult Open(nsIFile& aCacheDirectory);
+  void Close();
+  void Reset();  // Not including any intermediary for update.
+
+  /**
+   * Clear data for specific tables.
+   * If ClearType is Clear_Cache, this function will only clear cache in lookup
+   * cache, otherwise, it will clear data in lookup cache and data stored on
+   * disk.
+   */
+  enum ClearType {
+    Clear_Cache,
+    Clear_All,
+  };
+  void ResetTables(ClearType aType, const nsTArray<nsCString>& aTables);
+
+  /**
+   * Get the list of active tables and their chunks in a format
+   * suitable for an update request.
+   */
+  void TableRequest(nsACString& aResult);
+
+  /*
+   * Get all tables that we know about.
+   */
+  nsresult ActiveTables(nsTArray<nsCString>& aTables) const;
+
+  /**
+   * Check URL fragments against a specified table.
+   * The fragments should be generated by |LookupCache::GetLookupFragments|
+   */
+  nsresult CheckURIFragments(const nsTArray<nsCString>& aSpecFragments,
+                             const nsACString& table,
+                             LookupResultArray& aResults);
+
+  /**
+   * Asynchronously apply updates to the in-use databases. When the
+   * update is complete, the caller can be notified by |aCallback|, which
+   * will occur on the caller thread.
+   */
+  using AsyncUpdateCallback = std::function<void(nsresult)>;
+  nsresult AsyncApplyUpdates(const TableUpdateArray& aUpdates,
+                             const AsyncUpdateCallback& aCallback);
+
+  /**
+   * Wait until the ongoing async update is finished and callback
+   * is fired. Once this function returns, AsyncApplyUpdates is
+   * no longer available.
+   */
+  void FlushAndDisableAsyncUpdate();
+
+  /**
+   * Apply full hashes retrived from gethash to cache.
+   */
+  nsresult ApplyFullHashes(ConstTableUpdateArray& aUpdates);
+
+  /*
+   * Get a bunch of extra prefixes to query for completion
+   * and mask the real entry being requested
+   */
+  nsresult ReadNoiseEntries(const Prefix& aPrefix, const nsACString& aTableName,
+                            uint32_t aCount, PrefixArray& aNoiseEntries);
+
+#ifdef MOZ_SAFEBROWSING_DUMP_FAILED_UPDATES
+  nsresult DumpRawTableUpdates(const nsACString& aRawUpdates);
+#endif
+
+  static void SplitTables(const nsACString& str, nsTArray<nsCString>& tables);
+
+  // Given a root store directory, return a private store directory
+  // based on the table name. To avoid migration issue, the private
+  // store directory is only different from root directory for V4 tables.
+  //
+  // For V4 tables (suffixed by '-proto'), the private directory would
+  // be [root directory path]/[provider]. The provider of V4 tables is
+  // 'google4'.
+  //
+  // Note that if the table name is not owned by any provider, just use
+  // the root directory.
+  static nsresult GetPrivateStoreDirectory(nsIFile* aRootStoreDirectory,
+                                           const nsACString& aTableName,
+                                           const nsACString& aProvider,
+                                           nsIFile** aPrivateStoreDirectory);
+
+  // Swap in in-memory and on-disk database and remove all
+  // update intermediaries.
+  nsresult SwapInNewTablesAndCleanup();
+
+  RefPtr<LookupCache> GetLookupCache(const nsACString& aTable,
+                                     bool aForUpdate = false);
+
+  void GetCacheInfo(const nsACString& aTable,
+                    nsIUrlClassifierCacheInfo** aCache);
+
+  bool OnUpdateThread() const;
+
+ private:
+  ~Classifier();
+
+  void DropStores();
+  void DeleteTables(nsIFile* aDirectory, const nsTArray<nsCString>& aTables);
+
+  nsresult CreateStoreDirectory();
+  nsresult SetupPathNames();
+  nsresult RecoverBackups();
+  nsresult CleanToDelete();
+  nsresult CopyInUseDirForUpdate();
+  nsresult CopyDirectoryInterruptible(nsCOMPtr<nsIFile>& aDestDir,
+                                      nsCOMPtr<nsIFile>& aSourceDir);
+  nsresult RegenActiveTables();
+
+  void MergeNewLookupCaches();  // Merge mNewLookupCaches into mLookupCaches.
+
+  void CopyAndInvalidateFullHashCache();
+
+  // Remove any intermediary for update, including in-memory
+  // and on-disk data.
+  void RemoveUpdateIntermediaries();
+
+#ifdef MOZ_SAFEBROWSING_DUMP_FAILED_UPDATES
+  already_AddRefed<nsIFile> GetFailedUpdateDirectroy();
+  nsresult DumpFailedUpdate();
+#endif
+
+  nsresult ScanStoreDir(nsIFile* aDirectory,
+                        const nsTArray<nsCString>& aExtensions,
+                        nsTArray<nsCString>& aTables);
+
+  nsresult UpdateHashStore(TableUpdateArray& aUpdates,
+                           const nsACString& aTable);
+
+  nsresult UpdateTableV4(TableUpdateArray& aUpdates, const nsACString& aTable);
+
+  nsresult UpdateCache(RefPtr<const TableUpdate> aUpdates);
+
+  RefPtr<LookupCache> GetLookupCacheForUpdate(const nsACString& aTable) {
+    return GetLookupCache(aTable, true);
+  }
+
+  RefPtr<LookupCache> GetLookupCacheFrom(const nsACString& aTable,
+                                         LookupCacheArray& aLookupCaches,
+                                         nsIFile* aRootStoreDirectory);
+
+  bool CheckValidUpdate(TableUpdateArray& aUpdates, const nsACString& aTable);
+
+  nsresult LoadHashStore(nsIFile* aDirectory, nsACString& aResult,
+                         nsTArray<nsCString>& aFailedTableNames);
+
+  nsresult LoadMetadata(nsIFile* aDirectory, nsACString& aResult,
+                        nsTArray<nsCString>& aFailedTableNames);
+
+  static nsCString GetProvider(const nsACString& aTableName);
+
+  /**
+   * The "background" part of ApplyUpdates. Once the background update
+   * is called, the foreground update has to be called along with the
+   * background result no matter whether the background update is
+   * successful or not.
+   */
+  nsresult ApplyUpdatesBackground(TableUpdateArray& aUpdates,
+                                  nsTArray<nsCString>& aFailedTableNames);
+
+  /**
+   * The "foreground" part of ApplyUpdates. The in-use data (in-memory and
+   * on-disk) will be touched so this MUST be mutually exclusive to other
+   * member functions.
+   *
+   * If |aBackgroundRv| is successful, the return value is the result of
+   * bringing stuff to the foreground. Otherwise, the foreground table may
+   * be reset according to the background update failed reason and
+   * |aBackgroundRv| will be returned to forward the background update result.
+   */
+  nsresult ApplyUpdatesForeground(nsresult aBackgroundRv,
+                                  const nsTArray<nsCString>& aFailedTableNames);
+
+  // Used by worker thread and update thread to abort current operation.
+  bool ShouldAbort() const;
+
+  // Add built-in entries for testing.
+  nsresult AddMozEntries(nsTArray<nsCString>& aTables);
+
+  // Remove test files if exist
+  nsresult ClearLegacyFiles();
+
+  // Root dir of the Local profile.
+  nsCOMPtr<nsIFile> mCacheDirectory;
+  // Main directory where to store the databases.
+  nsCOMPtr<nsIFile> mRootStoreDirectory;
+  // Used for atomically updating the other dirs.
+  nsCOMPtr<nsIFile> mBackupDirectory;
+  nsCOMPtr<nsIFile> mUpdatingDirectory;  // For update only.
+  nsCOMPtr<nsIFile> mToDeleteDirectory;
+  LookupCacheArray mLookupCaches;  // For query only.
+  nsTArray<nsCString> mActiveTablesCache;
+  uint32_t mHashKey;
+
+  // In-memory cache for the result of TableRequest. See
+  // nsIUrlClassifierDBService.getTables for the format.
+  nsCString mTableRequestResult;
+
+  // Whether mTableRequestResult is outdated and needs to
+  // be reloaded from disk.
+  bool mIsTableRequestResultOutdated;
+
+  // The copy of mLookupCaches for update only.
+  LookupCacheArray mNewLookupCaches;
+
+  // True when Reset() is called.
+  bool mUpdateInterrupted;
+
+  // True once CLose() has been called
+  Atomic<bool> mIsClosed;
+
+  RefPtr<LazyIdleThread> mUpdateThread;  // For async update.
+
+  // Identical to mRootStoreDirectory but for update only because
+  // nsIFile is not thread safe and mRootStoreDirectory needs to
+  // be accessed in CopyInUseDirForUpdate().
+  // It will be initialized right before update on the worker thread.
+  nsCOMPtr<nsIFile> mRootStoreDirectoryForUpdate;
+};
+
+}  // namespace safebrowsing
+}  // namespace mozilla
+
+#endif