Bug 1305801 - Part 5: Support SafeBrowsing v4 partial update. r=gcp
authorDimi Lee <dlee@mozilla.com>
Tue, 04 Oct 2016 09:14:39 +0800
changeset 316425 2d4f6c10aacfb9a319624a9df50d7b0769cfafa9
parent 316424 122db2c234f18b7e9ae8e6487ab673874d887a50
child 316426 3470e326025c62381dc5f7c06629dbe5dbd7f242
push id30770
push userkwierso@gmail.com
push dateWed, 05 Oct 2016 00:00:48 +0000
treeherdermozilla-central@3470e326025c [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersgcp
bugs1305801
milestone52.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1305801 - Part 5: Support SafeBrowsing v4 partial update. r=gcp MozReview-Commit-ID: 7OEWLaZbotS
toolkit/components/telemetry/Histograms.json
toolkit/components/url-classifier/Classifier.cpp
toolkit/components/url-classifier/HashStore.cpp
toolkit/components/url-classifier/HashStore.h
toolkit/components/url-classifier/LookupCache.cpp
toolkit/components/url-classifier/LookupCache.h
toolkit/components/url-classifier/LookupCacheV4.cpp
toolkit/components/url-classifier/LookupCacheV4.h
toolkit/components/url-classifier/moz.build
toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
--- a/toolkit/components/telemetry/Histograms.json
+++ b/toolkit/components/telemetry/Histograms.json
@@ -3816,16 +3816,24 @@
   },
   "URLCLASSIFIER_COMPLETE_TIMEOUT": {
     "alert_emails": ["safebrowsing-telemetry@mozilla.org"],
     "expires_in_version": "56",
     "kind": "boolean",
     "bug_numbers": [1172688],
     "description": "This metric is recorded every time a gethash lookup is performed, `true` is recorded if the lookup times out."
   },
+  "URLCLASSIFIER_UPDATE_ERROR_TYPE": {
+    "alert_emails": ["safebrowsing-telemetry@mozilla.org"],
+    "expires_in_version": "58",
+    "kind": "enumerated",
+    "n_values": 10,
+    "bug_numbers": [1305801],
+    "description": "An error was encountered while parsing a partial update returned by a Safe Browsing V4 server (0 = addition of an already existing prefix, 1 = parser got into an infinite loop, 2 = removal index out of bounds)"
+  },
   "CSP_DOCUMENTS_COUNT": {
     "alert_emails": ["seceng@mozilla.com"],
     "bug_numbers": [1252829],
     "expires_in_version": "55",
     "kind": "count",
     "description": "Number of unique pages that contain a CSP"
   },
   "CSP_UNSAFE_INLINE_DOCUMENTS_COUNT": {
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -1,14 +1,15 @@
 //* -*- 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/. */
 
 #include "Classifier.h"
+#include "LookupCacheV4.h"
 #include "nsIPrefBranch.h"
 #include "nsIPrefService.h"
 #include "nsISimpleEnumerator.h"
 #include "nsIRandomGenerator.h"
 #include "nsIInputStream.h"
 #include "nsISeekableStream.h"
 #include "nsIFile.h"
 #include "nsNetCID.h"
@@ -941,44 +942,72 @@ Classifier::UpdateTableV4(nsTArray<Table
   }
 
   LookupCacheV4* lookupCache =
     LookupCache::Cast<LookupCacheV4>(GetLookupCache(aTable));
   if (!lookupCache) {
     return NS_ERROR_FAILURE;
   }
 
-  PrefixStringMap prefixes;
+  nsresult rv = NS_OK;
+
+  // prefixes2 is only used in partial update. If there are multiple
+  // updates for the same table, prefixes1 & prefixes2 will act as
+  // input and output in turn to reduce memory copy overhead.
+  PrefixStringMap prefixes1, prefixes2;
+  PrefixStringMap* output = &prefixes1;
+
   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();
+      TableUpdateV4::PrefixStdStringMap& map = updateV4->Prefixes();
 
+      output->Clear();
       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);
