Bug 1254763 - Use per-provider directory for V4 databases. r=gcp draft
authorHenry Chang <hchang@mozilla.com>
Fri, 02 Sep 2016 11:54:44 +0800
changeset 409060 2e3045926c523b361c5e4b8fa39dc2123d1cc237
parent 409041 d5f20820c80514476f596090292a5d77c4b41e3b
child 409061 4f52f29cc5180b4832111bbde05f335bc3e0c852
push id28371
push userhchang@mozilla.com
push dateFri, 02 Sep 2016 03:55:00 +0000
reviewersgcp
bugs1254763
milestone51.0a1
Bug 1254763 - Use per-provider directory for V4 databases. r=gcp MozReview-Commit-ID: L2pzFsNj53k
toolkit/components/url-classifier/Classifier.cpp
toolkit/components/url-classifier/Classifier.h
toolkit/components/url-classifier/HashStore.cpp
toolkit/components/url-classifier/HashStore.h
toolkit/components/url-classifier/LookupCache.cpp
toolkit/components/url-classifier/LookupCache.h
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -11,16 +11,17 @@
 #include "nsIInputStream.h"
 #include "nsISeekableStream.h"
 #include "nsIFile.h"
 #include "nsNetCID.h"
 #include "nsPrintfCString.h"
 #include "nsThreadUtils.h"
 #include "mozilla/Telemetry.h"
 #include "mozilla/Logging.h"
+#include "mozilla/SyncRunnable.h"
 
 // MOZ_LOG=UrlClassifierDbService:5
 extern mozilla::LazyLogModule gUrlClassifierDbServiceLog;
 #define LOG(args) MOZ_LOG(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug, args)
 #define LOG_ENABLED() MOZ_LOG_TEST(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug)
 
 #define STORE_DIRECTORY      NS_LITERAL_CSTRING("safebrowsing")
 #define TO_DELETE_DIR_SUFFIX NS_LITERAL_CSTRING("-to_delete")
@@ -73,40 +74,184 @@ Classifier::SplitTables(const nsACString
     }
     begin = iter;
     if (begin != end) {
       begin++;
     }
   }
 }
 
