Bug 1283009 - Store variable-length prefix to disk. r=francois, gcp draft
authorDimi Lee <dlee@mozilla.com>
Mon, 19 Sep 2016 11:51:01 +0800
changeset 414982 66207456cdcc328994de07e0e503778d6953b6ec
parent 414776 5986afdcfdd4307d6d4fb7ddd02f9c35d56a9e32
child 415026 c409a9694f76ba3f62d56fb1b8ceb684c08b3945
push id29767
push userdlee@mozilla.com
push dateMon, 19 Sep 2016 03:53:58 +0000
reviewersfrancois, gcp
bugs1283009
milestone51.0a1
Bug 1283009 - Store variable-length prefix to disk. r=francois, gcp MozReview-Commit-ID: BMTGtgMuQdg
toolkit/components/url-classifier/Classifier.cpp
toolkit/components/url-classifier/Classifier.h
toolkit/components/url-classifier/HashStore.h
toolkit/components/url-classifier/LookupCache.cpp
toolkit/components/url-classifier/LookupCache.h
toolkit/components/url-classifier/ProtocolParser.cpp
toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
toolkit/components/url-classifier/tests/gtest/moz.build
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -512,40 +512,35 @@ Classifier::ApplyUpdates(nsTArray<TableU
     clockStart = PR_IntervalNow();
   }
 
   nsresult rv;
 
   {
     ScopedUpdatesClearer scopedUpdatesClearer(aUpdates);
 
-    // In order to prevent any premature update code from being
-    // run against V4 updates, we bail out as early as possible
-    // if aUpdates is using V4.
-    for (auto update : *aUpdates) {
-      if (update && TableUpdate::Cast<TableUpdateV4>(update)) {
-        LOG(("V4 update is not supported yet."));
-        // TODO: Bug 1283009 - Supports applying table udpate V4.
-        return NS_ERROR_NOT_IMPLEMENTED;
-      }
-    }
-
     LOG(("Backup before update."));
 
     rv = BackupTables();
     NS_ENSURE_SUCCESS(rv, rv);
 
     LOG(("Applying %d table updates.", aUpdates->Length()));
 
     for (uint32_t i = 0; i < aUpdates->Length(); i++) {
       // Previous UpdateHashStore() may have consumed this update..
       if ((*aUpdates)[i]) {
         // Run all updates for one table
         nsCString updateTable(aUpdates->ElementAt(i)->TableName());
-        rv = UpdateHashStore(aUpdates, updateTable);
+
+        if (TableUpdate::Cast<TableUpdateV2>((*aUpdates)[i])) {
+          rv = UpdateHashStore(aUpdates, updateTable);
+        } else {
+          rv = UpdateTableV4(aUpdates, updateTable);
+        }
+
         if (NS_FAILED(rv)) {
           if (rv != NS_ERROR_OUT_OF_MEMORY) {
             Reset();
           }
           return rv;
         }
       }
     }
@@ -852,17 +847,18 @@ Classifier::UpdateHashStore(nsTArray<Tab
   }
 
   nsresult rv = store.Open();
   NS_ENSURE_SUCCESS(rv, rv);
   rv = store.BeginUpdate();
   NS_ENSURE_SUCCESS(rv, rv);
 
   // Read the part of the store that is (only) in the cache
-  LookupCache *lookupCache = GetLookupCache(store.TableName());
+  LookupCacheV2* lookupCache =
+    LookupCache::Cast<LookupCacheV2>(GetLookupCache(store.TableName()));
   if (!lookupCache) {
     return NS_ERROR_FAILURE;
   }
 
   // Clear cache when update
   lookupCache->ClearCache();
 
   FallibleTArray<uint32_t> AddPrefixHashes;
@@ -917,40 +913,98 @@ Classifier::UpdateHashStore(nsTArray<Tab
   NS_ENSURE_SUCCESS(rv, rv);
 
   // At this point the store is updated and written out to disk, but
   // the data is still in memory.  Build our quick-lookup table here.
   rv = lookupCache->Build(store.AddPrefixes(), store.AddCompletes());
   NS_ENSURE_SUCCESS(rv, rv);
 
 #if defined(DEBUG)
-  lookupCache->Dump();
+  lookupCache->DumpCompletions();
 #endif
   rv = lookupCache->WriteFile();
   NS_ENSURE_SUCCESS(rv, rv);
 
   int64_t now = (PR_Now() / PR_USEC_PER_SEC);
   LOG(("Successfully updated %s", store.TableName().get()));
   mTableFreshness.Put(store.TableName(), now);
 
   return NS_OK;
 }
 
 nsresult
+Classifier::UpdateTableV4(nsTArray<TableUpdate*>* aUpdates,
+                          const nsACString& aTable)
+{
+  LOG(("Classifier::UpdateTableV4(%s)", PromiseFlatCString(aTable).get()));
+
+  if (!CheckValidUpdate(aUpdates, aTable)) {
+    return NS_OK;
+  }
+
+  LookupCacheV4* lookupCache =
+    LookupCache::Cast<LookupCacheV4>(GetLookupCache(aTable));
+  if (!lookupCache) {
+    return NS_ERROR_FAILURE;
+  }
+
+  PrefixStringMap prefixes;
+  for (uint32_t i = 0; i < aUpdates->Length(); i++) {
+    TableUpdate *update = aUpdates->ElementAt(i);
+    if (!update || !update->TableName().Equals(aTable)) {
+      continue;
+    }
+
+    auto updateV4 = TableUpdate::Cast<TableUpdateV4>(update);
+    NS_ENSURE_TRUE(updateV4, NS_ERROR_FAILURE);
+
+    if (updateV4->IsFullUpdate()) {
+      prefixes.Clear();
+      TableUpdateV4::PrefixesStringMap& map = updateV4->Prefixes();
+
+      for (auto iter = map.Iter(); !iter.Done(); iter.Next()) {
+        // prefixes is an nsClassHashtable object stores prefix string.
+        // It will take the ownership of the put object.
+        nsCString* prefix = new nsCString(iter.Data()->GetPrefixString());
+        prefixes.Put(iter.Key(), prefix);
+      }
+    } else {
+      // TODO: Bug 1287058, partial update
+    }
+
+    aUpdates->ElementAt(i) = nullptr;
+  }
+
+  nsresult rv = lookupCache->Build(prefixes);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  rv = lookupCache->WriteFile();
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  int64_t now = (PR_Now() / PR_USEC_PER_SEC);
+  LOG(("Successfully updated %s\n", PromiseFlatCString(aTable).get()));
+  mTableFreshness.Put(aTable, now);
+
+  return NS_OK;
+}
+
+nsresult
 Classifier::UpdateCache(TableUpdate* aUpdate)
 {
   if (!aUpdate) {
     return NS_OK;
   }
 
   nsAutoCString table(aUpdate->TableName());
   LOG(("Classifier::UpdateCache(%s)", table.get()));
 
   LookupCache *lookupCache = GetLookupCache(table);
-  NS_ENSURE_TRUE(lookupCache, NS_ERROR_FAILURE);
+  if (!lookupCache) {
+    return NS_ERROR_FAILURE;
+  }
 
   auto updateV2 = TableUpdate::Cast<TableUpdateV2>(aUpdate);
   lookupCache->AddCompletionsToCache(updateV2->AddCompletes());
 
 #if defined(DEBUG)
   lookupCache->DumpCache();
 #endif
 
@@ -961,17 +1015,26 @@ 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, mRootStoreDirectory));
+  // TODO : Bug 1302600, It would be better if we have a more general non-main
+  //        thread method to convert table name to protocol version. Currently
+  //        we can only know this by checking if the table name ends with '-proto'.
+  UniquePtr<LookupCache> cache;
+  if (StringEndsWith(aTable, NS_LITERAL_CSTRING("-proto"))) {
+    cache = MakeUnique<LookupCacheV4>(aTable, mRootStoreDirectory);
+  } else {
+    cache = MakeUnique<LookupCacheV2>(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();
@@ -983,17 +1046,18 @@ Classifier::GetLookupCache(const nsACStr
 }
 
 nsresult
 Classifier::ReadNoiseEntries(const Prefix& aPrefix,
                              const nsACString& aTableName,
                              uint32_t aCount,
                              PrefixArray* aNoiseEntries)
 {
-  LookupCache *cache = GetLookupCache(aTableName);
+  // TODO : Bug 1297962, support adding noise for v4
+  LookupCacheV2 *cache = static_cast<LookupCacheV2*>(GetLookupCache(aTableName));
   if (!cache) {
     return NS_ERROR_FAILURE;
   }
 
   FallibleTArray<uint32_t> prefixes;
   nsresult rv = cache->GetPrefixes(prefixes);
   NS_ENSURE_SUCCESS(rv, rv);
 
--- a/toolkit/components/url-classifier/Classifier.h
+++ b/toolkit/components/url-classifier/Classifier.h
@@ -110,16 +110,19 @@ private:
   nsresult BackupTables();
   nsresult RemoveBackupTables();
   nsresult RegenActiveTables();
   nsresult ScanStoreDir(nsTArray<nsCString>& aTables);
 
   nsresult UpdateHashStore(nsTArray<TableUpdate*>* aUpdates,
                            const nsACString& aTable);
 
+  nsresult UpdateTableV4(nsTArray<TableUpdate*>* aUpdates,
+                         const nsACString& aTable);
+
   nsresult UpdateCache(TableUpdate* aUpdates);
 
   LookupCache *GetLookupCache(const nsACString& aTable);
 
   bool CheckValidUpdate(nsTArray<TableUpdate*>* aUpdates,
                         const nsACString& aTable);
 
   // Root dir of the Local profile.
--- a/toolkit/components/url-classifier/HashStore.h
+++ b/toolkit/components/url-classifier/HashStore.h
@@ -148,36 +148,40 @@ public:
   };
 
   typedef nsClassHashtable<nsUint32HashKey, PrefixString> PrefixesStringMap;
   typedef nsTArray<int32_t> RemovalIndiceArray;
 
 public:
   explicit TableUpdateV4(const nsACString& aTable)
     : TableUpdate(aTable)
+    , mFullUpdate(false)
   {
   }
 
   bool Empty() const override
   {
     return mPrefixesMap.IsEmpty() && mRemovalIndiceArray.IsEmpty();
   }
 
+  bool IsFullUpdate() const { return mFullUpdate; }
   PrefixesStringMap& Prefixes() { return mPrefixesMap; }
   RemovalIndiceArray& RemovalIndices() { return mRemovalIndiceArray; }
 
   // For downcasting.
   static const int TAG = 4;
 
+  void SetFullUpdate(bool aIsFullUpdate) { mFullUpdate = aIsFullUpdate; }
   void NewPrefixes(int32_t aSize, std::string& aPrefixes);
   void NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices);
 
 private:
   virtual int Tag() const override { return TAG; }
 
+  bool mFullUpdate;
   PrefixesStringMap mPrefixesMap;
   RemovalIndiceArray mRemovalIndiceArray;
 };
 
 // There is one hash store per table.
 class HashStore {
 public:
   HashStore(const nsACString& aTableName, nsIFile* aRootStoreFile);
--- a/toolkit/components/url-classifier/LookupCache.cpp
+++ b/toolkit/components/url-classifier/LookupCache.cpp
@@ -35,47 +35,32 @@
 // 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 {
 
+const int LookupCacheV2::VER = 2;
+const int LookupCacheV4::VER = 4;
+
 LookupCache::LookupCache(const nsACString& aTableName, nsIFile* aRootStoreDir)
   : mPrimed(false)
   , mTableName(aTableName)
   , mRootStoreDirectory(aRootStoreDir)
 {
   UpdateRootDirHandle(mRootStoreDirectory);
 }
 
 nsresult
-LookupCache::Init()
-{
-  mPrefixSet = new nsUrlClassifierPrefixSet();
-  nsresult rv = mPrefixSet->Init(mTableName);
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  return NS_OK;
-}
-
-LookupCache::~LookupCache()
-{
-}
-
-nsresult
 LookupCache::Open()
 {
-  LOG(("Reading Completions"));
-  nsresult rv = ReadCompletions();
-  NS_ENSURE_SUCCESS(rv, rv);
-
   LOG(("Loading PrefixSet"));
-  rv = LoadPrefixSet();
+  nsresult rv = LoadPrefixSet();
   NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
 nsresult
 LookupCache::UpdateRootDirHandle(nsIFile* aNewRootStoreDirectory)
 {
@@ -120,42 +105,16 @@ LookupCache::Reset()
   rv = prefixsetFile->Remove(false);
   NS_ENSURE_SUCCESS(rv, rv);
 
   ClearAll();
 
   return NS_OK;
 }
 
-
-nsresult
-LookupCache::Build(AddPrefixArray& aAddPrefixes,
-                   AddCompleteArray& aAddCompletes)
-{
-  Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_COMPLETIONS,
-                        static_cast<uint32_t>(aAddCompletes.Length()));
-
-  mUpdateCompletions.Clear();
-  mUpdateCompletions.SetCapacity(aAddCompletes.Length());
-  for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
-    mUpdateCompletions.AppendElement(aAddCompletes[i].CompleteHash());
-  }
-  aAddCompletes.Clear();
-  mUpdateCompletions.Sort();
-
-  Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_PREFIXES,
-                        static_cast<uint32_t>(aAddPrefixes.Length()));
-
-  nsresult rv = ConstructPrefixSet(aAddPrefixes);
-  NS_ENSURE_SUCCESS(rv, rv);
-  mPrimed = true;
-
-  return NS_OK;
-}
-
 nsresult
 LookupCache::AddCompletionsToCache(AddCompleteArray& aAddCompletes)
 {
   for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
     if (mGetHashCache.BinaryIndexOf(aAddCompletes[i].CompleteHash()) == mGetHashCache.NoIndex) {
       mGetHashCache.AppendElement(aAddCompletes[i].CompleteHash());
     }
   }