+        output->Put(iter.Key(), prefix);
       }
     } else {
-      // TODO: Bug 1287058, partial update
+      PrefixStringMap* input = nullptr;
+      // If both prefix sets are empty, this means we are doing a partial update
+      // without a prior full/partial update in the loop. In this case we should
+      // get prefixes from the lookup cache first.
+      if (prefixes1.IsEmpty() && prefixes2.IsEmpty()) {
+        lookupCache->GetPrefixes(prefixes1);
+        input = &prefixes1;
+        output = &prefixes2;
+      } else {
+        MOZ_ASSERT(prefixes1.IsEmpty() ^ prefixes2.IsEmpty());
+
+        // When there are multiple partial updates, input should always point
+        // to the non-empty prefix set(filled by previous full/partial update).
+        // output should always point to the empty prefix set.
+        input = prefixes1.IsEmpty() ? &prefixes2 : &prefixes1;
+        output = prefixes1.IsEmpty() ? &prefixes1 : &prefixes2;
+      }
+
+      rv = lookupCache->ApplyPartialUpdate(updateV4, *input, *output);
+      NS_ENSURE_SUCCESS(rv, rv);
+
+      input->Clear();
     }
 
     aUpdates->ElementAt(i) = nullptr;
   }
 
-  nsresult rv = lookupCache->Build(prefixes);
+  rv = lookupCache->Build(*output);
   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);
--- a/toolkit/components/url-classifier/HashStore.cpp
+++ b/toolkit/components/url-classifier/HashStore.cpp
@@ -163,17 +163,17 @@ TableUpdateV2::NewSubComplete(uint32_t a
 }
 
 void
 TableUpdateV4::NewPrefixes(int32_t aSize, std::string& aPrefixes)
 {
   NS_ENSURE_TRUE_VOID(aPrefixes.size() % aSize == 0);
   NS_ENSURE_TRUE_VOID(!mPrefixesMap.Get(aSize));
 
-  PrefixString* prefix = new PrefixString(aPrefixes);
+  PrefixStdString* prefix = new PrefixStdString(aPrefixes);
   mPrefixesMap.Put(aSize, prefix);
 }
 
 void
 TableUpdateV4::NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices)
 {
   for (size_t i = 0; i < aNumOfIndices; i++) {
     mRemovalIndiceArray.AppendElement(aIndices[i]);
--- a/toolkit/components/url-classifier/HashStore.h
+++ b/toolkit/components/url-classifier/HashStore.h
@@ -127,62 +127,62 @@ private:
   virtual int Tag() const override { return TAG; }
 };
 
 // Structure for DBService/HashStore/Classifiers to update.
 // It would contain the prefixes (both fixed and variable length)
 // for addition and indices to removal. See Bug 1283009.
 class TableUpdateV4 : public TableUpdate {
 public:
-  struct PrefixString {
+  struct PrefixStdString {
   private:
     std::string mStorage;
     nsDependentCSubstring mString;
 
   public:
-    explicit PrefixString(std::string& aString)
+    explicit PrefixStdString(std::string& aString)
     {
       aString.swap(mStorage);
       mString.Rebind(mStorage.data(), mStorage.size());
     };
 
     const nsACString& GetPrefixString() const { return mString; };
   };
 
-  typedef nsClassHashtable<nsUint32HashKey, PrefixString> PrefixesStringMap;
+  typedef nsClassHashtable<nsUint32HashKey, PrefixStdString> PrefixStdStringMap;
   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; }
+  PrefixStdStringMap& 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;
+  PrefixStdStringMap mPrefixesMap;
   RemovalIndiceArray mRemovalIndiceArray;
 };
 
 // There is one hash store per table.
 class HashStore {
 public:
   HashStore(const nsACString& aTableName, nsIFile* aRootStoreFile);
   ~HashStore();
--- a/toolkit/components/url-classifier/LookupCache.cpp
+++ b/toolkit/components/url-classifier/LookupCache.cpp
@@ -36,17 +36,16 @@
 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);
 }
@@ -583,61 +582,10 @@ LookupCacheV2::DumpCompletions()
   for (uint32_t i = 0; i < mUpdateCompletions.Length(); i++) {
     nsAutoCString str;
     mUpdateCompletions[i].ToHexString(str);
     LOG(("Update: %s", str.get()));
   }
 }
 #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
