Bug 1305484 - Save/load state and checksum to/from disk rather than prefs. r=dimi,francois
authorHenry Chang <hchang@mozilla.com>
Thu, 13 Oct 2016 15:22:08 +0800
changeset 425531 599f3dd65b6aea6bfb1749e0fbbdb50811107f31
parent 425530 a77e3aaf913c2534f6338d4de14a456a0bffbdb4
child 425532 f1d9a4dc244f1e7ff3b129fe04f9eef44c6308ac
push id32446
push userchevobbe.nicolas@gmail.com
push dateFri, 14 Oct 2016 22:50:51 +0000
reviewersdimi, francois
bugs1305484
milestone52.0a1
Bug 1305484 - Save/load state and checksum to/from disk rather than prefs. r=dimi,francois MozReview-Commit-ID: 4gmmrI9wY4c
toolkit/components/url-classifier/Classifier.cpp
toolkit/components/url-classifier/Classifier.h
toolkit/components/url-classifier/HashStore.cpp
toolkit/components/url-classifier/HashStore.h
toolkit/components/url-classifier/LookupCacheV4.cpp
toolkit/components/url-classifier/LookupCacheV4.h
toolkit/components/url-classifier/ProtocolParser.cpp
toolkit/components/url-classifier/content/listmanager.js
toolkit/components/url-classifier/nsIUrlClassifierDBService.idl
toolkit/components/url-classifier/tests/unit/test_listmanager.js
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -13,26 +13,29 @@
 #include "nsISeekableStream.h"
 #include "nsIFile.h"
 #include "nsNetCID.h"
 #include "nsPrintfCString.h"
 #include "nsThreadUtils.h"
 #include "mozilla/Telemetry.h"
 #include "mozilla/Logging.h"
 #include "mozilla/SyncRunnable.h"
+#include "mozilla/Base64.h"
 
 // MOZ_LOG=UrlClassifierDbService:5
 extern mozilla::LazyLogModule gUrlClassifierDbServiceLog;
 #define LOG(args) MOZ_LOG(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug, args)
 #define LOG_ENABLED() MOZ_LOG_TEST(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug)
 
 #define STORE_DIRECTORY      NS_LITERAL_CSTRING("safebrowsing")
 #define TO_DELETE_DIR_SUFFIX NS_LITERAL_CSTRING("-to_delete")
 #define BACKUP_DIR_SUFFIX    NS_LITERAL_CSTRING("-backup")
 
+#define METADATA_SUFFIX      NS_LITERAL_CSTRING(".metadata")
+
 namespace mozilla {
 namespace safebrowsing {
 
 namespace {
 
 // A scoped-clearer for nsTArray<TableUpdate*>.
 // The owning elements will be deleted and the array itself
 // will be cleared on exiting the scope.
@@ -386,16 +389,17 @@ Classifier::DeleteTables(const nsTArray<
     }
   }
   NS_ENSURE_SUCCESS_VOID(rv);
 }
 
 void
 Classifier::TableRequest(nsACString& aResult)
 {
+  // Generating v2 table info.
   nsTArray<nsCString> tables;
   ActiveTables(tables);
   for (uint32_t i = 0; i < tables.Length(); i++) {
     HashStore store(tables[i], mRootStoreDirectory);
 
     nsresult rv = store.Open();
     if (NS_FAILED(rv))
       continue;
@@ -419,16 +423,23 @@ Classifier::TableRequest(nsACString& aRe
       aResult.AppendLiteral("s:");
       nsAutoCString subList;
       subs.Serialize(subList);
       aResult.Append(subList);
     }
 
     aResult.Append('\n');
   }
+
+  // Load meta data from *.metadata files in the root directory.
+  // Specifically for v4 tables.
+  nsCString metadata;
+  nsresult rv = LoadMetadata(mRootStoreDirectory, metadata);
+  NS_ENSURE_SUCCESS_VOID(rv);
+  aResult.Append(metadata);
 }
 
 nsresult
 Classifier::Check(const nsACString& aSpec,
                   const nsACString& aTables,
                   uint32_t aFreshnessGuarantee,
                   LookupResultArray& aResults)
 {
@@ -950,16 +961,17 @@ Classifier::UpdateTableV4(nsTArray<Table
   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;
 
+  TableUpdateV4* lastAppliedUpdate = nullptr;
   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);
@@ -994,25 +1006,35 @@ Classifier::UpdateTableV4(nsTArray<Table
       }
 
       rv = lookupCache->ApplyPartialUpdate(updateV4, *input, *output);
       NS_ENSURE_SUCCESS(rv, rv);
 
       input->Clear();
     }
 
+    // Keep track of the last applied update.
+    lastAppliedUpdate = updateV4;
+
     aUpdates->ElementAt(i) = nullptr;
   }
 
   rv = lookupCache->Build(*output);
   NS_ENSURE_SUCCESS(rv, rv);
 
   rv = lookupCache->WriteFile();
   NS_ENSURE_SUCCESS(rv, rv);
 