+static nsresult
+DeriveProviderFromPref(const nsACString& aTableName, nsCString& aProviderName)
+{
+  // Check all preferences "browser.safebrowsing.provider.[provider].list"
+  // to see which one contains |aTableName|.
+
+  nsCOMPtr<nsIPrefService> prefs = do_GetService(NS_PREFSERVICE_CONTRACTID);
+  NS_ENSURE_TRUE(prefs, NS_ERROR_FAILURE);
+  nsCOMPtr<nsIPrefBranch> prefBranch;
+  nsresult rv = prefs->GetBranch("browser.safebrowsing.provider.",
+                                  getter_AddRefs(prefBranch));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  // We've got a pref branch for "browser.safebrowsing.provider.".
+  // Enumerate all children prefs and parse providers.
+  uint32_t childCount;
+  char** childArray;
+  rv = prefBranch->GetChildList("", &childCount, &childArray);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  // Collect providers from childArray.
+  nsTHashtable<nsCStringHashKey> providers;
+  for (uint32_t i = 0; i < childCount; i++) {
+    nsCString child(childArray[i]);
+    auto dotPos = child.FindChar('.');
+    if (dotPos < 0) {
+      continue;
+    }
+
+    nsDependentCSubstring provider = Substring(child, 0, dotPos);
+
+    providers.PutEntry(provider);
+  }
+  NS_FREE_XPCOM_ALLOCATED_POINTER_ARRAY(childCount, childArray);
+
+  // Now we have all providers. Check which one owns |aTableName|.
+  // e.g. The owning lists of provider "google" is defined in
+  // "browser.safebrowsing.provider.google.lists".
+  for (auto itr = providers.Iter(); !itr.Done(); itr.Next()) {
+    auto entry = itr.Get();
+    nsCString provider(entry->GetKey());
+    nsPrintfCString owninListsPref("%s.lists", provider.get());
+
+    nsXPIDLCString owningLists;
+    nsresult rv = prefBranch->GetCharPref(owninListsPref.get(),
+                                          getter_Copies(owningLists));
+    if (NS_FAILED(rv)) {
+      continue;
+    }
+
+    // We've got the owning lists (represented as string) of |provider|.
+    // Parse the string and see if |aTableName| is included.
+    nsTArray<nsCString> tables;
+    Classifier::SplitTables(owningLists, tables);
+    if (tables.Contains(aTableName)) {
+      aProviderName = provider;
+      return NS_OK;
+    }
+  }
+
+  return NS_ERROR_FAILURE;
+}
+
+// Lookup the provider name by table name on non-main thread.
+// On main thread, just call DeriveProviderFromPref() instead
+// but DeriveProviderFromPref is supposed to used internally.
+static nsCString
+GetProviderByTableName(const nsACString& aTableName)
+{
+  MOZ_ASSERT(!NS_IsMainThread(), "GetProviderByTableName MUST be called "
+                                 "on non-main thread.");
+  nsCString providerName;
+
+  nsCOMPtr<nsIRunnable> r = NS_NewRunnableFunction([&aTableName,
+                                                    &providerName] () -> void {
+    nsresult rv = DeriveProviderFromPref(aTableName, providerName);
+    if (NS_FAILED(rv)) {
+      LOG(("No provider found for %s", nsCString(aTableName).get()));
+    }
+  });
+
+  nsCOMPtr<nsIThread> mainThread = do_GetMainThread();
+  SyncRunnable::DispatchToThread(mainThread, r);
+
+  return providerName;
+}
+
+nsresult
+Classifier::GetPrivateStoreDirectory(nsIFile* aRootStoreDirectory,
+                                     const nsACString& aTableName,
+                                     nsIFile** aPrivateStoreDirectory)
+{
+  NS_ENSURE_ARG_POINTER(aPrivateStoreDirectory);
+
+  if (!StringEndsWith(aTableName, NS_LITERAL_CSTRING("-proto"))) {
+    // Only V4 table names (ends with '-proto') would be stored
+    // to per-provider sub-directory.
+    nsCOMPtr<nsIFile>(aRootStoreDirectory).forget(aPrivateStoreDirectory);
+    return NS_OK;
+  }
+
+  nsCString providerName = GetProviderByTableName(aTableName);
+  if (providerName.IsEmpty()) {
+    // When failing to get provider, just store in the root folder.
+    nsCOMPtr<nsIFile>(aRootStoreDirectory).forget(aPrivateStoreDirectory);
+    return NS_OK;
+  }
+
+  nsCOMPtr<nsIFile> providerDirectory;
+
+  // Clone first since we are gonna create a new directory.
+  nsresult rv = aRootStoreDirectory->Clone(getter_AddRefs(providerDirectory));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  // Append the provider name to the root store directory.
+  rv = providerDirectory->AppendNative(providerName);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  // Ensure existence of the provider directory.
+  bool dirExists;
+  rv = providerDirectory->Exists(&dirExists);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  if (!dirExists) {
+    LOG(("Creating private directory for %s", nsCString(aTableName).get()));
+    rv = providerDirectory->Create(nsIFile::DIRECTORY_TYPE, 0755);
+    NS_ENSURE_SUCCESS(rv, rv);
+    providerDirectory.forget(aPrivateStoreDirectory);
+    return rv;
+  }
+
+  // Store directory exists. Check if it's a directory.
+  bool isDir;
+  rv = providerDirectory->IsDirectory(&isDir);
+  NS_ENSURE_SUCCESS(rv, rv);
+  if (!isDir) {
+    return NS_ERROR_FILE_DESTINATION_NOT_DIR;
+  }
+
+  providerDirectory.forget(aPrivateStoreDirectory);
+
+  return NS_OK;
+}
+
 Classifier::Classifier()
 {
 }
 
 Classifier::~Classifier()
 {
   Close();
 }
 
 nsresult
 Classifier::SetupPathNames()
 {
   // Get the root directory where to store all the databases.
-  nsresult rv = mCacheDirectory->Clone(getter_AddRefs(mStoreDirectory));
+  nsresult rv = mCacheDirectory->Clone(getter_AddRefs(mRootStoreDirectory));
   NS_ENSURE_SUCCESS(rv, rv);
 
-  rv = mStoreDirectory->AppendNative(STORE_DIRECTORY);
+  rv = mRootStoreDirectory->AppendNative(STORE_DIRECTORY);
   NS_ENSURE_SUCCESS(rv, rv);
 
   // Make sure LookupCaches (which are persistent and survive updates)
   // are reading/writing in the right place. We will be moving their
   // files "underneath" them during backup/restore.
   for (uint32_t i = 0; i < mLookupCaches.Length(); i++) {
-    mLookupCaches[i]->UpdateDirHandle(mStoreDirectory);
+    mLookupCaches[i]->UpdateRootDirHandle(mRootStoreDirectory);
   }
 
   // Directory where to move a backup before an update.
   rv = mCacheDirectory->Clone(getter_AddRefs(mBackupDirectory));
   NS_ENSURE_SUCCESS(rv, rv);
 
   rv = mBackupDirectory->AppendNative(STORE_DIRECTORY + BACKUP_DIR_SUFFIX);
   NS_ENSURE_SUCCESS(rv, rv);
@@ -122,25 +267,25 @@ Classifier::SetupPathNames()
   return NS_OK;
 }
 
 nsresult
 Classifier::CreateStoreDirectory()
 {
   // Ensure the safebrowsing directory exists.
   bool storeExists;
-  nsresult rv = mStoreDirectory->Exists(&storeExists);
+  nsresult rv = mRootStoreDirectory->Exists(&storeExists);
   NS_ENSURE_SUCCESS(rv, rv);
 
   if (!storeExists) {
-    rv = mStoreDirectory->Create(nsIFile::DIRECTORY_TYPE, 0755);
+    rv = mRootStoreDirectory->Create(nsIFile::DIRECTORY_TYPE, 0755);
     NS_ENSURE_SUCCESS(rv, rv);
   } else {
     bool storeIsDir;
-    rv = mStoreDirectory->IsDirectory(&storeIsDir);
+    rv = mRootStoreDirectory->IsDirectory(&storeIsDir);
     NS_ENSURE_SUCCESS(rv, rv);
     if (!storeIsDir)
       return NS_ERROR_FILE_DESTINATION_NOT_DIR;
   }
 
   return NS_OK;
 }
 
