Bug 1284204 - Parse SafeBrowsing V4 updates to TableUpdateV4. r=francois. draft
authorHenry Chang <hchang@mozilla.com>
Fri, 12 Aug 2016 11:55:48 +0800
changeset 401662 0cf1c60713a938a72c87d140869a63b2afa12e3f
parent 401500 fe895421dfbe1f1f8f1fc6a39bb20774423a6d74
child 528556 4bef29c1b981b88c534e8b74559003365301a131
push id26534
push userhchang@mozilla.com
push dateWed, 17 Aug 2016 09:57:48 +0000
reviewersfrancois
bugs1284204
milestone51.0a1
Bug 1284204 - Parse SafeBrowsing V4 updates to TableUpdateV4. r=francois. MozReview-Commit-ID: 88AEIbosxKl
toolkit/components/url-classifier/Classifier.cpp
toolkit/components/url-classifier/HashStore.cpp
toolkit/components/url-classifier/HashStore.h
toolkit/components/url-classifier/ProtocolParser.cpp
toolkit/components/url-classifier/ProtocolParser.h
toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -648,25 +648,28 @@ Classifier::UpdateHashStore(nsTArray<Tab
     if (!update || !update->TableName().Equals(store.TableName()))
       continue;
 
     rv = store.ApplyUpdate(*update);
     NS_ENSURE_SUCCESS(rv, rv);
 
     applied++;
 
-    LOG(("Applied update to table %s:", store.TableName().get()));
-    LOG(("  %d add chunks", update->AddChunks().Length()));
-    LOG(("  %d add prefixes", update->AddPrefixes().Length()));
-    LOG(("  %d add completions", update->AddCompletes().Length()));
-    LOG(("  %d sub chunks", update->SubChunks().Length()));
-    LOG(("  %d sub prefixes", update->SubPrefixes().Length()));
-    LOG(("  %d sub completions", update->SubCompletes().Length()));
-    LOG(("  %d add expirations", update->AddExpirations().Length()));
-    LOG(("  %d sub expirations", update->SubExpirations().Length()));
+    auto updateV2 = TableUpdate::Cast<TableUpdateV2>(update);
+    if (updateV2) {
+      LOG(("Applied update to table %s:", store.TableName().get()));
+      LOG(("  %d add chunks", updateV2->AddChunks().Length()));
+      LOG(("  %d add prefixes", updateV2->AddPrefixes().Length()));
+      LOG(("  %d add completions", updateV2->AddCompletes().Length()));
+      LOG(("  %d sub chunks", updateV2->SubChunks().Length()));
+      LOG(("  %d sub prefixes", updateV2->SubPrefixes().Length()));
+      LOG(("  %d sub completions", updateV2->SubCompletes().Length()));
+      LOG(("  %d add expirations", updateV2->AddExpirations().Length()));
+      LOG(("  %d sub expirations", updateV2->SubExpirations().Length()));
+    }
 
     aUpdates->ElementAt(i) = nullptr;
     delete update;
   }
 
   LOG(("Applied %d update(s) to %s.", applied, store.TableName().get()));
 
   rv = store.Rebuild();
@@ -709,17 +712,18 @@ Classifier::UpdateCache(TableUpdate* aUp
   }
 
   nsAutoCString table(aUpdate->TableName());
   LOG(("Classifier::UpdateCache(%s)", table.get()));
 
   LookupCache *lookupCache = GetLookupCache(table);
   NS_ENSURE_TRUE(lookupCache, NS_ERROR_FAILURE);
 
-  lookupCache->AddCompletionsToCache(aUpdate->AddCompletes());
+  auto updateV2 = TableUpdate::Cast<TableUpdateV2>(aUpdate);
+  lookupCache->AddCompletionsToCache(updateV2->AddCompletes());
 
 #if defined(DEBUG)
   lookupCache->DumpCache();
 #endif
 
   return NS_OK;
 }
 
--- a/toolkit/components/url-classifier/HashStore.cpp
+++ b/toolkit/components/url-classifier/HashStore.cpp
@@ -115,57 +115,75 @@ extern mozilla::LazyLogModule gUrlClassi
 
 namespace mozilla {
 namespace safebrowsing {
 
 const uint32_t STORE_MAGIC = 0x1231af3b;
 const uint32_t CURRENT_VERSION = 3;
 
 nsresult
-TableUpdate::NewAddPrefix(uint32_t aAddChunk, const Prefix& aHash)
+TableUpdateV2::NewAddPrefix(uint32_t aAddChunk, const Prefix& aHash)
 {
   AddPrefix *add = mAddPrefixes.AppendElement(fallible);
   if (!add) return NS_ERROR_OUT_OF_MEMORY;
   add->addChunk = aAddChunk;
   add->prefix = aHash;
   return NS_OK;
 }
 
 nsresult
-TableUpdate::NewSubPrefix(uint32_t aAddChunk, const Prefix& aHash, uint32_t aSubChunk)
+TableUpdateV2::NewSubPrefix(uint32_t aAddChunk, const Prefix& aHash, uint32_t aSubChunk)
 {
   SubPrefix *sub = mSubPrefixes.AppendElement(fallible);
   if (!sub) return NS_ERROR_OUT_OF_MEMORY;
   sub->addChunk = aAddChunk;
   sub->prefix = aHash;
   sub->subChunk = aSubChunk;
   return NS_OK;
 }
 
 nsresult
-TableUpdate::NewAddComplete(uint32_t aAddChunk, const Completion& aHash)
+TableUpdateV2::NewAddComplete(uint32_t aAddChunk, const Completion& aHash)
 {
   AddComplete *add = mAddCompletes.AppendElement(fallible);
   if (!add) return NS_ERROR_OUT_OF_MEMORY;
   add->addChunk = aAddChunk;
   add->complete = aHash;
   return NS_OK;
 }
 
 nsresult
-TableUpdate::NewSubComplete(uint32_t aAddChunk, const Completion& aHash, uint32_t aSubChunk)
+TableUpdateV2::NewSubComplete(uint32_t aAddChunk, const Completion& aHash, uint32_t aSubChunk)
 {
   SubComplete *sub = mSubCompletes.AppendElement(fallible);
   if (!sub) return NS_ERROR_OUT_OF_MEMORY;
   sub->addChunk = aAddChunk;
   sub->complete = aHash;
   sub->subChunk = aSubChunk;
   return NS_OK;
 }
 
+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);
+  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]);
+  }
+}
+
 HashStore::HashStore(const nsACString& aTableName, nsIFile* aStoreDir)
   : mTableName(aTableName)
   , mStoreDirectory(aStoreDir)
   , mInUpdate(false)
   , mFileSize(0)
 {
 }
 