-LookupCacheV4::Has(const Completion& aCompletion,
-                   bool* aHas, bool* aComplete)
-{
-  *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
@@ -197,39 +197,12 @@ private:
 
   // 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
new file mode 100644
--- /dev/null
+++ b/toolkit/components/url-classifier/LookupCacheV4.cpp
@@ -0,0 +1,299 @@
+//* -*- 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/. */
+
+#include "LookupCacheV4.h"
+#include "HashStore.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)
+
+namespace mozilla {
+namespace safebrowsing {
+
+const int LookupCacheV4::VER = 4;
+
+// Prefixes coming from updates and VLPrefixSet are both stored in the HashTable
+// where the (key, value) pair is a prefix size and a lexicographic-sorted string.
+// The difference is prefixes from updates use std:string(to avoid additional copies)
+// and prefixes from VLPrefixSet use nsCString.
+// This class provides a common interface for the partial update algorithm to make it
+// easier to operate on two different kind prefix string map..
+class VLPrefixSet
+{
+public:
+  explicit VLPrefixSet(const PrefixStringMap& aMap);
+  explicit VLPrefixSet(const TableUpdateV4::PrefixStdStringMap& aMap);
+
+  // This function will merge the prefix map in VLPrefixSet to aPrefixMap.
+  void Merge(PrefixStringMap& aPrefixMap);
+
+  // Find the smallest string from the map in VLPrefixSet.
+  bool GetSmallestPrefix(nsDependentCSubstring& aOutString);
+
+  // Return the number of prefixes in the map
+  uint32_t Count() const { return mCount; }
+
+private:
+  // PrefixString structure contains a lexicographic-sorted string with
+  // a |pos| variable to indicate which substring we are pointing to right now.
+  // |pos| increases each time GetSmallestPrefix finds the smallest string.
+  struct PrefixString {
+    PrefixString(const nsACString& aStr, uint32_t aSize)
+     : pos(0)
+     , size(aSize)
+    {
+      data.Rebind(aStr.BeginReading(), aStr.Length());
+    }
+
+    const char* get() {
+      return pos < data.Length() ? data.BeginReading() + pos : nullptr;
+    }
+    void next() { pos += size; }
+    uint32_t remaining() { return data.Length() - pos; }
+
+    nsDependentCSubstring data;
+    uint32_t pos;
+    uint32_t size;
+  };
+
+  nsClassHashtable<nsUint32HashKey, PrefixString> mMap;
+  uint32_t mCount;
+};
+
+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
+LookupCacheV4::Has(const Completion& aCompletion,
+                   bool* aHas, bool* aComplete)
+{
+  *aHas = false;
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV4::Build(PrefixStringMap& aPrefixMap)
+{
+  return mVLPrefixSet->SetPrefixes(aPrefixMap);
+}
+
+nsresult
+LookupCacheV4::GetPrefixes(PrefixStringMap& aPrefixMap)
+{
+  return mVLPrefixSet->GetPrefixes(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);
+}
+
+static void
+AppendPrefixToMap(PrefixStringMap& prefixes, nsDependentCSubstring& prefix)
+{
+  if (!prefix.Length()) {
+    return;
+  }
+
+  nsCString* prefixString = prefixes.LookupOrAdd(prefix.Length());
+  prefixString->Append(prefix.BeginReading(), prefix.Length());
+}
+
+// Please see https://bug1287058.bmoattachments.org/attachment.cgi?id=8795366
+// for detail about partial update algorithm.
+nsresult
+LookupCacheV4::ApplyPartialUpdate(TableUpdateV4* aTableUpdate,
+                                  PrefixStringMap& aInputMap,
+                                  PrefixStringMap& aOutputMap)
+{
+  MOZ_ASSERT(aOutputMap.IsEmpty());
+
+  // oldPSet contains prefixes we already have or we just merged last round.
+  // addPSet contains prefixes stored in tableUpdate which should be merged with oldPSet.
+  VLPrefixSet oldPSet(aInputMap);
+  VLPrefixSet addPSet(aTableUpdate->Prefixes());
+
+  // RemovalIndiceArray is a sorted integer array indicating the index of prefix we should
+  // remove from the old prefix set(according to lexigraphic order).
+  // |removalIndex| is the current index of RemovalIndiceArray.
+  // |numOldPrefixPicked| is used to record how many prefixes we picked from the old map.
+  TableUpdateV4::RemovalIndiceArray& removalArray = aTableUpdate->RemovalIndices();
+  uint32_t removalIndex = 0;
+  int32_t numOldPrefixPicked = -1;
+
+  nsDependentCSubstring smallestOldPrefix;
+  nsDependentCSubstring smallestAddPrefix;
+
+  // This is used to avoid infinite loop for partial update algorithm.
+  // The maximum loops will be the number of old prefixes plus the number of add prefixes.
+  uint32_t index = oldPSet.Count() + addPSet.Count() + 1;
+  for(;index > 0; index--) {
+    // Get smallest prefix from the old prefix set if we don't have one
+    if (smallestOldPrefix.IsEmpty()) {
+      // If prefixes from the old prefix set are all merged,
+      // then we can merge the entire add prefix set directly.
+      if (!oldPSet.GetSmallestPrefix(smallestOldPrefix)) {
+        AppendPrefixToMap(aOutputMap, smallestAddPrefix);
+        addPSet.Merge(aOutputMap);
+        break;
+      }
+    }
+
+    // Get smallest prefix from add prefix set if we don't have one
+    if (smallestAddPrefix.IsEmpty()) {
+      // If add prefixes are all merged and there is no removalIndices left,
+      // then merge the entire old prefix set directly. If there are still
+      // removalIndices left, we should still merge prefixes one by one
+      // to know which prefix from old prefix set should be removed.
+      if (!addPSet.GetSmallestPrefix(smallestAddPrefix) &&
+        removalIndex >= removalArray.Length()) {
+        AppendPrefixToMap(aOutputMap, smallestOldPrefix);
+        oldPSet.Merge(aOutputMap);
+        break;
+      }
+    }
+
+    // Compare the smallest string in old prefix set and add prefix set, merge the
+    // smaller one into new map to ensure merged string still follows
+    // lexigraphic order.
+    if (smallestOldPrefix < smallestAddPrefix ||
+        smallestAddPrefix.IsEmpty()) {
+      numOldPrefixPicked++;
+
+      // If the number of picks from old map matches the removalIndex, then this prefix
+      // will be removed by not merging it to new map.
+      if (removalIndex < removalArray.Length() &&
+          numOldPrefixPicked == removalArray[removalIndex]) {
+        removalIndex++;
+      } else {
+        AppendPrefixToMap(aOutputMap, smallestOldPrefix);
+      }
+      smallestOldPrefix.SetLength(0);
+    } else if (smallestOldPrefix > smallestAddPrefix ||
+               smallestOldPrefix.IsEmpty()){
+      AppendPrefixToMap(aOutputMap, smallestAddPrefix);
+      smallestAddPrefix.SetLength(0);
+    } else {
+      NS_WARNING("Add prefix should not exist in the original prefix set.");
+      Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
+                            DUPLICATE_PREFIX);
+      return NS_ERROR_FAILURE;
+    }
+  }
+
+  // We expect index will be greater to 0 because max number of runs will be
+  // the number of original prefix plus add prefix.
+  if (index <= 0) {
+    NS_WARNING("There are still prefixes remaining after reaching maximum runs.");
+    Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
+                          INFINITE_LOOP);
+    return NS_ERROR_FAILURE;
+  }
+
+  if (removalIndex < removalArray.Length()) {
+    NS_WARNING("There are still prefixes to remove after exhausting the old PrefixSet.");
+    Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
+                          WRONG_REMOVAL_INDICES);
+    return NS_ERROR_FAILURE;
+  }
+
+  return NS_OK;
+}
+
+VLPrefixSet::VLPrefixSet(const PrefixStringMap& aMap)
+  : mCount(0)
+{
+  for (auto iter = aMap.ConstIter(); !iter.Done(); iter.Next()) {
+    uint32_t size = iter.Key();
+    mMap.Put(size, new PrefixString(*iter.Data(), size));
+    mCount += iter.Data()->Length() / size;
+  }
+}
+
+VLPrefixSet::VLPrefixSet(const TableUpdateV4::PrefixStdStringMap& aMap)
+  : mCount(0)
+{
+  for (auto iter = aMap.ConstIter(); !iter.Done(); iter.Next()) {
+    uint32_t size = iter.Key();
+    mMap.Put(size, new PrefixString(iter.Data()->GetPrefixString(), size));
+    mCount += iter.Data()->GetPrefixString().Length() / size;
+  }
+}
+
+void
+VLPrefixSet::Merge(PrefixStringMap& aPrefixMap) {
+  for (auto iter = mMap.ConstIter(); !iter.Done(); iter.Next()) {
+    nsCString* prefixString = aPrefixMap.LookupOrAdd(iter.Key());
+    PrefixString* str = iter.Data();
+
+    if (str->get()) {
+      prefixString->Append(str->get(), str->remaining());
+    }
+  }
+}
+
+bool
+VLPrefixSet::GetSmallestPrefix(nsDependentCSubstring& aOutString) {
+  PrefixString* pick = nullptr;
+  for (auto iter = mMap.ConstIter(); !iter.Done(); iter.Next()) {
+    PrefixString* str = iter.Data();
+
+    if (!str->get()) {
+      continue;
+    }
+
+    if (aOutString.IsEmpty()) {
+      aOutString.Rebind(str->get(), iter.Key());
+      pick = str;
+      continue;
+    }
+
+    nsDependentCSubstring cur(str->get(), iter.Key());
+    if (cur < aOutString) {
+      aOutString.Rebind(str->get(), iter.Key());
+      pick = str;
+    }
+  }
+
+  if (pick) {
+    pick->next();
+  }
+
+  return pick != nullptr;
+}
+
+} // namespace safebrowsing
+} // namespace mozilla
new file mode 100644
--- /dev/null
+++ b/toolkit/components/url-classifier/LookupCacheV4.h
@@ -0,0 +1,61 @@
+//* -*- 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 LookupCacheV4_h__
+#define LookupCacheV4_h__
+
+#include "LookupCache.h"
+
+namespace mozilla {
+namespace safebrowsing {
+
+// Forward declaration.
+class TableUpdateV4;
+
+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);
+
+  nsresult GetPrefixes(PrefixStringMap& aPrefixMap);
+
+  // ApplyPartialUpdate will merge partial update data stored in aTableUpdate
+  // with prefixes in aInputMap.
+  nsresult ApplyPartialUpdate(TableUpdateV4* aTableUpdate,
+                              PrefixStringMap& aInputMap,
+                              PrefixStringMap& aOutputMap);
+
+  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; }
+
+  enum UPDATE_ERROR_TYPES {
+    DUPLICATE_PREFIX = 0,
+    INFINITE_LOOP = 1,
+    WRONG_REMOVAL_INDICES = 2,
+  };
+
+  RefPtr<VariableLengthPrefixSet> mVLPrefixSet;
+};
+
+} // namespace safebrowsing
+} // namespace mozilla
+
+#endif
--- a/toolkit/components/url-classifier/moz.build
+++ b/toolkit/components/url-classifier/moz.build
@@ -19,16 +19,17 @@ XPIDL_MODULE = 'url-classifier'
 
 # Disable RTTI in google protocol buffer
 DEFINES['GOOGLE_PROTOBUF_NO_RTTI'] = True
 
 UNIFIED_SOURCES += [
     'ChunkSet.cpp',
     'Classifier.cpp',
     'LookupCache.cpp',
+    'LookupCacheV4.cpp',
     'nsCheckSummedOutputStream.cpp',
     'nsUrlClassifierDBService.cpp',
     'nsUrlClassifierProxies.cpp',
     'nsUrlClassifierUtils.cpp',
     'protobuf/safebrowsing.pb.cc',
     'ProtocolParser.cpp',
 ]
 