@@ -184,17 +329,17 @@ Classifier::Close()
   DropStores();
 }
 
 void
 Classifier::Reset()
 {
   DropStores();
 
-  mStoreDirectory->Remove(true);
+  mRootStoreDirectory->Remove(true);
   mBackupDirectory->Remove(true);
   mToDeleteDirectory->Remove(true);
 
   CreateStoreDirectory();
 
   mTableFreshness.Clear();
   RegenActiveTables();
 }
@@ -210,17 +355,17 @@ Classifier::ResetTables(const nsTArray<n
 
   RegenActiveTables();
 }
 
 void
 Classifier::DeleteTables(const nsTArray<nsCString>& aTables)
 {
   nsCOMPtr<nsISimpleEnumerator> entries;
-  nsresult rv = mStoreDirectory->GetDirectoryEntries(getter_AddRefs(entries));
+  nsresult rv = mRootStoreDirectory->GetDirectoryEntries(getter_AddRefs(entries));
   NS_ENSURE_SUCCESS_VOID(rv);
 
   bool hasMore;
   while (NS_SUCCEEDED(rv = entries->HasMoreElements(&hasMore)) && hasMore) {
     nsCOMPtr<nsISupports> supports;
     rv = entries->GetNext(getter_AddRefs(supports));
     NS_ENSURE_SUCCESS_VOID(rv);
 
@@ -243,17 +388,17 @@ Classifier::DeleteTables(const nsTArray<
 }
 
 void
 Classifier::TableRequest(nsACString& aResult)
 {
   nsTArray<nsCString> tables;
   ActiveTables(tables);
   for (uint32_t i = 0; i < tables.Length(); i++) {
-    HashStore store(tables[i], mStoreDirectory);
+    HashStore store(tables[i], mRootStoreDirectory);
 
     nsresult rv = store.Open();
     if (NS_FAILED(rv))
       continue;
 
     aResult.Append(store.TableName());
     aResult.Append(';');
 
@@ -497,17 +642,17 @@ Classifier::RegenActiveTables()
 {
   mActiveTablesCache.Clear();
 
   nsTArray<nsCString> foundTables;
   ScanStoreDir(foundTables);
 
   for (uint32_t i = 0; i < foundTables.Length(); i++) {
     nsCString table(foundTables[i]);
-    HashStore store(table, mStoreDirectory);
+    HashStore store(table, mRootStoreDirectory);
 
     nsresult rv = store.Open();
     if (NS_FAILED(rv))
       continue;
 
     LookupCache *lookupCache = GetLookupCache(store.TableName());
     if (!lookupCache) {
       continue;
@@ -528,17 +673,17 @@ Classifier::RegenActiveTables()
 
   return NS_OK;
 }
 
 nsresult
 Classifier::ScanStoreDir(nsTArray<nsCString>& aTables)
 {
   nsCOMPtr<nsISimpleEnumerator> entries;
-  nsresult rv = mStoreDirectory->GetDirectoryEntries(getter_AddRefs(entries));
+  nsresult rv = mRootStoreDirectory->GetDirectoryEntries(getter_AddRefs(entries));
   NS_ENSURE_SUCCESS(rv, rv);
 
   bool hasMore;
   while (NS_SUCCEEDED(rv = entries->HasMoreElements(&hasMore)) && hasMore) {
     nsCOMPtr<nsISupports> supports;
     rv = entries->GetNext(getter_AddRefs(supports));
     NS_ENSURE_SUCCESS(rv, rv);
 
@@ -592,23 +737,23 @@ Classifier::BackupTables()
   // dir always has a valid, complete copy, instead of a partial one,
   // because that's the one we will copy over the normal store dir.
 
   nsCString backupDirName;
   nsresult rv = mBackupDirectory->GetNativeLeafName(backupDirName);
   NS_ENSURE_SUCCESS(rv, rv);
 
   nsCString storeDirName;
-  rv = mStoreDirectory->GetNativeLeafName(storeDirName);
+  rv = mRootStoreDirectory->GetNativeLeafName(storeDirName);
   NS_ENSURE_SUCCESS(rv, rv);
 
-  rv = mStoreDirectory->MoveToNative(nullptr, backupDirName);
+  rv = mRootStoreDirectory->MoveToNative(nullptr, backupDirName);
   NS_ENSURE_SUCCESS(rv, rv);
 
-  rv = mStoreDirectory->CopyToNative(nullptr, storeDirName);
+  rv = mRootStoreDirectory->CopyToNative(nullptr, storeDirName);
   NS_ENSURE_SUCCESS(rv, rv);
 
   // We moved some things to new places, so move the handles around, too.
   rv = SetupPathNames();
   NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
@@ -635,25 +780,25 @@ Classifier::RecoverBackups()
 {
   bool backupExists;
   nsresult rv = mBackupDirectory->Exists(&backupExists);
   NS_ENSURE_SUCCESS(rv, rv);
 
   if (backupExists) {
     // Remove the safebrowsing dir if it exists
     nsCString storeDirName;
-    rv = mStoreDirectory->GetNativeLeafName(storeDirName);
+    rv = mRootStoreDirectory->GetNativeLeafName(storeDirName);
     NS_ENSURE_SUCCESS(rv, rv);
 
     bool storeExists;
-    rv = mStoreDirectory->Exists(&storeExists);
+    rv = mRootStoreDirectory->Exists(&storeExists);
     NS_ENSURE_SUCCESS(rv, rv);
 
     if (storeExists) {
-      rv = mStoreDirectory->Remove(true);
+      rv = mRootStoreDirectory->Remove(true);
       NS_ENSURE_SUCCESS(rv, rv);
     }
 
     // Move the backup to the store location
     rv = mBackupDirectory->MoveToNative(nullptr, storeDirName);
     NS_ENSURE_SUCCESS(rv, rv);
 
     // mBackupDirectory now points to storeDir, fix up.
@@ -695,17 +840,17 @@ Classifier::CheckValidUpdate(nsTArray<Ta
  * This will consume+delete updates from the passed nsTArray.
 */
 nsresult
 Classifier::UpdateHashStore(nsTArray<TableUpdate*>* aUpdates,
                             const nsACString& aTable)
 {
   LOG(("Classifier::UpdateHashStore(%s)", PromiseFlatCString(aTable).get()));
 
-  HashStore store(aTable, mStoreDirectory);
+  HashStore store(aTable, mRootStoreDirectory);
 
   if (!CheckValidUpdate(aUpdates, store.TableName())) {
     return NS_OK;
   }
 
   nsresult rv = store.Open();
   NS_ENSURE_SUCCESS(rv, rv);
   rv = store.BeginUpdate();
@@ -816,17 +961,17 @@ LookupCache *
 Classifier::GetLookupCache(const nsACString& aTable)
 {
   for (uint32_t i = 0; i < mLookupCaches.Length(); i++) {
     if (mLookupCaches[i]->TableName().Equals(aTable)) {
       return mLookupCaches[i];
     }
   }
 
-  UniquePtr<LookupCache> cache(new LookupCache(aTable, mStoreDirectory));
+  UniquePtr<LookupCache> cache(new LookupCache(aTable, mRootStoreDirectory));
   nsresult rv = cache->Init();
   if (NS_FAILED(rv)) {
     return nullptr;
   }
   rv = cache->Open();
   if (NS_FAILED(rv)) {
     if (rv == NS_ERROR_FILE_CORRUPTED) {
       Reset();
--- a/toolkit/components/url-classifier/Classifier.h
+++ b/toolkit/components/url-classifier/Classifier.h
@@ -80,16 +80,30 @@ public:
    * and mask the real entry being requested
    */
   nsresult ReadNoiseEntries(const Prefix& aPrefix,
                             const nsACString& aTableName,
                             uint32_t aCount,
                             PrefixArray* aNoiseEntries);
   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,
+                                           nsIFile** aPrivateStoreDirectory);
+
 private:
   void DropStores();
   void DeleteTables(const nsTArray<nsCString>& aTables);
 
   nsresult CreateStoreDirectory();
   nsresult SetupPathNames();
   nsresult RecoverBackups();
   nsresult CleanToDelete();
@@ -106,17 +120,17 @@ private:
   LookupCache *GetLookupCache(const nsACString& aTable);
 
   bool CheckValidUpdate(nsTArray<TableUpdate*>* aUpdates,
                         const nsACString& aTable);
 
   // Root dir of the Local profile.
   nsCOMPtr<nsIFile> mCacheDirectory;
   // Main directory where to store the databases.
-  nsCOMPtr<nsIFile> mStoreDirectory;
+  nsCOMPtr<nsIFile> mRootStoreDirectory;
   // Used for atomically updating the other dirs.
   nsCOMPtr<nsIFile> mBackupDirectory;
   nsCOMPtr<nsIFile> mToDeleteDirectory;
   nsCOMPtr<nsICryptoHash> mCryptoHash;
   nsTArray<LookupCache*> mLookupCaches;
   nsTArray<nsCString> mActiveTablesCache;
   uint32_t mHashKey;
   // Stores the last time a given table was updated (seconds).
--- a/toolkit/components/url-classifier/HashStore.cpp
+++ b/toolkit/components/url-classifier/HashStore.cpp
@@ -33,16 +33,17 @@
 #include "nsICryptoHash.h"
 #include "nsISeekableStream.h"
 #include "nsIStreamConverterService.h"
 #include "nsNetUtil.h"
 #include "nsCheckSummedOutputStream.h"
 #include "prio.h"
 #include "mozilla/Logging.h"
 #include "zlib.h"
+#include "Classifier.h"
 
 // Main store for SafeBrowsing protocol data. We store
 // known add/sub chunks, prefixes and completions in memory
 // during an update, and serialize to disk.
 // We do not store the add prefixes, those are retrieved by
 // decompressing the PrefixSet cache whenever we need to apply
 // an update.
 //
@@ -174,22 +175,28 @@ TableUpdateV4::NewPrefixes(int32_t aSize
 void
 TableUpdateV4::NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices)
 {
   for (size_t i = 0; i < aNumOfIndices; i++) {
     mRemovalIndiceArray.AppendElement(aIndices[i]);
   }
 }
 
-HashStore::HashStore(const nsACString& aTableName, nsIFile* aStoreDir)
+HashStore::HashStore(const nsACString& aTableName, nsIFile* aRootStoreDir)
   : mTableName(aTableName)
-  , mStoreDirectory(aStoreDir)
   , mInUpdate(false)
   , mFileSize(0)
 {
+  nsresult rv = Classifier::GetPrivateStoreDirectory(aRootStoreDir,
+                                                     aTableName,
+                                                     getter_AddRefs(mStoreDirectory));
+  if (NS_FAILED(rv)) {
+    LOG(("Failed to get private store directory for %s", mTableName.get()));
+    mStoreDirectory = aRootStoreDir;
+  }
 }
 
 HashStore::~HashStore()
 {
 }
 
 nsresult
 HashStore::Reset()
--- a/toolkit/components/url-classifier/HashStore.h
+++ b/toolkit/components/url-classifier/HashStore.h
@@ -175,17 +175,17 @@ private:
 
   PrefixesStringMap mPrefixesMap;
   RemovalIndiceArray mRemovalIndiceArray;
 };
 
 // There is one hash store per table.
 class HashStore {
 public:
-  HashStore(const nsACString& aTableName, nsIFile* aStoreFile);
+  HashStore(const nsACString& aTableName, nsIFile* aRootStoreFile);
   ~HashStore();
 
   const nsCString& TableName() const { return mTableName; }
 
   nsresult Open();
   // Add Prefixes are stored partly in the PrefixSet (contains the
   // Prefix data organized for fast lookup/low RAM usage) and partly in the
   // HashStore (Add Chunk numbers - only used for updates, slow retrieval).
--- a/toolkit/components/url-classifier/LookupCache.cpp
+++ b/toolkit/components/url-classifier/LookupCache.cpp
@@ -5,16 +5,18 @@
 
 #include "LookupCache.h"
 #include "HashStore.h"
 #include "nsISeekableStream.h"
 #include "mozilla/Telemetry.h"
 #include "mozilla/Logging.h"
 #include "nsNetUtil.h"
 #include "prprf.h"
+#include "nsIUrlListManager.h"
+#include "Classifier.h"
 
 // We act as the main entry point for all the real lookups,
 // so note that those are not done to the actual HashStore.
 // The latter solely exists to store the data needed to handle
 // the updates from the protocol.
 
 // This module provides a front for PrefixSet, mUpdateCompletions,
 // and mGetHashCache, which together contain everything needed to
@@ -34,21 +36,22 @@
 // MOZ_LOG=UrlClassifierDbService:5
 extern mozilla::LazyLogModule gUrlClassifierDbServiceLog;
 #define LOG(args) MOZ_LOG(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug, args)
 #define LOG_ENABLED() MOZ_LOG_TEST(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug)
 
 namespace mozilla {
 namespace safebrowsing {
 
-LookupCache::LookupCache(const nsACString& aTableName, nsIFile* aStoreDir)
+LookupCache::LookupCache(const nsACString& aTableName, nsIFile* aRootStoreDir)
   : mPrimed(false)
   , mTableName(aTableName)
-  , mStoreDirectory(aStoreDir)
+  , mRootStoreDirectory(aRootStoreDir)
 {
+  UpdateRootDirHandle(mRootStoreDirectory);
 }
 
 nsresult
 LookupCache::Init()
 {
   mPrefixSet = new nsUrlClassifierPrefixSet();
   nsresult rv = mPrefixSet->Init(mTableName);
   NS_ENSURE_SUCCESS(rv, rv);
@@ -70,19 +73,42 @@ LookupCache::Open()
   LOG(("Loading PrefixSet"));
   rv = LoadPrefixSet();
   NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
 nsresult
-LookupCache::UpdateDirHandle(nsIFile* aStoreDirectory)
+LookupCache::UpdateRootDirHandle(nsIFile* aNewRootStoreDirectory)
 {
-  return aStoreDirectory->Clone(getter_AddRefs(mStoreDirectory));
+  nsresult rv;
+
+  if (aNewRootStoreDirectory != mRootStoreDirectory) {
+    rv = aNewRootStoreDirectory->Clone(getter_AddRefs(mRootStoreDirectory));
+    NS_ENSURE_SUCCESS(rv, rv);
+  }
+
+  rv = Classifier::GetPrivateStoreDirectory(mRootStoreDirectory,
+                                            mTableName,
+                                            getter_AddRefs(mStoreDirectory));
+
+  if (NS_FAILED(rv)) {
+    LOG(("Failed to get private store directory for %s", mTableName.get()));
+    mStoreDirectory = mRootStoreDirectory;
+  }
+
+  if (LOG_ENABLED()) {
+    nsString path;
+    mStoreDirectory->GetPath(path);
+    LOG(("Private store directory for %s is %s", mTableName.get(),
+                                                 NS_ConvertUTF16toUTF8(path).get()));
+  }
+
+  return rv;
 }
 
 nsresult
 LookupCache::Reset()
 {
   LOG(("LookupCache resetting"));
 
   nsCOMPtr<nsIFile> prefixsetFile;
@@ -231,17 +257,17 @@ void
 LookupCache::ClearCache()
 {
   mGetHashCache.Clear();
 }
 
 nsresult
 LookupCache::ReadCompletions()
 {
-  HashStore store(mTableName, mStoreDirectory);
+  HashStore store(mTableName, mRootStoreDirectory);
 
   nsresult rv = store.Open();
   NS_ENSURE_SUCCESS(rv, rv);
 
   mUpdateCompletions.Clear();
 
   const AddCompleteArray& addComplete = store.AddCompletes();
   for (uint32_t i = 0; i < addComplete.Length(); i++) {
--- a/toolkit/components/url-classifier/LookupCache.h
+++ b/toolkit/components/url-classifier/LookupCache.h
@@ -99,17 +99,17 @@ public:
   ~LookupCache();
 
   const nsCString &TableName() const { return mTableName; }
 
   nsresult Init();
   nsresult Open();
   // The directory handle where we operate will
   // be moved away when a backup is made.
-  nsresult UpdateDirHandle(nsIFile* aStoreDirectory);
+  nsresult UpdateRootDirHandle(nsIFile* aRootStoreDirectory);
   // This will Clear() the passed arrays when done.
   nsresult Build(AddPrefixArray& aAddPrefixes,
                  AddCompleteArray& aAddCompletes);
   nsresult AddCompletionsToCache(AddCompleteArray& aAddCompletes);
   nsresult GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes);
   void ClearUpdatedCompletions();
   void ClearCache();
 
@@ -129,16 +129,17 @@ private:
   nsresult LoadPrefixSet();
   nsresult LoadCompletions();
   // Construct a Prefix Set with known prefixes.
   // This will Clear() aAddPrefixes when done.
   nsresult ConstructPrefixSet(AddPrefixArray& aAddPrefixes);
 
   bool mPrimed;
   nsCString mTableName;
+  nsCOMPtr<nsIFile> mRootStoreDirectory;
   nsCOMPtr<nsIFile> mStoreDirectory;
   // Set of prefixes known to be in the database
   RefPtr<nsUrlClassifierPrefixSet> mPrefixSet;
   // Full length hashes obtained in update request
   CompletionArray mUpdateCompletions;
   // Full length hashes obtained in gethash request
   CompletionArray mGetHashCache;
 };