+  if (lastAppliedUpdate) {
+    LOG(("Write meta data of the last applied update."));
+    rv = lookupCache->WriteMetadata(lastAppliedUpdate);
+    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
@@ -1105,10 +1127,81 @@ Classifier::ReadNoiseEntries(const Prefi
     if (newPref != aPrefix) {
       aNoiseEntries->AppendElement(newPref);
     }
   }
 
   return NS_OK;
 }
 
+nsresult
+Classifier::LoadMetadata(nsIFile* aDirectory, nsACString& aResult)
+{
+  nsCOMPtr<nsISimpleEnumerator> entries;
+  nsresult rv = aDirectory->GetDirectoryEntries(getter_AddRefs(entries));
+  NS_ENSURE_SUCCESS(rv, rv);
+  NS_ENSURE_ARG_POINTER(entries);
+
+  bool hasMore;
+  while (NS_SUCCEEDED(rv = entries->HasMoreElements(&hasMore)) && hasMore) {
+    nsCOMPtr<nsISupports> supports;
+    rv = entries->GetNext(getter_AddRefs(supports));
+    NS_ENSURE_SUCCESS(rv, rv);
+
+    nsCOMPtr<nsIFile> file = do_QueryInterface(supports);
+
+    // If |file| is a directory, recurse to find its entries as well.
+    bool isDirectory;
+    if (NS_FAILED(file->IsDirectory(&isDirectory))) {
+      continue;
+    }
+    if (isDirectory) {
+      LoadMetadata(file, aResult);
+      continue;
+    }
+
+    // Truncate file extension to get the table name.
+    nsCString tableName;
+    rv = file->GetNativeLeafName(tableName);
+    NS_ENSURE_SUCCESS(rv, rv);
+
+    int32_t dot = tableName.RFind(METADATA_SUFFIX, 0);
+    if (dot == -1) {
+      continue;
+    }
+    tableName.Cut(dot, METADATA_SUFFIX.Length());
+
+    LookupCacheV4* lookupCache =
+      LookupCache::Cast<LookupCacheV4>(GetLookupCache(tableName));
+    if (!lookupCache) {
+      continue;
+    }
+
+    nsCString state;
+    nsCString checksum;
+    rv = lookupCache->LoadMetadata(state, checksum);
+    if (NS_FAILED(rv)) {
+      LOG(("Failed to get metadata for table %s", tableName.get()));
+      continue;
+    }
+
+    // The state might include '\n' so that we have to encode.
+    nsAutoCString stateBase64;
+    rv = Base64Encode(state, stateBase64);
+    NS_ENSURE_SUCCESS(rv, rv);
+
+    nsAutoCString checksumBase64;
+    rv = Base64Encode(checksum, checksumBase64);
+    NS_ENSURE_SUCCESS(rv, rv);
+
+    LOG(("Appending state '%s' and checksum '%s' for table %s",
+         stateBase64.get(), checksumBase64.get(), tableName.get()));
+
+    aResult.AppendPrintf("%s;%s:%s\n", tableName.get(),
+                                       stateBase64.get(),
+                                       checksumBase64.get());
+  }
+
+  return rv;
+}
+
 } // namespace safebrowsing
 } // namespace mozilla
--- a/toolkit/components/url-classifier/Classifier.h
+++ b/toolkit/components/url-classifier/Classifier.h
@@ -120,16 +120,18 @@ private:
 
   nsresult UpdateCache(TableUpdate* aUpdates);
 
   LookupCache *GetLookupCache(const nsACString& aTable);
 
   bool CheckValidUpdate(nsTArray<TableUpdate*>* aUpdates,
                         const nsACString& aTable);
 