@@ -172,115 +131,48 @@ LookupCache::DumpCache()
     return;
 
   for (uint32_t i = 0; i < mGetHashCache.Length(); i++) {
     nsAutoCString str;
     mGetHashCache[i].ToHexString(str);
     LOG(("Caches: %s", str.get()));
   }
 }
-
-void
-LookupCache::Dump()
-{
-  if (!LOG_ENABLED())
-    return;
-
-  for (uint32_t i = 0; i < mUpdateCompletions.Length(); i++) {
-    nsAutoCString str;
-    mUpdateCompletions[i].ToHexString(str);
-    LOG(("Update: %s", str.get()));
-  }
-}
 #endif
 
 nsresult
-LookupCache::Has(const Completion& aCompletion,
-                 bool* aHas, bool* aComplete)
-{
-  *aHas = *aComplete = false;
-
-  uint32_t prefix = aCompletion.ToUint32();
-
-  bool found;
-  nsresult rv = mPrefixSet->Contains(prefix, &found);
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  LOG(("Probe in %s: %X, found %d", mTableName.get(), prefix, found));
-
-  if (found) {
-    *aHas = true;
-  }
-
-  // TODO: We may need to distinguish completions found in cache or update in the future
-  if ((mGetHashCache.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex) ||
-      (mUpdateCompletions.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex)) {
-    LOG(("Complete in %s", mTableName.get()));
-    *aComplete = true;
-    *aHas = true;
-  }
-
-  return NS_OK;
-}
-
-nsresult
 LookupCache::WriteFile()
 {
   nsCOMPtr<nsIFile> psFile;
   nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
   NS_ENSURE_SUCCESS(rv, rv);
 
   rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
   NS_ENSURE_SUCCESS(rv, rv);
 
-  rv = mPrefixSet->StoreToFile(psFile);
+  rv = StoreToFile(psFile);
   NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "failed to store the prefixset");
 
   return NS_OK;
 }
 
 void
 LookupCache::ClearAll()
 {
   ClearCache();
-  ClearUpdatedCompletions();
-  mPrefixSet->SetPrefixes(nullptr, 0);
+  ClearPrefixes();
   mPrimed = false;
 }
 
 void