@@ -57,16 +58,17 @@ EXTRA_PP_COMPONENTS += [
 
 EXTRA_JS_MODULES += [
     'SafeBrowsing.jsm',
 ]
 
 EXPORTS += [
     'Entries.h',
     'LookupCache.h',
+    'LookupCacheV4.h',
     'nsUrlClassifierPrefixSet.h',
     'protobuf/safebrowsing.pb.h',
     'VariableLengthPrefixSet.h',
 ]
 
 FINAL_LIBRARY = 'xul'
 
 LOCAL_INCLUDES += [
--- a/toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
+++ b/toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
@@ -1,9 +1,10 @@
 #include "LookupCache.h"
+#include "LookupCacheV4.h"
 #include "HashStore.h"
 #include "gtest/gtest.h"
 #include "nsIThread.h"
 #include "nsAppDirectoryServiceDefs.h"
 #include "nsThreadUtils.h"
 
 namespace mozilla {
 namespace safebrowsing {
--- a/toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
+++ b/toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
@@ -1,38 +1,64 @@
 #include "Classifier.h"
 #include "HashStore.h"
 #include "nsAppDirectoryServiceDefs.h"
 #include "nsIFile.h"
 #include "nsIThread.h"
 #include "string.h"
 #include "gtest/gtest.h"
+#include "nsThreadUtils.h"
 
 using namespace mozilla;
 using namespace mozilla::safebrowsing;
 
 typedef nsCString _Prefix;
 typedef nsTArray<_Prefix> _PrefixArray;
 
+// This function removes common elements of inArray and outArray from
+// outArray. This is used by partial update testcase to ensure partial update
+// data won't contain prefixes we already have.
+static void
+RemoveIntersection(const _PrefixArray& inArray, _PrefixArray& outArray)
+{
+  for (uint32_t i = 0; i < inArray.Length(); i++) {
+    int32_t idx = outArray.BinaryIndexOf(inArray[i]);
+    if (idx >= 0) {
+      outArray.RemoveElementAt(idx);
+    }
+  }
+}
+
+// This fucntion removes elements from outArray by index specified in
+// removal array.
+static void
+RemoveElements(const nsTArray<uint32_t>& removal, _PrefixArray& outArray)
+{
+  for (int32_t i = removal.Length() - 1; i >= 0; i--) {
+    outArray.RemoveElementAt(removal[i]);
+  }
+}
+
+// This function converts lexigraphic-sorted prefixes to a hashtable
+// which key is prefix size and value is concatenated prefix string.
 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);
 
@@ -54,16 +80,30 @@ CreateRandomSortedPrefixArray(uint32_t N
         break;
       }
     }
   }
 
   outArray.Sort();
 }
 