+  nsresult LoadMetadata(nsIFile* aDirectory, nsACString& aResult);
+
   // Root dir of the Local profile.
   nsCOMPtr<nsIFile> mCacheDirectory;
   // Main directory where to store the databases.
   nsCOMPtr<nsIFile> mRootStoreDirectory;
   // Used for atomically updating the other dirs.
   nsCOMPtr<nsIFile> mBackupDirectory;
   nsCOMPtr<nsIFile> mToDeleteDirectory;
   nsCOMPtr<nsICryptoHash> mCryptoHash;
--- a/toolkit/components/url-classifier/HashStore.cpp
+++ b/toolkit/components/url-classifier/HashStore.cpp
@@ -195,16 +195,22 @@ TableUpdateV4::NewPrefixes(int32_t aSize
 void
 TableUpdateV4::NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices)
 {
   for (size_t i = 0; i < aNumOfIndices; i++) {
     mRemovalIndiceArray.AppendElement(aIndices[i]);
   }
 }
 
+void
+TableUpdateV4::NewChecksum(const std::string& aChecksum)
+{
+  mChecksum.Assign(aChecksum.data(), aChecksum.size());
+}
+
 HashStore::HashStore(const nsACString& aTableName, nsIFile* aRootStoreDir)
   : mTableName(aTableName)
   , mInUpdate(false)
   , mFileSize(0)
 {
   nsresult rv = Classifier::GetPrivateStoreDirectory(aRootStoreDir,
                                                      aTableName,
                                                      getter_AddRefs(mStoreDirectory));
--- a/toolkit/components/url-classifier/HashStore.h
+++ b/toolkit/components/url-classifier/HashStore.h
@@ -160,30 +160,36 @@ public:
   bool Empty() const override
   {
     return mPrefixesMap.IsEmpty() && mRemovalIndiceArray.IsEmpty();
   }
 
   bool IsFullUpdate() const { return mFullUpdate; }
   PrefixStdStringMap& Prefixes() { return mPrefixesMap; }
   RemovalIndiceArray& RemovalIndices() { return mRemovalIndiceArray; }
+  const nsACString& ClientState() const { return mClientState; }
+  const nsACString& Checksum() const { return mChecksum; }
 
   // 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);
+  void SetNewClientState(const nsACString& aState) { mClientState = aState; }
+  void NewChecksum(const std::string& aChecksum);
 
 private:
   virtual int Tag() const override { return TAG; }
 
   bool mFullUpdate;
   PrefixStdStringMap mPrefixesMap;
   RemovalIndiceArray mRemovalIndiceArray;
+  nsCString mClientState;
+  nsCString mChecksum;
 };
 
 // There is one hash store per table.
 class HashStore {
 public:
   HashStore(const nsACString& aTableName, nsIFile* aRootStoreFile);
   ~HashStore();
 
--- a/toolkit/components/url-classifier/LookupCacheV4.cpp
+++ b/toolkit/components/url-classifier/LookupCacheV4.cpp
@@ -6,16 +6,18 @@
 #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)
 
+#define METADATA_SUFFIX NS_LITERAL_CSTRING(".metadata")
+
 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)
@@ -228,16 +230,175 @@ LookupCacheV4::ApplyPartialUpdate(TableU
     Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
                           WRONG_REMOVAL_INDICES);
     return NS_ERROR_FAILURE;
   }
 
   return NS_OK;
 }
 