-LookupCache::ClearUpdatedCompletions()
-{
-  mUpdateCompletions.Clear();
-}
-
-void
 LookupCache::ClearCache()
 {
   mGetHashCache.Clear();
 }
 
-nsresult
-LookupCache::ReadCompletions()
-{
-  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++) {
-    mUpdateCompletions.AppendElement(addComplete[i].complete);
-  }
-
-  return NS_OK;
-}
-
 /* static */ bool
 LookupCache::IsCanonicalizedIP(const nsACString& aHost)
 {
   // The canonicalization process will have left IP addresses in dotted
   // decimal with no surprises.
   uint32_t i1, i2, i3, i4;
   char c;
   if (PR_sscanf(PromiseFlatCString(aHost).get(), "%u.%u.%u.%u%c",
@@ -447,19 +339,188 @@ LookupCache::GetHostKeys(const nsACStrin
     lookupHost2->Assign(hostComponents[last - 2]);
     lookupHost2->Append(".");
     lookupHost2->Append(*lookupHost);
   }
 
   return NS_OK;
 }
 
-bool LookupCache::IsPrimed()
+nsresult
+LookupCache::LoadPrefixSet()
+{
+  nsCOMPtr<nsIFile> psFile;
+  nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  bool exists;
+  rv = psFile->Exists(&exists);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  if (exists) {
+    LOG(("stored PrefixSet exists, loading from disk"));
+    rv = LoadFromFile(psFile);
+    if (NS_FAILED(rv)) {
+      if (rv == NS_ERROR_FILE_CORRUPTED) {
+        Reset();
+      }
+      return rv;
+    }
+    mPrimed = true;
+  } else {
+    LOG(("no (usable) stored PrefixSet found"));
+  }
+
+#ifdef DEBUG
+  if (mPrimed) {
+    uint32_t size = SizeOfPrefixSet();
+    LOG(("SB tree done, size = %d bytes\n", size));
+  }
+#endif
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::Init()
+{
+  mPrefixSet = new nsUrlClassifierPrefixSet();
+  nsresult rv = mPrefixSet->Init(mTableName);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::Open()
+{
+  nsresult rv = LookupCache::Open();
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  LOG(("Reading Completions"));
+  rv = ReadCompletions();
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  return NS_OK;
+}
+
+void
+LookupCacheV2::ClearAll()
+{
+  LookupCache::ClearAll();
+  mUpdateCompletions.Clear();
+}
+
+nsresult
+LookupCacheV2::Has(const Completion& aCompletion,
+                   bool* aHas, bool* aComplete)
 {
-  return mPrimed;
+  *aHas = *aComplete = false;
+
+  uint32_t prefix = aCompletion.ToUint32();
+
+  bool found;
+  nsresult rv = mPrefixSet->Contains(prefix, &found);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  LOG(("Probe in %s: %X, found %d", mTableName.get(), prefix, found));
+
+  if (found) {
+    *aHas = true;
+  }
+
+  if ((mGetHashCache.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex) ||
+      (mUpdateCompletions.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex)) {
+    LOG(("Complete in %s", mTableName.get()));
+    *aComplete = true;
+    *aHas = true;
+  }
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::Build(AddPrefixArray& aAddPrefixes,
+                     AddCompleteArray& aAddCompletes)
+{
+  Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_COMPLETIONS,
+                        static_cast<uint32_t>(aAddCompletes.Length()));
+
+  mUpdateCompletions.Clear();
+  mUpdateCompletions.SetCapacity(aAddCompletes.Length());
+  for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
+    mUpdateCompletions.AppendElement(aAddCompletes[i].CompleteHash());
+  }
+  aAddCompletes.Clear();
+  mUpdateCompletions.Sort();
+
+  Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_PREFIXES,
+                        static_cast<uint32_t>(aAddPrefixes.Length()));
+
+  nsresult rv = ConstructPrefixSet(aAddPrefixes);
+  NS_ENSURE_SUCCESS(rv, rv);
+  mPrimed = true;
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes)
+{
+  if (!mPrimed) {
+    // This can happen if its a new table, so no error.
+    LOG(("GetPrefixes from empty LookupCache"));
+    return NS_OK;
+  }
+  return mPrefixSet->GetPrefixesNative(aAddPrefixes);
+}
+
+nsresult
+LookupCacheV2::ReadCompletions()
+{
+  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++) {
+    mUpdateCompletions.AppendElement(addComplete[i].complete);
+  }
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::ClearPrefixes()
+{
+  return mPrefixSet->SetPrefixes(nullptr, 0);
+}
+
+nsresult
+LookupCacheV2::StoreToFile(nsIFile* aFile)
+{
+  return mPrefixSet->StoreToFile(aFile);
+}
+
+nsresult
+LookupCacheV2::LoadFromFile(nsIFile* aFile)
+{
+  return mPrefixSet->LoadFromFile(aFile);
+}
+
+size_t
+LookupCacheV2::SizeOfPrefixSet()
+{
+  return mPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
 }
 
 #ifdef DEBUG
 template <class T>
 static void EnsureSorted(T* aArray)
 {
   typename T::elem_type* start = aArray->Elements();
   typename T::elem_type* end = aArray->Elements() + aArray->Length();
@@ -473,17 +534,17 @@ static void EnsureSorted(T* aArray)
       MOZ_ASSERT(*previous <= *iter);
     }
   }
   return;
 }
 #endif
 
 nsresult