@@ -535,18 +553,23 @@ Merge(ChunkSet* aStoreChunks,
 
   aStorePrefixes->AppendElements(adds, fallible);
   EntrySort(*aStorePrefixes);
 
   return NS_OK;
 }
 
 nsresult
-HashStore::ApplyUpdate(TableUpdate &update)
+HashStore::ApplyUpdate(TableUpdate &aUpdate)
 {
+  auto updateV2 = TableUpdate::Cast<TableUpdateV2>(&aUpdate);
+  NS_ENSURE_TRUE(updateV2, NS_ERROR_FAILURE);
+
+  TableUpdateV2& update = *updateV2;
+
   nsresult rv = mAddExpirations.Merge(update.AddExpirations());
   NS_ENSURE_SUCCESS(rv, rv);
 
   rv = mSubExpirations.Merge(update.SubExpirations());
   NS_ENSURE_SUCCESS(rv, rv);
 
   rv = Expire();
   NS_ENSURE_SUCCESS(rv, rv);
--- a/toolkit/components/url-classifier/HashStore.h
+++ b/toolkit/components/url-classifier/HashStore.h
@@ -8,30 +8,61 @@
 #include "Entries.h"
 #include "ChunkSet.h"
 
 #include "nsString.h"
 #include "nsTArray.h"
 #include "nsIFile.h"
 #include "nsIFileStreams.h"
 #include "nsCOMPtr.h"
+#include "nsClassHashtable.h"
+#include "safebrowsing.pb.h"
+#include <string>
 
 namespace mozilla {
 namespace safebrowsing {
 
+// The abstract class of TableUpdateV2 and TableUpdateV4. This
+// is convenient for passing the TableUpdate* around associated
+// with v2 and v4 instance.
+class TableUpdate {
+public:
+  TableUpdate(const nsACString& aTable)
+    : mTable(aTable)
+  {
+  }
+
+  virtual ~TableUpdate() {}
+
+  // To be overriden.
+  virtual bool Empty() const = 0;
+
+  // Common interfaces.
+  const nsCString& TableName() const { return mTable; }
+
+  template<typename T>
+  static T* Cast(TableUpdate* aThat) {
+    return (T::TAG == aThat->Tag() ? reinterpret_cast<T*>(aThat) : nullptr);
+  }
+
+private:
+  virtual int Tag() const = 0;
+
+  nsCString mTable;
+};
+
 // A table update is built from a single update chunk from the server. As the
 // protocol parser processes each chunk, it constructs a table update with the
 // new hashes.
-class TableUpdate {
+class TableUpdateV2 : public TableUpdate {
 public:
-  explicit TableUpdate(const nsACString& aTable)
-      : mTable(aTable) {}
-  const nsCString& TableName() const { return mTable; }
+  explicit TableUpdateV2(const nsACString& aTable)
+    : TableUpdate(aTable) {}
 
-  bool Empty() const {
+  bool Empty() const override {
     return mAddChunks.Length() == 0 &&
       mSubChunks.Length() == 0 &&
       mAddExpirations.Length() == 0 &&
       mSubExpirations.Length() == 0 &&
       mAddPrefixes.Length() == 0 &&
       mSubPrefixes.Length() == 0 &&
       mAddCompletes.Length() == 0 &&
       mSubCompletes.Length() == 0;
@@ -69,32 +100,86 @@ public:
   ChunkSet& SubExpirations() { return mSubExpirations; }
 
   // Hashes associated with this chunk.
   AddPrefixArray& AddPrefixes() { return mAddPrefixes; }
   SubPrefixArray& SubPrefixes() { return mSubPrefixes; }
   AddCompleteArray& AddCompletes() { return mAddCompletes; }
   SubCompleteArray& SubCompletes() { return mSubCompletes; }
 
+  // For downcasting.
+  static const int TAG = 2;
+
 private:
-  nsCString mTable;
 
   // The list of chunk numbers that we have for each of the type of chunks.
   ChunkSet mAddChunks;
   ChunkSet mSubChunks;
   ChunkSet mAddExpirations;
   ChunkSet mSubExpirations;
 
   // 4-byte sha256 prefixes.
   AddPrefixArray mAddPrefixes;
   SubPrefixArray mSubPrefixes;
 
   // 32-byte hashes.
   AddCompleteArray mAddCompletes;
   SubCompleteArray mSubCompletes;
+
+  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 {
+  private:
+    std::string mStorage;
+    nsDependentCSubstring mString;
+
+  public:
+    explicit PrefixString(std::string& aString)
+    {
+      aString.swap(mStorage);
+      mString.Rebind(mStorage.data(), mStorage.size());
+    };
+
+    const nsACString& GetPrefixString() const { return mString; };
+  };
+
+  typedef nsClassHashtable<nsUint32HashKey, PrefixString> PrefixesStringMap;
+  typedef nsTArray<int32_t> RemovalIndiceArray;
+
+public:
+  explicit TableUpdateV4(const nsACString& aTable)
+    : TableUpdate(aTable)
+  {
+  }
+
+  bool Empty() const override
+  {
+    return mPrefixesMap.IsEmpty() && mRemovalIndiceArray.IsEmpty();
+  }
+
+  PrefixesStringMap& Prefixes() { return mPrefixesMap; }
+  RemovalIndiceArray& RemovalIndices() { return mRemovalIndiceArray; }
+
+  // For downcasting.
+  static const int TAG = 4;
+
+  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; }
+
+  PrefixesStringMap mPrefixesMap;
+  RemovalIndiceArray mRemovalIndiceArray;
 };
 
 // There is one hash store per table.
 class HashStore {
 public:
   HashStore(const nsACString& aTableName, nsIFile* aStoreFile);
   ~HashStore();
 
--- a/toolkit/components/url-classifier/ProtocolParser.cpp
+++ b/toolkit/components/url-classifier/ProtocolParser.cpp
@@ -56,45 +56,87 @@ ParseChunkRange(nsACString::const_iterat
   if (numRead == 1) {
     *aLast = *aFirst;
     return true;
   }
 
   return false;
 }
 
+///////////////////////////////////////////////////////////////
+// ProtocolParser implementation
+
 ProtocolParser::ProtocolParser()
   : mUpdateStatus(NS_OK)
-  , mState(PROTOCOL_STATE_CONTROL)
-  , mUpdateWait(0)
-  , mResetRequested(false)
-  , mTableUpdate(nullptr)
 {
 }
 
 ProtocolParser::~ProtocolParser()
 {
   CleanupUpdates();
 }
 
 nsresult
 ProtocolParser::Init(nsICryptoHash* aHasher)
 {
   mCryptoHash = aHasher;
   return NS_OK;
 }
 
 void
-ProtocolParser::SetCurrentTable(const nsACString& aTable)
+ProtocolParser::CleanupUpdates()
+{
+  for (uint32_t i = 0; i < mTableUpdates.Length(); i++) {
+    delete mTableUpdates[i];
+  }
+  mTableUpdates.Clear();
+}
+
+TableUpdate *
+ProtocolParser::GetTableUpdate(const nsACString& aTable)
 {
-  mTableUpdate = GetTableUpdate(aTable);
+  for (uint32_t i = 0; i < mTableUpdates.Length(); i++) {
+    if (aTable.Equals(mTableUpdates[i]->TableName())) {
+      return mTableUpdates[i];
+    }
+  }
+
+  // We free automatically on destruction, ownership of these
+  // updates can be transferred to DBServiceWorker, which passes
+  // them back to Classifier when doing the updates, and that
+  // will free them.
+  TableUpdate *update = CreateTableUpdate(aTable);
+  mTableUpdates.AppendElement(update);
+  return update;
+}
+
+///////////////////////////////////////////////////////////////////////
+// ProtocolParserV2
+
+ProtocolParserV2::ProtocolParserV2()
+  : mState(PROTOCOL_STATE_CONTROL)
+  , mUpdateWait(0)
+  , mResetRequested(false)
+  , mTableUpdate(nullptr)
+{
+}
+
+ProtocolParserV2::~ProtocolParserV2()
+{
+}
+
+void
+ProtocolParserV2::SetCurrentTable(const nsACString& aTable)
+{
+  auto update = GetTableUpdate(aTable);
+  mTableUpdate = TableUpdate::Cast<TableUpdateV2>(update);
 }
 
 nsresult
-ProtocolParser::AppendStream(const nsACString& aData)
+ProtocolParserV2::AppendStream(const nsACString& aData)
 {
   if (NS_FAILED(mUpdateStatus))
     return mUpdateStatus;
 
   nsresult rv;
   mPending.Append(aData);
 
   bool done = false;
@@ -111,23 +153,23 @@ ProtocolParser::AppendStream(const nsACS
       mUpdateStatus = rv;
       return rv;
     }
   }
   return NS_OK;
 }
 
 void
-ProtocolParser::End()
+ProtocolParserV2::End()
 {
   // Inbound data has already been processed in every AppendStream() call.
 }
 
 nsresult
-ProtocolParser::ProcessControl(bool* aDone)
+ProtocolParserV2::ProcessControl(bool* aDone)
 {
   nsresult rv;
 
   nsAutoCString line;
   *aDone = true;
   while (NextLine(line)) {
     PARSER_LOG(("Processing %s\n", line.get()));
 
@@ -157,17 +199,17 @@ ProtocolParser::ProcessControl(bool* aDo
     }
   }
 
   *aDone = true;
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessExpirations(const nsCString& aLine)
+ProtocolParserV2::ProcessExpirations(const nsCString& aLine)
 {
   if (!mTableUpdate) {
     NS_WARNING("Got an expiration without a table.");
     return NS_ERROR_FAILURE;
   }
   const nsCSubstring &list = Substring(aLine, 3);
   nsACString::const_iterator begin, end;
   list.BeginReading(begin);
@@ -193,17 +235,17 @@ ProtocolParser::ProcessExpirations(const
     } else {
       return NS_ERROR_FAILURE;
     }
   }
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessChunkControl(const nsCString& aLine)
+ProtocolParserV2::ProcessChunkControl(const nsCString& aLine)
 {
   if (!mTableUpdate) {
     NS_WARNING("Got a chunk before getting a table.");
     return NS_ERROR_FAILURE;
   }
 
   mState = PROTOCOL_STATE_CHUNK;
   char command;
@@ -267,39 +309,39 @@ ProtocolParser::ProcessChunkControl(cons
       }
       break;
   }
 
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessForward(const nsCString& aLine)
+ProtocolParserV2::ProcessForward(const nsCString& aLine)
 {
   const nsCSubstring &forward = Substring(aLine, 2);
   return AddForward(forward);
 }
 
 nsresult
-ProtocolParser::AddForward(const nsACString& aUrl)
+ProtocolParserV2::AddForward(const nsACString& aUrl)
 {
   if (!mTableUpdate) {
     NS_WARNING("Forward without a table name.");
     return NS_ERROR_FAILURE;
   }
 
   ForwardedUpdate *forward = mForwards.AppendElement();
   forward->table = mTableUpdate->TableName();
   forward->url.Assign(aUrl);
 
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessChunk(bool* aDone)
+ProtocolParserV2::ProcessChunk(bool* aDone)
 {
   if (!mTableUpdate) {
     NS_WARNING("Processing chunk without an active table.");
     return NS_ERROR_FAILURE;
   }
 
   NS_ASSERTION(mChunkState.num != 0, "Must have a chunk number.");
 
@@ -326,17 +368,17 @@ ProtocolParser::ProcessChunk(bool* aDone
   }
   return ProcessPlaintextChunk(chunk);
 }
 
 /**
  * Process a plaintext chunk (currently only used in unit tests).
  */
 nsresult
-ProtocolParser::ProcessPlaintextChunk(const nsACString& aChunk)
+ProtocolParserV2::ProcessPlaintextChunk(const nsACString& aChunk)
 {
   if (!mTableUpdate) {
     NS_WARNING("Chunk received with no table.");
     return NS_ERROR_FAILURE;
   }
 
   PARSER_LOG(("Handling a %d-byte simple chunk", aChunk.Length()));
 
@@ -395,17 +437,17 @@ ProtocolParser::ProcessPlaintextChunk(co
       }
     }
   }
 
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessShaChunk(const nsACString& aChunk)
+ProtocolParserV2::ProcessShaChunk(const nsACString& aChunk)
 {
   uint32_t start = 0;
   while (start < aChunk.Length()) {
     // First four bytes are the domain key.
     Prefix domain;
     domain.Assign(Substring(aChunk, start, DOMAIN_SIZE));
     start += DOMAIN_SIZE;
 
@@ -435,31 +477,31 @@ ProtocolParser::ProcessShaChunk(const ns
     }
     NS_ENSURE_SUCCESS(rv, rv);
   }
 
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessDigestChunk(const nsACString& aChunk)
+ProtocolParserV2::ProcessDigestChunk(const nsACString& aChunk)
 {
   PARSER_LOG(("Handling a %d-byte digest256 chunk", aChunk.Length()));
 
   if (mChunkState.type == CHUNK_ADD_DIGEST) {
     return ProcessDigestAdd(aChunk);
   }
   if (mChunkState.type == CHUNK_SUB_DIGEST) {
     return ProcessDigestSub(aChunk);
   }
   return NS_ERROR_UNEXPECTED;
 }
 
 nsresult
-ProtocolParser::ProcessDigestAdd(const nsACString& aChunk)
+ProtocolParserV2::ProcessDigestAdd(const nsACString& aChunk)
 {
   // The ABNF format for add chunks is (HASH)+, where HASH is 32 bytes.
   MOZ_ASSERT(aChunk.Length() % 32 == 0,
              "Chunk length in bytes must be divisible by 4");
   uint32_t start = 0;
   while (start < aChunk.Length()) {
     Completion hash;
     hash.Assign(Substring(aChunk, start, COMPLETE_SIZE));
@@ -468,17 +510,17 @@ ProtocolParser::ProcessDigestAdd(const n
     if (NS_FAILED(rv)) {
       return rv;
     }
   }
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessDigestSub(const nsACString& aChunk)
+ProtocolParserV2::ProcessDigestSub(const nsACString& aChunk)
 {
   // The ABNF format for sub chunks is (ADDCHUNKNUM HASH)+, where ADDCHUNKNUM
   // is a 4 byte chunk number, and HASH is 32 bytes.
   MOZ_ASSERT(aChunk.Length() % 36 == 0,
              "Chunk length in bytes must be divisible by 36");
   uint32_t start = 0;
   while (start < aChunk.Length()) {
     // Read ADDCHUNKNUM
@@ -498,17 +540,17 @@ ProtocolParser::ProcessDigestSub(const n
     if (NS_FAILED(rv)) {
       return rv;
     }
   }
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessHostAdd(const Prefix& aDomain, uint8_t aNumEntries,
+ProtocolParserV2::ProcessHostAdd(const Prefix& aDomain, uint8_t aNumEntries,
                                const nsACString& aChunk, uint32_t* aStart)
 {
   NS_ASSERTION(mChunkState.hashSize == PREFIX_SIZE,
                "ProcessHostAdd should only be called for prefix hashes.");
 
   if (aNumEntries == 0) {
     nsresult rv = mTableUpdate->NewAddPrefix(mChunkState.num, aDomain);
     if (NS_FAILED(rv)) {
@@ -532,17 +574,17 @@ ProtocolParser::ProcessHostAdd(const Pre
     }
     *aStart += PREFIX_SIZE;
   }
 
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessHostSub(const Prefix& aDomain, uint8_t aNumEntries,
+ProtocolParserV2::ProcessHostSub(const Prefix& aDomain, uint8_t aNumEntries,
                                const nsACString& aChunk, uint32_t *aStart)
 {
   NS_ASSERTION(mChunkState.hashSize == PREFIX_SIZE,
                "ProcessHostSub should only be called for prefix hashes.");
 
   if (aNumEntries == 0) {
     if ((*aStart) + 4 > aChunk.Length()) {
       NS_WARNING("Received a zero-entry sub chunk without an associated add.");
@@ -587,17 +629,17 @@ ProtocolParser::ProcessHostSub(const Pre
       return rv;
     }
   }
 
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessHostAddComplete(uint8_t aNumEntries,
+ProtocolParserV2::ProcessHostAddComplete(uint8_t aNumEntries,
                                        const nsACString& aChunk, uint32_t* aStart)
 {
   NS_ASSERTION(mChunkState.hashSize == COMPLETE_SIZE,
                "ProcessHostAddComplete should only be called for complete hashes.");
 
   if (aNumEntries == 0) {
     // this is totally comprehensible.
     // My sarcasm detector is going off!
@@ -619,17 +661,17 @@ ProtocolParser::ProcessHostAddComplete(u
     }
     *aStart += COMPLETE_SIZE;
   }
 
   return NS_OK;
 }
 
 nsresult
-ProtocolParser::ProcessHostSubComplete(uint8_t aNumEntries,
+ProtocolParserV2::ProcessHostSubComplete(uint8_t aNumEntries,
                                        const nsACString& aChunk, uint32_t* aStart)
 {
   NS_ASSERTION(mChunkState.hashSize == COMPLETE_SIZE,
                "ProcessHostSubComplete should only be called for complete hashes.");
 
   if (aNumEntries == 0) {
     // this is totally comprehensible.
     NS_WARNING("Expected > 0 entries for a 32-byte hash sub.");
@@ -658,65 +700,58 @@ ProtocolParser::ProcessHostSubComplete(u
       return rv;
     }
   }
 
   return NS_OK;
 }
 
 bool
-ProtocolParser::NextLine(nsACString& aLine)
+ProtocolParserV2::NextLine(nsACString& aLine)
 {
   int32_t newline = mPending.FindChar('\n');
   if (newline == kNotFound) {
     return false;
   }
   aLine.Assign(Substring(mPending, 0, newline));
   mPending.Cut(0, newline + 1);
   return true;
 }
 
-void
-ProtocolParser::CleanupUpdates()
-{
-  for (uint32_t i = 0; i < mTableUpdates.Length(); i++) {
-    delete mTableUpdates[i];
-  }
-  mTableUpdates.Clear();
-}
-
-TableUpdate *
-ProtocolParser::GetTableUpdate(const nsACString& aTable)
+TableUpdate*
+ProtocolParserV2::CreateTableUpdate(const nsACString& aTableName) const
 {
-  for (uint32_t i = 0; i < mTableUpdates.Length(); i++) {
-    if (aTable.Equals(mTableUpdates[i]->TableName())) {
-      return mTableUpdates[i];
-    }
-  }
-
-  // We free automatically on destruction, ownership of these
-  // updates can be transferred to DBServiceWorker, which passes
-  // them back to Classifier when doing the updates, and that
-  // will free them.
-  TableUpdate *update = new TableUpdate(aTable);
-  mTableUpdates.AppendElement(update);
-  return update;
+  return new TableUpdateV2(aTableName);
 }
 
 ///////////////////////////////////////////////////////////////////////
 // ProtocolParserProtobuf
 
 ProtocolParserProtobuf::ProtocolParserProtobuf()
 {
 }
 
 ProtocolParserProtobuf::~ProtocolParserProtobuf()
 {
 }
 
+void
+ProtocolParserProtobuf::SetCurrentTable(const nsACString& aTable)
+{
+  // Should never occur.
+  MOZ_ASSERT_UNREACHABLE("SetCurrentTable shouldn't be called");
+}
+
+
+TableUpdate*
+ProtocolParserProtobuf::CreateTableUpdate(const nsACString& aTableName) const
+{
+  return new TableUpdateV4(aTableName);
+}
+
 nsresult
 ProtocolParserProtobuf::AppendStream(const nsACString& aData)
 {
   // Protobuf data cannot be parsed progressively. Just save the incoming data.
   mPending.Append(aData);
   return NS_OK;
 }
 
@@ -776,29 +811,34 @@ ProtocolParserProtobuf::ProcessOneRespon
   }
 
   // Warn if there's no new state.
   if (!aResponse.has_new_client_state()) {
     NS_WARNING("New state not initialized.");
     return NS_ERROR_FAILURE;
   }
 
+  auto tu = GetTableUpdate(nsCString(listName.get()));
+  auto tuV4 = TableUpdate::Cast<TableUpdateV4>(tu);
+  NS_ENSURE_TRUE(tuV4, NS_ERROR_FAILURE);
+
   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")));
-  ProcessAdditionOrRemoval(aResponse.additions(), true /*aIsAddition*/);
-  ProcessAdditionOrRemoval(aResponse.removals(), false);
+  ProcessAdditionOrRemoval(*tuV4, aResponse.additions(), true /*aIsAddition*/);
+  ProcessAdditionOrRemoval(*tuV4, aResponse.removals(), false);
   PARSER_LOG(("\n\n"));
 
   return NS_OK;
 }
 
 nsresult
-ProtocolParserProtobuf::ProcessAdditionOrRemoval(const ThreatEntrySetList& aUpdate,
+ProtocolParserProtobuf::ProcessAdditionOrRemoval(TableUpdateV4& aTableUpdate,
+                                                 const ThreatEntrySetList& aUpdate,
                                                  bool aIsAddition)
 {
   nsresult ret = NS_OK;
 
   for (int i = 0; i < aUpdate.size(); i++) {
     auto update = aUpdate.Get(i);
     if (!update.has_compression_type()) {
       NS_WARNING(nsPrintfCString("%s with no compression type.",
@@ -807,32 +847,33 @@ ProtocolParserProtobuf::ProcessAdditionO
     }
 
     switch (update.compression_type()) {
     case COMPRESSION_TYPE_UNSPECIFIED:
       NS_WARNING("Unspecified compression type.");
       break;
 
     case RAW:
-      ret = (aIsAddition ? ProcessRawAddition(update)
-                         : ProcessRawRemoval(update));
+      ret = (aIsAddition ? ProcessRawAddition(aTableUpdate, update)
+                         : ProcessRawRemoval(aTableUpdate, update));
       break;
 
     case RICE:
       // Not implemented yet (see bug 1285848),
       NS_WARNING("Encoded table update is not supported yet.");
       break;
     }
   }
 
   return ret;
 }
 
 nsresult
-ProtocolParserProtobuf::ProcessRawAddition(const ThreatEntrySet& aAddition)
+ProtocolParserProtobuf::ProcessRawAddition(TableUpdateV4& aTableUpdate,
+                                           const ThreatEntrySet& aAddition)
 {
   if (!aAddition.has_raw_hashes()) {
     PARSER_LOG(("* No raw addition."));
     return NS_OK;
   }
 
   auto rawHashes = aAddition.raw_hashes();
   if (!rawHashes.has_prefix_size()) {
@@ -849,30 +890,42 @@ ProtocolParserProtobuf::ProcessRawAdditi
     PARSER_LOG(("  - # of prefixes: %d", numOfFixedLengthPrefixes));
     PARSER_LOG(("  - Memory address: 0x%p", fixedLengthPrefixes));
   } else {
     // TODO: Process variable length prefixes including full hashes.
     // See Bug 1283009.
     PARSER_LOG((" Raw addition (%d bytes)", rawHashes.prefix_size()));
   }
 
+  if (!rawHashes.mutable_raw_hashes()) {
+    PARSER_LOG(("Unable to get mutable raw hashes. Can't perform a string move."));
+    return NS_ERROR_FAILURE;
+  }
+
+  aTableUpdate.NewPrefixes(rawHashes.prefix_size(),
+                           *rawHashes.mutable_raw_hashes());
+
   return NS_OK;
 }
 
 nsresult
-ProtocolParserProtobuf::ProcessRawRemoval(const ThreatEntrySet& aRemoval)
+ProtocolParserProtobuf::ProcessRawRemoval(TableUpdateV4& aTableUpdate,
+                                          const ThreatEntrySet& aRemoval)
 {
   if (!aRemoval.has_raw_indices()) {
     NS_WARNING("A removal has no indices.");
     return NS_OK;
   }
 
   // indices is an array of int32.
   auto indices = aRemoval.raw_indices().indices();
   PARSER_LOG(("* Raw removal"));
   PARSER_LOG(("  - # of removal: %d", indices.size()));
 
+  aTableUpdate.NewRemovalIndices((const uint32_t*)indices.data(),
+                                 indices.size());
+
   return NS_OK;
 }
 
 
 } // namespace safebrowsing
 } // namespace mozilla
--- a/toolkit/components/url-classifier/ProtocolParser.h
+++ b/toolkit/components/url-classifier/ProtocolParser.h
@@ -9,53 +9,91 @@
 #include "HashStore.h"
 #include "nsICryptoHMAC.h"
 #include "safebrowsing.pb.h"
 
 namespace mozilla {
 namespace safebrowsing {
 
 /**
- * Helpers to parse the "shavar", "digest256" and "simple" list formats.
+ * Abstract base class for parsing update data in multiple formats.
  */
 class ProtocolParser {
 public:
   struct ForwardedUpdate {
     nsCString table;
     nsCString url;
   };
 
   ProtocolParser();
   virtual ~ProtocolParser();
 
   nsresult Status() const { return mUpdateStatus; }
 
   nsresult Init(nsICryptoHash* aHasher);
 
-  void SetCurrentTable(const nsACString& aTable);
+  virtual void SetCurrentTable(const nsACString& aTable) = 0;
 
   nsresult Begin();
-  virtual nsresult AppendStream(const nsACString& aData);
+  virtual nsresult AppendStream(const nsACString& aData) = 0;
 
   // Notify that the inbound data is ready for parsing if progressive
   // parsing is not supported, for example in V4.
-  virtual void End();
+  virtual void End() = 0;
 
   // Forget the table updates that were created by this pass.  It
   // becomes the caller's responsibility to free them.  This is shitty.
   TableUpdate *GetTableUpdate(const nsACString& aTable);
   void ForgetTableUpdates() { mTableUpdates.Clear(); }
   nsTArray<TableUpdate*> &GetTableUpdates() { return mTableUpdates; }
 
-  // Update information.
-  const nsTArray<ForwardedUpdate> &Forwards() const { return mForwards; }
-  int32_t UpdateWait() { return mUpdateWait; }
-  bool ResetRequested() { return mResetRequested; }
+  // These are only meaningful to V2. Since they were originally public,
+  // moving them to ProtocolParserV2 requires a dymamic cast in the call
+  // sites. As a result, we will leave them until retirely removing V2
+  // codes.
+  virtual const nsTArray<ForwardedUpdate> &Forwards() const { return mForwards; }
+  virtual int32_t UpdateWait() { return 0; }
+  virtual bool ResetRequested() { return false; }
+
+protected:
+  virtual TableUpdate* CreateTableUpdate(const nsACString& aTableName) const = 0;
+
+  nsCString mPending;
+  nsresult mUpdateStatus;
+
+  // Keep track of updates to apply before passing them to the DBServiceWorkers.
+  nsTArray<TableUpdate*> mTableUpdates;
+
+  nsTArray<ForwardedUpdate> mForwards;
+  nsCOMPtr<nsICryptoHash> mCryptoHash;
 
 private:
+  void CleanupUpdates();
+};
+
+/**
+ * Helpers to parse the "shavar", "digest256" and "simple" list formats.
+ */
+class ProtocolParserV2 final : public ProtocolParser {
+public:
+  ProtocolParserV2();
+  virtual ~ProtocolParserV2();
+
+  virtual void SetCurrentTable(const nsACString& aTable) override;
+  virtual nsresult AppendStream(const nsACString& aData) override;
+  virtual void End() override;
+
+  // Update information.
+  virtual const nsTArray<ForwardedUpdate> &Forwards() const override { return mForwards; }
+  virtual int32_t UpdateWait() override { return mUpdateWait; }
+  virtual bool ResetRequested() override { return mResetRequested; }
+
+private:
+  virtual TableUpdate* CreateTableUpdate(const nsACString& aTableName) const override;
+
   nsresult ProcessControl(bool* aDone);
   nsresult ProcessExpirations(const nsCString& aLine);
   nsresult ProcessChunkControl(const nsCString& aLine);
   nsresult ProcessForward(const nsCString& aLine);
   nsresult AddForward(const nsACString& aUrl);
   nsresult ProcessChunk(bool* done);
   // Remove this, it's only used for testing
   nsresult ProcessPlaintextChunk(const nsACString& aChunk);
@@ -71,23 +109,16 @@ private:
   // Digest chunks are very similar to shavar chunks, except digest chunks
   // always contain the full hash, so there is no need for chunk data to
   // contain prefix sizes.
   nsresult ProcessDigestChunk(const nsACString& aChunk);
   nsresult ProcessDigestAdd(const nsACString& aChunk);
   nsresult ProcessDigestSub(const nsACString& aChunk);
   bool NextLine(nsACString& aLine);
 
-  void CleanupUpdates();
-
-protected:
-  nsCString mPending;
-  nsresult mUpdateStatus;
-
-private:
   enum ParserState {
     PROTOCOL_STATE_CONTROL,
     PROTOCOL_STATE_CHUNK
   };
   ParserState mState;
 
   enum ChunkType {
     // Types for shavar tables.
@@ -103,47 +134,51 @@ private:
     ChunkType type;
     uint32_t num;
     uint32_t hashSize;
     uint32_t length;
     void Clear() { num = 0; hashSize = 0; length = 0; }
   };
   ChunkState mChunkState;
 
-  nsCOMPtr<nsICryptoHash> mCryptoHash;
-
   uint32_t mUpdateWait;
   bool mResetRequested;
 
-  nsTArray<ForwardedUpdate> mForwards;
-  // Keep track of updates to apply before passing them to the DBServiceWorkers.
-  nsTArray<TableUpdate*> mTableUpdates;
   // Updates to apply to the current table being parsed.
-  TableUpdate *mTableUpdate;
+  TableUpdateV2 *mTableUpdate;
 };
 
 // Helpers to parse the "proto" list format.
 class ProtocolParserProtobuf final : public ProtocolParser {
 public:
   typedef FetchThreatListUpdatesResponse_ListUpdateResponse ListUpdateResponse;
   typedef google::protobuf::RepeatedPtrField<ThreatEntrySet> ThreatEntrySetList;
 
 public:
   ProtocolParserProtobuf();
 
+  virtual void SetCurrentTable(const nsACString& aTable) override;
   virtual nsresult AppendStream(const nsACString& aData) override;
   virtual void End() override;
 
 private:
   virtual ~ProtocolParserProtobuf();
 
+  virtual TableUpdate* CreateTableUpdate(const nsACString& aTableName) const override;
+
   // For parsing update info.
   nsresult ProcessOneResponse(const ListUpdateResponse& aResponse);
-  nsresult ProcessAdditionOrRemoval(const ThreatEntrySetList& aUpdate,
+
+  nsresult ProcessAdditionOrRemoval(TableUpdateV4& aTableUpdate,
+                                    const ThreatEntrySetList& aUpdate,
                                     bool aIsAddition);
-  nsresult ProcessRawAddition(const ThreatEntrySet& aAddition);
-  nsresult ProcessRawRemoval(const ThreatEntrySet& aRemoval);
+
+  nsresult ProcessRawAddition(TableUpdateV4& aTableUpdate,
+                              const ThreatEntrySet& aAddition);
+
+  nsresult ProcessRawRemoval(TableUpdateV4& aTableUpdate,
+                             const ThreatEntrySet& aRemoval);
 };
 
 } // namespace safebrowsing
 } // namespace mozilla
 
 #endif
--- a/toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
+++ b/toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
@@ -455,18 +455,18 @@ nsUrlClassifierDBServiceWorker::BeginStr
 
     if (useProtobuf != isCurProtobuf) {
       NS_WARNING("Cannot mix 'proto' tables with other types "
                  "within the same provider.");
       break;
     }
   }
 
-  mProtocolParser = (useProtobuf ? new ProtocolParserProtobuf()
-                                 : new ProtocolParser());
+  mProtocolParser = (useProtobuf ? static_cast<ProtocolParser*>(new ProtocolParserProtobuf())
+                                 : static_cast<ProtocolParser*>(new ProtocolParserV2()));
   if (!mProtocolParser)
     return NS_ERROR_OUT_OF_MEMORY;
 
   mProtocolParser->Init(mCryptoHash);
 
   if (!table.IsEmpty()) {
     mProtocolParser->SetCurrentTable(table);
   }
@@ -714,17 +714,17 @@ nsUrlClassifierDBServiceWorker::CacheCom
   // Ownership is transferred in to us
   nsAutoPtr<CacheResultArray> resultsPtr(results);
 
   if (mLastResults == *resultsPtr) {
     LOG(("Skipping completions that have just been cached already."));
     return NS_OK;
   }
 
-  nsAutoPtr<ProtocolParser> pParse(new ProtocolParser());
+  nsAutoPtr<ProtocolParserV2> pParse(new ProtocolParserV2());
   nsTArray<TableUpdate*> updates;
 
   // Only cache results for tables that we have, don't take
   // in tables we might accidentally have hit during a completion.
   // This happens due to goog vs googpub lists existing.
   nsTArray<nsCString> tables;
   nsresult rv = mClassifier->ActiveTables(tables);
   NS_ENSURE_SUCCESS(rv, rv);
@@ -733,31 +733,35 @@ nsUrlClassifierDBServiceWorker::CacheCom
     bool activeTable = false;
     for (uint32_t table = 0; table < tables.Length(); table++) {
       if (tables[table].Equals(resultsPtr->ElementAt(i).table)) {
         activeTable = true;
         break;
       }
     }
     if (activeTable) {
-      TableUpdate * tu = pParse->GetTableUpdate(resultsPtr->ElementAt(i).table);
+      TableUpdateV2* tuV2 = TableUpdate::Cast<TableUpdateV2>(
+        pParse->GetTableUpdate(resultsPtr->ElementAt(i).table));
+
+      NS_ENSURE_TRUE(tuV2, NS_ERROR_FAILURE);
+
       LOG(("CacheCompletion Addchunk %d hash %X", resultsPtr->ElementAt(i).entry.addChunk,
            resultsPtr->ElementAt(i).entry.ToUint32()));
-      rv = tu->NewAddComplete(resultsPtr->ElementAt(i).entry.addChunk,
-                              resultsPtr->ElementAt(i).entry.complete);
+      rv = tuV2->NewAddComplete(resultsPtr->ElementAt(i).entry.addChunk,
+                                resultsPtr->ElementAt(i).entry.complete);
       if (NS_FAILED(rv)) {
         // We can bail without leaking here because ForgetTableUpdates
         // hasn't been called yet.
         return rv;
       }
-      rv = tu->NewAddChunk(resultsPtr->ElementAt(i).entry.addChunk);
+      rv = tuV2->NewAddChunk(resultsPtr->ElementAt(i).entry.addChunk);
       if (NS_FAILED(rv)) {
         return rv;
       }
-      updates.AppendElement(tu);
+      updates.AppendElement(tuV2);
       pParse->ForgetTableUpdates();
     } else {
       LOG(("Completion received, but table is not active, so not caching."));
     }
    }
 
   mClassifier->ApplyFullHashes(&updates);
   mLastResults = *resultsPtr;