+//////////////////////////////////////////////////////////////////////////
+// A set of lightweight functions for reading/writing value from/to file.
+
+namespace {
+
+template<typename T>
+struct ValueTraits
+{
+  static uint32_t Length(const T& aValue) { return sizeof(T); }
+  static char* WritePtr(T& aValue, uint32_t aLength) { return (char*)&aValue; }
+  static const char* ReadPtr(const T& aValue) { return (char*)&aValue; }
+  static bool IsFixedLength() { return true; }
+};
+
+template<>
+struct ValueTraits<nsACString>
+{
+  static bool IsFixedLength() { return false; }
+
+  static uint32_t Length(const nsACString& aValue)
+  {
+    return aValue.Length();
+  }
+
+  static char* WritePtr(nsACString& aValue, uint32_t aLength)
+  {
+    aValue.SetLength(aLength);
+    return aValue.BeginWriting();
+  }
+
+  static const char* ReadPtr(const nsACString& aValue)
+  {
+    return aValue.BeginReading();
+  }
+};
+
+template<typename T> static nsresult
+WriteValue(nsIOutputStream *aOutputStream, const T& aValue)
+{
+  uint32_t writeLength = ValueTraits<T>::Length(aValue);
+  if (!ValueTraits<T>::IsFixedLength()) {
+    // We need to write out the variable value length.
+    nsresult rv = WriteValue(aOutputStream, writeLength);
+    NS_ENSURE_SUCCESS(rv, rv);
+  }
+
+  // Write out the value.
+  auto valueReadPtr = ValueTraits<T>::ReadPtr(aValue);
+  uint32_t written;
+  nsresult rv = aOutputStream->Write(valueReadPtr, writeLength, &written);
+  if (NS_FAILED(rv) || written != writeLength) {
+    LOG(("Failed to write the value."));
+    return NS_FAILED(rv) ? rv : NS_ERROR_FAILURE;
+  }
+
+  return rv;
+}
+
+template<typename T> static nsresult
+ReadValue(nsIInputStream* aInputStream, T& aValue)
+{
+  nsresult rv;
+
+  uint32_t readLength;
+  if (ValueTraits<T>::IsFixedLength()) {
+    readLength = ValueTraits<T>::Length(aValue);
+  } else {
+    // Read the variable value length from file.
+    nsresult rv = ReadValue(aInputStream, readLength);
+    NS_ENSURE_SUCCESS(rv, rv);
+  }
+
+  // Read the value.
+  uint32_t read;
+  auto valueWritePtr = ValueTraits<T>::WritePtr(aValue, readLength);
+  rv = aInputStream->Read(valueWritePtr, readLength, &read);
+  if (NS_FAILED(rv) || read != readLength) {
+    LOG(("Failed to read the value."));
+    return NS_FAILED(rv) ? rv : NS_ERROR_FAILURE;
+  }
+
+  return rv;
+}
+
+} // end of unnamed namespace.
+////////////////////////////////////////////////////////////////////////
+
+nsresult
+LookupCacheV4::WriteMetadata(TableUpdateV4* aTableUpdate)
+{
+  NS_ENSURE_ARG_POINTER(aTableUpdate);
+
+  nsCOMPtr<nsIFile> metaFile;
+  nsresult rv = mStoreDirectory->Clone(getter_AddRefs(metaFile));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  rv = metaFile->AppendNative(mTableName + METADATA_SUFFIX);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  nsCOMPtr<nsIOutputStream> outputStream;
+  rv = NS_NewLocalFileOutputStream(getter_AddRefs(outputStream), metaFile,
+                                   PR_WRONLY | PR_TRUNCATE | PR_CREATE_FILE);
+  if (!NS_SUCCEEDED(rv)) {
+    LOG(("Unable to create file to store metadata."));
+    return rv;
+  }
+
+  // Write the state.
+  rv = WriteValue(outputStream, aTableUpdate->ClientState());
+  if (NS_FAILED(rv)) {
+    LOG(("Failed to write the list state."));
+    return rv;
+  }
+
+  // Write the checksum.
+  rv = WriteValue(outputStream, aTableUpdate->Checksum());
+  if (NS_FAILED(rv)) {
+    LOG(("Failed to write the list checksum."));
+    return rv;
+  }
+
+  return rv;
+}
+
+nsresult
+LookupCacheV4::LoadMetadata(nsACString& aState, nsACString& aChecksum)
+{
+  nsCOMPtr<nsIFile> metaFile;
+  nsresult rv = mStoreDirectory->Clone(getter_AddRefs(metaFile));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  rv = metaFile->AppendNative(mTableName + METADATA_SUFFIX);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  nsCOMPtr<nsIInputStream> localInFile;
+  rv = NS_NewLocalFileInputStream(getter_AddRefs(localInFile), metaFile,
+                                  PR_RDONLY | nsIFile::OS_READAHEAD);
+  if (NS_FAILED(rv)) {
+    LOG(("Unable to open metadata file."));
+    return rv;
+  }
+
+  // Read the list state.
+  rv = ReadValue(localInFile, aState);
+  if (NS_FAILED(rv)) {
+    LOG(("Failed to read state."));
+    return rv;
+  }
+
+  // Read the checksum.
+  rv = ReadValue(localInFile, aChecksum);
+  if (NS_FAILED(rv)) {
+    LOG(("Failed to read checksum."));
+    return rv;
+  }
+
+  return rv;
+}
+
 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;
   }