-LookupCache::ConstructPrefixSet(AddPrefixArray& aAddPrefixes)
+LookupCacheV2::ConstructPrefixSet(AddPrefixArray& aAddPrefixes)
 {
   Telemetry::AutoTimer<Telemetry::URLCLASSIFIER_PS_CONSTRUCT_TIME> timer;
 
   nsTArray<uint32_t> array;
   if (!array.SetCapacity(aAddPrefixes.Length(), fallible)) {
     return NS_ERROR_OUT_OF_MEMORY;
   }
 
@@ -507,60 +568,76 @@ LookupCache::ConstructPrefixSet(AddPrefi
   LOG(("SB tree done, size = %d bytes\n", size));
 #endif
 
   mPrimed = true;
 
   return NS_OK;
 }
 
-nsresult
-LookupCache::LoadPrefixSet()
+#if defined(DEBUG)
+void
+LookupCacheV2::DumpCompletions()
 {
-  nsCOMPtr<nsIFile> psFile;
-  nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  bool exists;
-  rv = psFile->Exists(&exists);
-  NS_ENSURE_SUCCESS(rv, rv);
+  if (!LOG_ENABLED())
+    return;
 
-  if (exists) {
-    LOG(("stored PrefixSet exists, loading from disk"));
-    rv = mPrefixSet->LoadFromFile(psFile);
-    if (NS_FAILED(rv)) {
-      if (rv == NS_ERROR_FILE_CORRUPTED) {
-        Reset();
-      }
-      return rv;
-    }
-    mPrimed = true;
-  } else {
-    LOG(("no (usable) stored PrefixSet found"));
+  for (uint32_t i = 0; i < mUpdateCompletions.Length(); i++) {
+    nsAutoCString str;
+    mUpdateCompletions[i].ToHexString(str);
+    LOG(("Update: %s", str.get()));
   }
+}
+#endif
 
-#ifdef DEBUG
-  if (mPrimed) {
-    uint32_t size = mPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
-    LOG(("SB tree done, size = %d bytes\n", size));
-  }
-#endif
+nsresult
+LookupCacheV4::Init()
+{
+  mVLPrefixSet = new VariableLengthPrefixSet();
+  nsresult rv = mVLPrefixSet->Init(mTableName);
+  NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
+// TODO : Bug 1298257, Implement url matching for variable-length prefix set
 nsresult
-LookupCache::GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes)
+LookupCacheV4::Has(const Completion& aCompletion,
+                   bool* aHas, bool* aComplete)
 {
-  if (!mPrimed) {
-    // This can happen if its a new table, so no error.
-    LOG(("GetPrefixes from empty LookupCache"));
-    return NS_OK;
-  }
-  return mPrefixSet->GetPrefixesNative(aAddPrefixes);
+  *aHas = false;
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV4::Build(PrefixStringMap& aPrefixMap)
+{
+  return mVLPrefixSet->SetPrefixes(aPrefixMap);
 }
 
+nsresult
+LookupCacheV4::ClearPrefixes()
+{
+  // Clear by seting a empty map
+  PrefixStringMap map;
+  return mVLPrefixSet->SetPrefixes(map);
+}
+
+nsresult
+LookupCacheV4::StoreToFile(nsIFile* aFile)
+{
+  return mVLPrefixSet->StoreToFile(aFile);
+}
+
+nsresult
+LookupCacheV4::LoadFromFile(nsIFile* aFile)
+{
+  return mVLPrefixSet->LoadFromFile(aFile);
+}
+
+size_t
+LookupCacheV4::SizeOfPrefixSet()
+{
+  return mVLPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
+}
 
 } // namespace safebrowsing
 } // namespace mozilla
--- a/toolkit/components/url-classifier/LookupCache.h
+++ b/toolkit/components/url-classifier/LookupCache.h
@@ -9,16 +9,17 @@
 #include "Entries.h"
 #include "nsString.h"
 #include "nsTArray.h"
 #include "nsCOMPtr.h"
 #include "nsIFile.h"
 #include "nsIFileStreams.h"
 #include "mozilla/RefPtr.h"
 #include "nsUrlClassifierPrefixSet.h"
+#include "VariableLengthPrefixSet.h"
 #include "mozilla/Logging.h"
 
 namespace mozilla {
 namespace safebrowsing {
 
 #define MAX_HOST_COMPONENTS 5
 #define MAX_PATH_COMPONENTS 4
 
@@ -91,63 +92,144 @@ public:
   // two keys will be returned:
   //  hostname.com/foo/bar -> [hostname.com]
   //  mail.hostname.com/foo/bar -> [hostname.com, mail.hostname.com]
   //  www.mail.hostname.com/foo/bar -> [hostname.com, mail.hostname.com]
   static nsresult GetHostKeys(const nsACString& aSpec,
                               nsTArray<nsCString>* aHostKeys);
 
   LookupCache(const nsACString& aTableName, nsIFile* aStoreFile);
-  ~LookupCache();
+  virtual ~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 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();
+
+  // Write data stored in lookup cache to disk.
+  nsresult WriteFile();
+
+  // Clear completions retrieved from gethash request.
   void ClearCache();
 
+  bool IsPrimed() const { return mPrimed; };
+
 #if DEBUG
   void DumpCache();
-  void Dump();
 #endif
-  nsresult WriteFile();
-  nsresult Has(const Completion& aCompletion,
-               bool* aHas, bool* aComplete);
-  bool IsPrimed();
+
+  virtual nsresult Open();
+  virtual nsresult Init() = 0;
+  virtual nsresult ClearPrefixes() = 0;
+  virtual nsresult Has(const Completion& aCompletion,
+                       bool* aHas, bool* aComplete) = 0;
+
+  template<typename T>
+  static T* Cast(LookupCache* aThat) {
+    return (T::VER == aThat->Ver() ? reinterpret_cast<T*>(aThat) : nullptr);
+  }
 
 private:
-  void ClearAll();
   nsresult Reset();
-  nsresult ReadCompletions();
   nsresult LoadPrefixSet();
-  nsresult LoadCompletions();
-  // Construct a Prefix Set with known prefixes.
-  // This will Clear() aAddPrefixes when done.
-  nsresult ConstructPrefixSet(AddPrefixArray& aAddPrefixes);
+
+  virtual nsresult StoreToFile(nsIFile* aFile) = 0;
+  virtual nsresult LoadFromFile(nsIFile* aFile) = 0;
+  virtual size_t SizeOfPrefixSet() = 0;
+
+  virtual int Ver() const = 0;
+
+protected:
+  virtual void ClearAll();
 
   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;
 
   // For gtest to inspect private members.
   friend class PerProviderDirectoryTestUtils;
 };
 
+class LookupCacheV2 final : public LookupCache
+{
+public:
+  explicit LookupCacheV2(const nsACString& aTableName, nsIFile* aStoreFile)
+    : LookupCache(aTableName, aStoreFile) {}
+  ~LookupCacheV2() {}
+
+  virtual nsresult Init() override;
+  virtual nsresult Open() override;
+  virtual void ClearAll() override;
+  virtual nsresult Has(const Completion& aCompletion,
+                       bool* aHas, bool* aComplete) override;
+
+  nsresult Build(AddPrefixArray& aAddPrefixes,
+                 AddCompleteArray& aAddCompletes);
+
+  nsresult GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes);
+
+#if DEBUG
+  void DumpCompletions();
+#endif
+
+  static const int VER;
+
+protected:
+  nsresult ReadCompletions();
+
+  virtual nsresult ClearPrefixes() override;
+  virtual nsresult StoreToFile(nsIFile* aFile) override;
+  virtual nsresult LoadFromFile(nsIFile* aFile) override;
+  virtual size_t SizeOfPrefixSet() override;
+
+private:
+  virtual int Ver() const override { return VER; }
+
+  // Construct a Prefix Set with known prefixes.
+  // This will Clear() aAddPrefixes when done.
+  nsresult ConstructPrefixSet(AddPrefixArray& aAddPrefixes);
+
+  // Full length hashes obtained in update request
+  CompletionArray mUpdateCompletions;
+
+  // Set of prefixes known to be in the database
+  RefPtr<nsUrlClassifierPrefixSet> mPrefixSet;
+};
+
+class LookupCacheV4 final : public LookupCache
+{
+public:
+  explicit LookupCacheV4(const nsACString& aTableName, nsIFile* aStoreFile)
+    : LookupCache(aTableName, aStoreFile) {}
+  ~LookupCacheV4() {}
+
+  virtual nsresult Init() override;
+  virtual nsresult Has(const Completion& aCompletion,
+                       bool* aHas, bool* aComplete) override;
+
+  nsresult Build(PrefixStringMap& aPrefixMap);
+
+  static const int VER;
+
+protected:
+  virtual nsresult ClearPrefixes() override;
+  virtual nsresult StoreToFile(nsIFile* aFile) override;
+  virtual nsresult LoadFromFile(nsIFile* aFile) override;
+  virtual size_t SizeOfPrefixSet() override;
+
+private:
+  virtual int Ver() const override { return VER; }
+
+  RefPtr<VariableLengthPrefixSet> mVLPrefixSet;
+};
+
 } // namespace safebrowsing
 } // namespace mozilla
 
 #endif