+// N: Number of removal indices, MAX: maximum index
+static void
+CreateRandomRemovalIndices(uint32_t N,
+                           uint32_t MAX,
+                           nsTArray<uint32_t>& outArray)
+{
+  for (uint32_t i = 0; i < N; i++) {
+    uint32_t idx = rand() % MAX;
+    if (!outArray.Contains(idx)) {
+      outArray.InsertElementSorted(idx);
+    }
+  }
+}
+
 // 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"));
@@ -84,41 +124,60 @@ GenerateUpdateData(bool fullUpdate,
 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"));
+  file->AppendNative(NS_LITERAL_CSTRING("safebrowsing"));
+  file->AppendNative(NS_LITERAL_CSTRING("gtest-malware-proto.pset"));
 
   RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;
   load->Init(NS_LITERAL_CSTRING("gtest-malware-proto"));
 
-  PrefixStringMap out;
+  PrefixStringMap prefixesInFile;
   load->LoadFromFile(file);
-  load->GetPrefixes(out);
+  load->GetPrefixes(prefixesInFile);
 
   for (auto iter = expected.ConstIter(); !iter.Done(); iter.Next()) {
-    nsCString* str1 = iter.Data();
-    nsCString* str2 = out.Get(iter.Key());
+    nsCString* expectedPrefix = iter.Data();
+    nsCString* resultPrefix = prefixesInFile.Get(iter.Key());
 
-    ASSERT_TRUE(*str1 == *str2);
+    ASSERT_TRUE(*resultPrefix == *expectedPrefix);
   }
 }
 
 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);