--- a/toolkit/components/url-classifier/LookupCacheV4.h
+++ b/toolkit/components/url-classifier/LookupCacheV4.h
@@ -30,16 +30,19 @@ public:
   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);
 
+  nsresult WriteMetadata(TableUpdateV4* aTableUpdate);
+  nsresult LoadMetadata(nsACString& aState, nsACString& aChecksum);
+
   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;
 
--- a/toolkit/components/url-classifier/ProtocolParser.cpp
+++ b/toolkit/components/url-classifier/ProtocolParser.cpp
@@ -777,37 +777,16 @@ ProtocolParserProtobuf::End()
     if (NS_SUCCEEDED(rv)) {
       mUpdateStatus = rv;
     } else {
       NS_WARNING("Failed to process one response.");
     }
   }
 }
 
-// Save state of |aListName| to the following pref:
-//
-//   "browser.safebrowsing.provider.google4.state.[aListName]"
-//
-static nsresult
-SaveStateToPref(const nsACString& aListName, const nsACString& aState)
-{
-  nsresult rv;
-  nsCOMPtr<nsIPrefBranch> prefs(do_GetService(NS_PREFSERVICE_CONTRACTID, &rv));
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  nsCString prefName("browser.safebrowsing.provider.google4.state.");
-  prefName.Append(aListName);
-
-  nsCString stateBase64;
-  rv = Base64Encode(aState, stateBase64);
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  return prefs->SetCharPref(prefName.get(), stateBase64.get());
-}
-
 nsresult
 ProtocolParserProtobuf::ProcessOneResponse(const ListUpdateResponse& aResponse)
 {
   // A response must have a threat type.
   if (!aResponse.has_threat_type()) {
     NS_WARNING("Threat type not initialized. This seems to be an invalid response.");
     return NS_ERROR_FAILURE;
   }
@@ -857,29 +836,29 @@ ProtocolParserProtobuf::ProcessOneRespon
     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);
 
-  // See Bug 1287059. We save the state to prefs until we support
-  // "saving states to HashStore".
   nsCString state(aResponse.new_client_state().c_str(),
                   aResponse.new_client_state().size());
-  NS_DispatchToMainThread(NS_NewRunnableFunction([listName, state] () {
-    DebugOnly<nsresult> rv = SaveStateToPref(listName, state);
-    NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "SaveStateToPref failed");
-  }));
+  tuV4->SetNewClientState(state);
+
+  if (aResponse.has_checksum()) {
+    tuV4->NewChecksum(aResponse.checksum().sha256());
+  }
 
   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")));
+  PARSER_LOG(("* hasChecksum: %s\n", (aResponse.has_checksum() ? "yes" : "no")));
 
   tuV4->SetFullUpdate(isFullUpdate);
   ProcessAdditionOrRemoval(*tuV4, aResponse.additions(), true /*aIsAddition*/);
   ProcessAdditionOrRemoval(*tuV4, aResponse.removals(), false);
   PARSER_LOG(("\n\n"));
 
   return NS_OK;
 }