--- a/toolkit/components/url-classifier/ProtocolParser.cpp
+++ b/toolkit/components/url-classifier/ProtocolParser.cpp
@@ -868,16 +868,18 @@ ProtocolParserProtobuf::ProcessOneRespon
     DebugOnly<nsresult> rv = SaveStateToPref(listName, state);
     NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "SaveStateToPref failed");
   }));
 
   PARSER_LOG(("==== Update for threat type '%d' ====", aResponse.threat_type()));
   PARSER_LOG(("* listName: %s\n", listName.get()));
   PARSER_LOG(("* newState: %s\n", aResponse.new_client_state().c_str()));
   PARSER_LOG(("* isFullUpdate: %s\n", (isFullUpdate ? "yes" : "no")));
+
+  tuV4->SetFullUpdate(isFullUpdate);
   ProcessAdditionOrRemoval(*tuV4, aResponse.additions(), true /*aIsAddition*/);
   ProcessAdditionOrRemoval(*tuV4, aResponse.removals(), false);
   PARSER_LOG(("\n\n"));
 
   return NS_OK;
 }
 
 nsresult
--- a/toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
+++ b/toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
@@ -76,22 +76,22 @@ void VerifyPrivateStorePath(const char* 
 TEST(PerProviderDirectory, LookupCache)
 {
   RunTestInNewThread([] () -> void {
     nsCOMPtr<nsIFile> rootDir;
     NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(rootDir));
 
     // For V2 tables (NOT ending with '-proto'), root directory should be
     // used as the private store.
-    VerifyPrivateStorePath<LookupCache>("goog-phish-shavar", "google", rootDir, false);
+    VerifyPrivateStorePath<LookupCacheV2>("goog-phish-shavar", "google", rootDir, false);
 
     // For V4 tables, if provider is found, use per-provider subdirectory;
     // If not found, use root directory.
-    VerifyPrivateStorePath<LookupCache>("goog-noprovider-proto", "", rootDir, false);
-    VerifyPrivateStorePath<LookupCache>("goog-phish-proto", "google4", rootDir, true);
+    VerifyPrivateStorePath<LookupCacheV4>("goog-noprovider-proto", "", rootDir, false);
+    VerifyPrivateStorePath<LookupCacheV4>("goog-phish-proto", "google4", rootDir, true);
   });
 }
 
 TEST(PerProviderDirectory, HashStore)
 {
   RunTestInNewThread([] () -> void {
     nsCOMPtr<nsIFile> rootDir;
     NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(rootDir));
new file mode 100644
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
@@ -0,0 +1,190 @@
+#include "Classifier.h"
+#include "HashStore.h"
+#include "nsAppDirectoryServiceDefs.h"
+#include "nsIFile.h"
+#include "nsIThread.h"
+#include "string.h"
+#include "gtest/gtest.h"
+
+using namespace mozilla;
+using namespace mozilla::safebrowsing;
+
+typedef nsCString _Prefix;
+typedef nsTArray<_Prefix> _PrefixArray;
+
+static void
+PrefixArrayToPrefixStringMap(const _PrefixArray& prefixArray,
+                             PrefixStringMap& outMap)
+{
+  outMap.Clear();
+
+  for (uint32_t i = 0; i < prefixArray.Length(); i++) {
+    const _Prefix& prefix = prefixArray[i];
+    nsCString* prefixString = outMap.LookupOrAdd(prefix.Length());
+    prefixString->Append(prefix.BeginReading(), prefix.Length());
+  }
+}
+
+// N: Number of prefixes, MIN/MAX: minimum/maximum prefix size
+// This function will append generated prefixes to outArray.
+// All elements in the generated array will be different.
+static void
+CreateRandomSortedPrefixArray(uint32_t N,
+                              uint32_t MIN,
+                              uint32_t MAX,
+                              _PrefixArray& outArray)
+{
+  outArray.SetCapacity(outArray.Length() + N);
+
+  const uint32_t range = (MAX - MIN + 1);
+
+  for (uint32_t i = 0; i < N; i++) {
+    uint32_t prefixSize = (rand() % range) + MIN;
+    _Prefix prefix;
+    prefix.SetLength(prefixSize);
+
+    while (true) {
+      char* dst = prefix.BeginWriting();
+      for (uint32_t j = 0; j < prefixSize; j++) {
+        dst[j] = rand() % 256;
+      }
+
+      if (!outArray.Contains(prefix)) {
+        outArray.AppendElement(prefix);
+        break;
+      }
+    }
+  }
+
+  outArray.Sort();
+}
+
+// Function to generate TableUpdateV4.
+static void
+GenerateUpdateData(bool fullUpdate,
+                   PrefixStringMap& add,
+                   nsTArray<uint32_t>& removal,
+                   nsTArray<TableUpdate*>& tableUpdates)
+{
+  TableUpdateV4* tableUpdate = new TableUpdateV4(NS_LITERAL_CSTRING("gtest-malware-proto"));
+  tableUpdate->SetFullUpdate(fullUpdate);
+
+  for (auto iter = add.ConstIter(); !iter.Done(); iter.Next()) {
+    nsCString* pstring = iter.Data();
+    std::string str(pstring->BeginReading(), pstring->Length());
+
+    tableUpdate->NewPrefixes(iter.Key(), str);
+  }
+
+  tableUpdate->NewRemovalIndices(removal.Elements(), removal.Length());
+
+  tableUpdates.AppendElement(tableUpdate);
+}
+
+static void
+VerifyPrefixSet(PrefixStringMap& expected)
+{
+  // Verify the prefix set is written to disk.
+  nsCOMPtr<nsIFile> file;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                         getter_AddRefs(file));
+
+  file->AppendRelativeNativePath(NS_LITERAL_CSTRING("safebrowsing/gtest-malware-proto.pset"));
+
+  RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;
+  load->Init(NS_LITERAL_CSTRING("gtest-malware-proto"));
+
+  PrefixStringMap out;
+  load->LoadFromFile(file);
+  load->GetPrefixes(out);
+
+  for (auto iter = expected.ConstIter(); !iter.Done(); iter.Next()) {
+    nsCString* str1 = iter.Data();
+    nsCString* str2 = out.Get(iter.Key());
+
+    ASSERT_TRUE(*str1 == *str2);
+  }
+}
+
+static void
+Clear()
+{
+  nsCOMPtr<nsIFile> file;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                         getter_AddRefs(file));
+  file->AppendRelativeNativePath(NS_LITERAL_CSTRING("safebrowsing/gtest-malware-proto.pset"));
+  file->Remove(false);
+}
+
+static void
+testUpdate(nsTArray<TableUpdate*>& tableUpdates,
+           PrefixStringMap& expected)
+{
+  nsCOMPtr<nsIFile> file;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                         getter_AddRefs(file));
+
+  UniquePtr<Classifier> classifier(new Classifier());
+  classifier->Open(*file);
+
+  RunTestInNewThread([&] () -> void {
+    nsresult rv = classifier->ApplyUpdates(&tableUpdates);
+    ASSERT_TRUE(rv == NS_OK);
+
+    VerifyPrefixSet(expected);
+  });
+}
+
+static void
+testFullUpdate(PrefixStringMap& add)
+{
+  nsTArray<uint32_t> empty;
+  nsTArray<TableUpdate*> tableUpdates;
+  GenerateUpdateData(true, add, empty, tableUpdates);
+
+  testUpdate(tableUpdates, add);
+}
+
+TEST(UrlClassifierTableUpdateV4, FixLenghtPSetFullUpdate)
+{
+  srand(time(NULL));
+
+   _PrefixArray array;
+  PrefixStringMap map;
+
+  CreateRandomSortedPrefixArray(5000, 4, 4, array);
+  PrefixArrayToPrefixStringMap(array, map);
+
+  testFullUpdate(map);
+
+  Clear();
+}
+
+
+TEST(UrlClassifierTableUpdateV4, VariableLenghtPSetFullUpdate)
+{
+   _PrefixArray array;
+  PrefixStringMap map;
+
+  CreateRandomSortedPrefixArray(5000, 5, 32, array);
+  PrefixArrayToPrefixStringMap(array, map);
+
+  testFullUpdate(map);
+
+  Clear();
+}
+
+// This test contain both variable length prefix set and fixed-length prefix set
+TEST(UrlClassifierTableUpdateV4, MixedPSetFullUpdate)
+{
+   _PrefixArray array;
+  PrefixStringMap map;
+
+  CreateRandomSortedPrefixArray(5000, 4, 4, array);
+  CreateRandomSortedPrefixArray(1000, 5, 32, array);
+  PrefixArrayToPrefixStringMap(array, map);
+
+  testFullUpdate(map);
+
+  Clear();
+}
--- a/toolkit/components/url-classifier/tests/gtest/moz.build
+++ b/toolkit/components/url-classifier/tests/gtest/moz.build
@@ -9,13 +9,14 @@ LOCAL_INCLUDES += [
 ]
 
 UNIFIED_SOURCES += [
     'TestChunkSet.cpp',
     'TestPerProviderDirectory.cpp',
     'TestSafebrowsingHash.cpp',
     'TestSafeBrowsingProtobuf.cpp',
     'TestTable.cpp',
+    'TestUrlClassifierTableUpdateV4.cpp',
     'TestUrlClassifierUtils.cpp',
     'TestVariableLengthPrefixSet.cpp',
 ]
 
 FINAL_LIBRARY = 'xul-gtest'