+
+  UniquePtr<Classifier> classifier(new Classifier());
+  classifier->Open(*file);
+  classifier->Reset();
+}
+
+static void
+testUpdateFail(nsTArray<TableUpdate*>& tableUpdates)
+{
+  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(NS_FAILED(rv));
+  });
 }
 
 static void
 testUpdate(nsTArray<TableUpdate*>& tableUpdates,
            PrefixStringMap& expected)
 {
   nsCOMPtr<nsIFile> file;
   NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
@@ -140,16 +199,28 @@ testFullUpdate(PrefixStringMap& add)
 {
   nsTArray<uint32_t> empty;
   nsTArray<TableUpdate*> tableUpdates;
   GenerateUpdateData(true, add, empty, tableUpdates);
 
   testUpdate(tableUpdates, add);
 }
 
+static void
+testPartialUpdate(PrefixStringMap& add,
+                  nsTArray<uint32_t>& removal,
+                  PrefixStringMap& expected)
+{
+  nsTArray<TableUpdate*> tableUpdates;
+  GenerateUpdateData(false, add, removal, tableUpdates);
+
+  testUpdate(tableUpdates, expected);
+}
+
+
 TEST(UrlClassifierTableUpdateV4, FixLenghtPSetFullUpdate)
 {
   srand(time(NULL));
 
    _PrefixArray array;
   PrefixStringMap map;
 
   CreateRandomSortedPrefixArray(5000, 4, 4, array);
@@ -183,8 +254,296 @@ TEST(UrlClassifierTableUpdateV4, MixedPS
   CreateRandomSortedPrefixArray(5000, 4, 4, array);
   CreateRandomSortedPrefixArray(1000, 5, 32, array);
   PrefixArrayToPrefixStringMap(array, map);
 
   testFullUpdate(map);
 
   Clear();
 }
+
+TEST(UrlClassifierTableUpdateV4, PartialUpdateWithRemoval)
+{
+  _PrefixArray fArray, pArray, mergedArray;
+  PrefixStringMap fMap, pMap, mergedMap;
+
+  {
+    CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
+    CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
+    PrefixArrayToPrefixStringMap(fArray, fMap);
+
+    testFullUpdate(fMap);
+  }
+
+  {
+    CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
+    CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
+    RemoveIntersection(fArray, pArray);
+    PrefixArrayToPrefixStringMap(pArray, pMap);
+
+    // Remove 1/5 of elements of original prefix set.
+    nsTArray<uint32_t> removal;
+    CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
+    RemoveElements(removal, fArray);
+
+    // Calculate the expected prefix map.
+    mergedArray.AppendElements(fArray);
+    mergedArray.AppendElements(pArray);
+    mergedArray.Sort();
+    PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
+
+    testPartialUpdate(pMap, removal, mergedMap);
+  }
+
+  Clear();
+}
+
+TEST(UrlClassifierTableUpdateV4, PartialUpdateWithoutRemoval)
+{
+  _PrefixArray fArray, pArray, mergedArray;
+  PrefixStringMap fMap, pMap, mergedMap;
+
+  {
+    CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
+    CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
+    PrefixArrayToPrefixStringMap(fArray, fMap);
+
+    testFullUpdate(fMap);
+  }
+
+  {
+    nsTArray<uint32_t> empty;
+
+    CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
+    CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
+    RemoveIntersection(fArray, pArray);
+    PrefixArrayToPrefixStringMap(pArray, pMap);
+
+    // Calculate the expected prefix map.
+    mergedArray.AppendElements(fArray);
+    mergedArray.AppendElements(pArray);
+    mergedArray.Sort();
+    PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
+
+    testPartialUpdate(pMap, empty, mergedMap);
+  }
+
+  Clear();
+}
+
+// Expect failure because partial update contains prefix already
+// in old prefix set.
+TEST(UrlClassifierTableUpdateV4, PartialUpdatePrefixAlreadyExist)
+{
+  _PrefixArray fArray, pArray;
+  PrefixStringMap fMap, pMap;
+
+  {
+    CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
+    PrefixArrayToPrefixStringMap(fArray, fMap);
+
+    testFullUpdate(fMap);
+  }
+
+  {
+    nsTArray<uint32_t> empty;
+    nsTArray<TableUpdate*> tableUpdates;
+
+    // Pick one prefix from full update prefix and add it to partial update.
+    // This should result a failure when call ApplyUpdates.
+    pArray.AppendElement(fArray[rand() % fArray.Length()]);
+    CreateRandomSortedPrefixArray(200, 4, 32, pArray);
+    PrefixArrayToPrefixStringMap(pArray, pMap);
+
+    GenerateUpdateData(false, pMap, empty, tableUpdates);
+    testUpdateFail(tableUpdates);
+  }
+
+  Clear();
+}
+
+// Test apply partial update directly without applying an full update first.
+TEST(UrlClassifierTableUpdateV4, OnlyPartialUpdate)
+{
+  _PrefixArray pArray;
+  PrefixStringMap pMap;
+  nsTArray<uint32_t> empty;
+
+  CreateRandomSortedPrefixArray(5000, 4, 4, pArray);
+  CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
+  PrefixArrayToPrefixStringMap(pArray, pMap);
+
+  testPartialUpdate(pMap, empty, pMap);
+
+  Clear();
+}
+
+// Test partial update without any ADD prefixes, only removalIndices.
+TEST(UrlClassifierTableUpdateV4, PartialUpdateOnlyRemoval)
+{
+  _PrefixArray fArray, pArray;
+  PrefixStringMap fMap, pMap, mergedMap;
+
+  {
+    CreateRandomSortedPrefixArray(5000, 4, 4, fArray);
+    CreateRandomSortedPrefixArray(1000, 5, 32, fArray);
+    PrefixArrayToPrefixStringMap(fArray, fMap);
+
+    testFullUpdate(fMap);
+  }
+
+  {
+    // Remove 1/5 of elements of original prefix set.
+    nsTArray<uint32_t> removal;
+    CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
+    RemoveElements(removal, fArray);
+
+    PrefixArrayToPrefixStringMap(fArray, mergedMap);
+
+    testPartialUpdate(pMap, removal, mergedMap);
+  }
+
+  Clear();
+}
+
+// Test one tableupdate array contains full update and multiple partial updates.
+TEST(UrlClassifierTableUpdateV4, MultipleTableUpdates)
+{
+  _PrefixArray fArray, pArray, mergedArray;
+  PrefixStringMap fMap, pMap, mergedMap;
+
+  {
+    nsTArray<uint32_t> empty;
+    nsTArray<TableUpdate*> tableUpdates;
+
+    // Generate first full udpate
+    CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
+    CreateRandomSortedPrefixArray(2000, 5, 32, fArray);
+    PrefixArrayToPrefixStringMap(fArray, fMap);
+
+    GenerateUpdateData(true, fMap, empty, tableUpdates);
+
+    // Generate second partial update
+    CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
+    CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
+    RemoveIntersection(fArray, pArray);
+    PrefixArrayToPrefixStringMap(pArray, pMap);
+
+    GenerateUpdateData(false, pMap, empty, tableUpdates);
+
+    // Generate thrid partial update
+    fArray.AppendElements(pArray);
+    fArray.Sort();
+    pArray.Clear();
+    CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
+    CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
+    RemoveIntersection(fArray, pArray);
+    PrefixArrayToPrefixStringMap(pArray, pMap);
+
+    // Remove 1/5 of elements of original prefix set.
+    nsTArray<uint32_t> removal;
+    CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
+    RemoveElements(removal, fArray);
+
+    GenerateUpdateData(false, pMap, removal, tableUpdates);
+
+    mergedArray.AppendElements(fArray);
+    mergedArray.AppendElements(pArray);
+    mergedArray.Sort();
+    PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
+
+    testUpdate(tableUpdates, mergedMap);
+  }
+
+  Clear();
+}
+
+// Test apply full update first, and then apply multiple partial updates
+// in one tableupdate array.
+TEST(UrlClassifierTableUpdateV4, MultiplePartialUpdateTableUpdates)
+{
+  _PrefixArray fArray, pArray, mergedArray;
+  PrefixStringMap fMap, pMap, mergedMap;
+
+  {
+    // Generate first full udpate
+    CreateRandomSortedPrefixArray(10000, 4, 4, fArray);
+    CreateRandomSortedPrefixArray(3000, 5, 32, fArray);
+    PrefixArrayToPrefixStringMap(fArray, fMap);
+
+    testFullUpdate(fMap);
+  }
+
+  {
+    nsTArray<uint32_t> removal;
+    nsTArray<TableUpdate*> tableUpdates;
+
+    // Generate first partial update
+    CreateRandomSortedPrefixArray(3000, 4, 4, pArray);
+    CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
+    RemoveIntersection(fArray, pArray);
+    PrefixArrayToPrefixStringMap(pArray, pMap);
+
+    // Remove 1/5 of elements of original prefix set.
+    CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
+    RemoveElements(removal, fArray);
+
+    GenerateUpdateData(false, pMap, removal, tableUpdates);
+
+    fArray.AppendElements(pArray);
+    fArray.Sort();
+    pArray.Clear();
+    removal.Clear();
+
+    // Generate second partial update.
+    CreateRandomSortedPrefixArray(2000, 4, 4, pArray);
+    CreateRandomSortedPrefixArray(1000, 5, 32, pArray);
+    RemoveIntersection(fArray, pArray);
+    PrefixArrayToPrefixStringMap(pArray, pMap);
+
+    // Remove 1/5 of elements of original prefix set.
+    CreateRandomRemovalIndices(fArray.Length() / 5, fArray.Length(), removal);
+    RemoveElements(removal, fArray);
+
+    GenerateUpdateData(false, pMap, removal, tableUpdates);
+
+    mergedArray.AppendElements(fArray);
+    mergedArray.AppendElements(pArray);
+    mergedArray.Sort();
+    PrefixArrayToPrefixStringMap(mergedArray, mergedMap);
+
+    testUpdate(tableUpdates, mergedMap);
+  }
+
+  Clear();
+}
+
+// Test removal indices are larger than the original prefix set.
+TEST(UrlClassifierTableUpdateV4, RemovalIndexTooLarge)
+{
+  _PrefixArray fArray, pArray;
+  PrefixStringMap fMap, pMap;
+
+  {
+    CreateRandomSortedPrefixArray(1000, 4, 32, fArray);
+    PrefixArrayToPrefixStringMap(fArray, fMap);
+
+    testFullUpdate(fMap);
+  }
+
+  {
+    nsTArray<uint32_t> removal;
+    nsTArray<TableUpdate*> tableUpdates;
+
+    CreateRandomSortedPrefixArray(200, 4, 32, pArray);
+    RemoveIntersection(fArray, pArray);
+    PrefixArrayToPrefixStringMap(pArray, pMap);
+
+    for (uint32_t i = 0; i < fArray.Length() + 1 ;i++) {
+      removal.AppendElement(i);
+    }
+
+    GenerateUpdateData(false, pMap, removal, tableUpdates);
+    testUpdateFail(tableUpdates);
+  }
+
+  Clear();
+}