--- a/toolkit/components/url-classifier/content/listmanager.js
+++ b/toolkit/components/url-classifier/content/listmanager.js
@@ -390,30 +390,42 @@ PROT_ListManager.prototype.makeUpdateReq
   if (useProtobuf) {
     let tableArray = [];
     Object.keys(streamerMap.tableNames).forEach(aTableName => {
       if (streamerMap.tableNames[aTableName]) {
         tableArray.push(aTableName);
       }
     });
 
+    // Build the <tablename, stateBase64> mapping.
+    let tableState = {};
+    tableData.split("\n").forEach(line => {
+      let p = line.indexOf(";");
+      if (-1 === p) {
+        return;
+      }
+      let tableName = line.substring(0, p);
+      let metadata = line.substring(p + 1).split(":");
+      let stateBase64 = metadata[0];
+      log(tableName + " ==> " + stateBase64);
+      tableState[tableName] = stateBase64;
+    });
+
     // The state is a byte stream which server told us from the
     // last table update. The state would be used to do the partial
     // update and the empty string means the table has
     // never been downloaded. See Bug 1287058 for supporting
     // partial update.
     let stateArray = [];
     tableArray.forEach(listName => {
-      // See Bug 1287059. We save the state to prefs until we support
-      // "saving states to HashStore".
-      let statePrefName = "browser.safebrowsing.provider.google4.state." + listName;
-      let stateBase64 = this.prefs_.getPref(statePrefName, "");
-      stateArray.push(stateBase64);
+      stateArray.push(tableState[listName] || "");
     });
 
+    log("stateArray: " + stateArray);
+
     let urlUtils = Cc["@mozilla.org/url-classifier/utils;1"]
                      .getService(Ci.nsIUrlClassifierUtils);
 
     streamerMap.requestPayload = urlUtils.makeUpdateRequestV4(tableArray,
                                                               stateArray,
                                                               tableArray.length);
     streamerMap.isPostRequest = false;
   } else {
--- a/toolkit/components/url-classifier/nsIUrlClassifierDBService.idl
+++ b/toolkit/components/url-classifier/nsIUrlClassifierDBService.idl
@@ -76,24 +76,34 @@ interface nsIUrlClassifierDBService : ns
    * @param c: The callback will be called with a comma-separated list
    *        of tables to which the key belongs.
    */
   void lookup(in nsIPrincipal principal,
               in ACString tables,
               in nsIUrlClassifierCallback c);
 
   /**
-   * Lists the tables along with which chunks are available in each table.
-   * This list is in the format of the request body:
-   *   tablename;chunkdata\n
-   *   tablename2;chunkdata2\n
+   * Lists the tables along with their meta info in the following format:
+   *
+   *   tablename;[metadata]\n
+   *   tablename2;[metadata]\n
+   *
+   * For v2 tables, the metadata is the chunks info such as
    *
-   * For example:
-   *   goog-phish-regexp;a:10,14,30-40s:56,67
-   *   goog-white-regexp;a:1-3,5
+   *   goog-phish-shavar;a:10,14,30-40s:56,67
+   *   goog-unwanted-shavar;a:1-3,5
+   *
+   * For v4 tables, base64 encoded state is currently the only info in the
+   * metadata (can be extended whenever necessary). For exmaple,
+   *
+   *   goog-phish-proto;Cg0IARAGGAEiAzAwMTABEKqTARoCGAjT1gDD:oCGAjT1gDD\n
+   *   goog-malware-proto;Cg0IAhAGGAEiAzAwMTABENCQARoCGAjx5Yty:BENCQARoCGAj\n
+   *
+   * Note that the metadata is colon-separated.
+   *
    */
   void getTables(in nsIUrlClassifierCallback c);
 
   /**
    * Set the nsIUrlClassifierCompleter object for a given table.  This
    * object will be used to request complete versions of partial
    * hashes.
    */
--- a/toolkit/components/url-classifier/tests/unit/test_listmanager.js
+++ b/toolkit/components/url-classifier/tests/unit/test_listmanager.js
@@ -64,16 +64,17 @@ let gExpectedQueryV4 = "";
 let gHttpServV4 = null;
 
 // These two variables are used to synchronize the last two racing updates
 // (in terms of "update URL") in test_update_all_tables().
 let gUpdatedCntForTableData = 0; // For TEST_TABLE_DATA_LIST.
 let gIsV4Updated = false;   // For TEST_TABLE_DATA_V4.
 
 const NEW_CLIENT_STATE = 'sta\0te';
+const CHECKSUM = 'check\0sum';
 
 prefBranch.setBoolPref("browser.safebrowsing.debug", true);
 
 // The "\xFF\xFF" is to generate a base64 string with "/".
 prefBranch.setCharPref("browser.safebrowsing.id", "Firefox\xFF\xFF");
 
 // Register tables.
 TEST_TABLE_DATA_LIST.forEach(function(t) {
@@ -255,36 +256,37 @@ function run_test() {
     response.setStatusLine(request.httpVersion, 200, "OK");
 
     // The protobuf binary represention of response:
     //
     // [
     //   {
     //     'threat_type': 2, // SOCIAL_ENGINEERING_PUBLIC
     //     'response_type': 2, // FULL_UPDATE
-    //     'new_client_state': 'sta\x00te' // NEW_CLIENT_STATE
+    //     'new_client_state': 'sta\x00te', // NEW_CLIENT_STATE
+    //     'checksum': { "sha256": 'check\x00sum' }, // CHECKSUM
+    //     'additions': { 'compression_type': RAW,
+    //                    'prefix_size': 4,
+    //                    'raw_hashes': "00000001000000020000000300000004"}
     //   }
     // ]
     //
-    let content = "\x0A\x0C\x08\x02\x20\x02\x3A\x06\x73\x74\x61\x00\x74\x65";
+    let content = "\x0A\x33\x08\x02\x20\x02\x2A\x18\x08\x01\x12\x14\x08\x04\x12\x10\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x02\x00\x00\x00\x03\x3A\x06\x73\x74\x61\x00\x74\x65\x42\x0B\x0A\x09\x63\x68\x65\x63\x6B\x00\x73\x75\x6D\x12\x08\x08\x08\x10\x80\x94\xEB\xDC\x03";
 
     response.bodyOutputStream.write(content, content.length);
 
     if (gIsV4Updated) {
       // This falls to the case where test_partialUpdateV4 is running.
       // We are supposed to have verified the update request contains
       // the state we set in the previous request.
       run_next_test();
       return;
     }
 
-    // See Bug 1284204. We save the state to pref at the moment to
-    // support partial update until "storing to HashStore" is supported.
-    // Here we poll the pref until the state has been saved.
-    waitUntilStateSavedToPref(NEW_CLIENT_STATE, () => {
+    waitUntilMetaDataSaved(NEW_CLIENT_STATE, CHECKSUM, () => {
       gIsV4Updated = true;
 
       if (gUpdatedCntForTableData === SERVER_INVOLVED_TEST_CASE_LIST.length) {
         // All tests are done!
         run_next_test();
         return;
       }
 
@@ -328,26 +330,47 @@ function readFileToString(aFilename) {
   let f = do_get_file(aFilename);
   let stream = Cc["@mozilla.org/network/file-input-stream;1"]
     .createInstance(Ci.nsIFileInputStream);
   stream.init(f, -1, 0, 0);
   let buf = NetUtil.readInputStreamToString(stream, stream.available());
   return buf;
 }
 
-function waitUntilStateSavedToPref(expectedState, callback) {
-  const STATE_PREF_NAME_PREFIX = 'browser.safebrowsing.provider.google4.state.';
+function waitUntilMetaDataSaved(expectedState, expectedChecksum, callback) {
+  let dbService = Cc["@mozilla.org/url-classifier/dbservice;1"]
+                     .getService(Ci.nsIUrlClassifierDBService);
 
-  let stateBase64 = '';
+  dbService.getTables(metaData => {
+    do_print("metadata: " + metaData);
+    let didCallback = false;
+    metaData.split("\n").some(line => {
+      // Parse [tableName];[stateBase64]
+      let p = line.indexOf(";");
+      if (-1 === p) {
+        return false; // continue.
+      }
+      let tableName = line.substring(0, p);
+      let metadata = line.substring(p + 1).split(":");
+      let stateBase64 = metadata[0];
+      let checksumBase64 = metadata[1];
 
-  try {
-    stateBase64 =
-      prefBranch.getCharPref(STATE_PREF_NAME_PREFIX + 'test-phish-proto');
-  } catch (e) {}
+      if (tableName !== 'test-phish-proto') {
+        return false; // continue.
+      }
+
+      if (stateBase64 === btoa(expectedState) &&
+          checksumBase64 === btoa(expectedChecksum)) {
+        do_print('State has been saved to disk!');
+        callback();
+        didCallback = true;
+      }
 
-  if (stateBase64 === btoa(expectedState)) {
-    do_print('State has been saved to pref!');
-    callback();
-    return;
-  }
+      return true; // break no matter whether the state is matching.
+    });
 
-  do_timeout(1000, waitUntilStateSavedToPref.bind(null, expectedState, callback));
+    if (!didCallback) {
+      do_timeout(1000, waitUntilMetaDataSaved.bind(null, expectedState,
+                                                         expectedChecksum,
+                                                         callback));
+    }
+  });
 }