fune/toolkit/components/url-classifier/Classifier.h

//* -*- 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