Bug 1276595 - Use protobuf API to parse v4 update response. r=francois draft
authorHenry Chang <hchang@mozilla.com>
Thu, 04 Aug 2016 14:34:52 +0800
changeset 396636 f6a7d39c48265614c48aa831926eda678a7c0ec5
parent 396634 1988d3ecc0349a70bc1b18443a7caa4a5dfd1bb6
child 527255 c6424e25b2127c0284ae82a23cb4cf7640870f0a
push id25062
push userhchang@mozilla.com
push dateThu, 04 Aug 2016 06:34:40 +0000
reviewersfrancois
bugs1276595
milestone51.0a1
Bug 1276595 - Use protobuf API to parse v4 update response. r=francois MozReview-Commit-ID: Gl3CFZI9vNc
toolkit/components/url-classifier/ProtocolParser.cpp
toolkit/components/url-classifier/ProtocolParser.h
toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
toolkit/components/url-classifier/nsUrlClassifierUtils.h
--- a/toolkit/components/url-classifier/ProtocolParser.cpp
+++ b/toolkit/components/url-classifier/ProtocolParser.cpp
@@ -109,16 +109,22 @@ ProtocolParser::AppendStream(const nsACS
     if (NS_FAILED(rv)) {
       mUpdateStatus = rv;
       return rv;
     }
   }
   return NS_OK;
 }
 
+void
+ProtocolParser::End()
+{
+  // Inbound data has already been processed in every AppendStream() call.
+}
+
 nsresult
 ProtocolParser::ProcessControl(bool* aDone)
 {
   nsresult rv;
 
   nsAutoCString line;
   *aDone = true;
   while (NextLine(line)) {
@@ -689,10 +695,183 @@ ProtocolParser::GetTableUpdate(const nsA
   // 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;
 }
 
+///////////////////////////////////////////////////////////////////////
+// ProtocolParserProtobuf
+
+ProtocolParserProtobuf::ProtocolParserProtobuf()
+{
+}
+
+ProtocolParserProtobuf::~ProtocolParserProtobuf()
+{
+}
+
+nsresult
+ProtocolParserProtobuf::AppendStream(const nsACString& aData)
+{
+  // Protobuf data cannot be parsed progressively. Just save the incoming data.
+  mPending.Append(aData);
+  return NS_OK;
+}
+
+void
+ProtocolParserProtobuf::End()
+{
+  // mUpdateStatus will be updated to success as long as not all
+  // the responses are invalid.
+  mUpdateStatus = NS_ERROR_FAILURE;
+
+  FetchThreatListUpdatesResponse response;
+  if (!response.ParseFromArray(mPending.get(), mPending.Length())) {
+    NS_WARNING("ProtocolParserProtobuf failed parsing data.");
+    return;
+  }
+
+  for (int i = 0; i < response.list_update_responses_size(); i++) {
+    auto r = response.list_update_responses(i);
+    nsresult rv = ProcessOneResponse(r);
+    if (NS_SUCCEEDED(rv)) {
+      mUpdateStatus = rv;
+    } else {
+      NS_WARNING("Failed to process one response.");
+    }
+  }
+}
+
+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;
+  }
+
+  // Convert threat type to list name.
+  nsCOMPtr<nsIUrlClassifierUtils> urlUtil =
+    do_GetService(NS_URLCLASSIFIERUTILS_CONTRACTID);
+  nsCString listName;
+  nsresult rv = urlUtil->ConvertThreatTypeToListName(aResponse.threat_type(),
+                                                     listName);
+  if (NS_FAILED(rv)) {
+    NS_WARNING(nsPrintfCString("Threat type to list name conversion error: %d",
+                               aResponse.threat_type().get()));
+    return NS_ERROR_FAILURE;
+  }
+
+  // Test if this is a full update.
+  bool isFullUpdate = false;
+  if (aResponse.has_response_type()) {
+    isFullUpdate =
+      aResponse.response_type() == ListUpdateResponse::FULL_UPDATE;
+  } else {
+    NS_WARNING("Response type not initialized.");
+    return NS_ERROR_FAILURE;
+  }
+
+  // Warn if there's no new state.
+  if (!aResponse.has_new_client_state()) {
+    NS_WARNING("New state not initialized.");
+    return 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);
+  PARSER_LOG(("\n\n"));
+
+  return NS_OK;
+}
+
+nsresult
+ProtocolParserProtobuf::ProcessAdditionOrRemoval(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.",
+                                  aIsAddition ? "Addition" : "Removal").get());
+      continue;
+    }
+
+    switch (update.compression_type()) {
+    case COMPRESSION_TYPE_UNSPECIFIED:
+      NS_WARNING("Unspecified compression type.");
+      break;
+
+    case RAW:
+      ret = (aIsAddition ? ProcessRawAddition(update)
+                         : ProcessRawRemoval(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)
+{
+  if (!aAddition.has_raw_hashes()) {
+    PARSER_LOG(("* No raw addition."));
+    return NS_OK;
+  }
+
+  auto rawHashes = aAddition.raw_hashes();
+  if (!rawHashes.has_prefix_size()) {
+    NS_WARNING("Raw hash has no prefix size");
+    return NS_OK;
+  }
+
+  auto prefixes = rawHashes.raw_hashes();
+  if (4 == rawHashes.prefix_size()) {
+    // Process fixed length prefixes separately.
+    uint32_t* fixedLengthPrefixes = (uint32_t*)prefixes.c_str();
+    size_t numOfFixedLengthPrefixes = prefixes.size() / 4;
+    PARSER_LOG(("* Raw addition (4 bytes)"));
+    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()));
+  }
+
+  return NS_OK;
+}
+
+nsresult
+ProtocolParserProtobuf::ProcessRawRemoval(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()));
+
+  return NS_OK;
+}
+
+
 } // namespace safebrowsing
 } // namespace mozilla
--- a/toolkit/components/url-classifier/ProtocolParser.h
+++ b/toolkit/components/url-classifier/ProtocolParser.h
@@ -3,41 +3,46 @@
  * 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 ProtocolParser_h__
 #define ProtocolParser_h__
 
 #include "HashStore.h"
 #include "nsICryptoHMAC.h"
+#include "safebrowsing.pb.h"
 
 namespace mozilla {
 namespace safebrowsing {
 
 /**
- * Some helpers for parsing the safe
+ * Helpers to parse the "shavar", "digest256" and "simple" list formats.
  */
 class ProtocolParser {
 public:
   struct ForwardedUpdate {
     nsCString table;
     nsCString url;
   };
 
   ProtocolParser();
-  ~ProtocolParser();
+  virtual ~ProtocolParser();
 
   nsresult Status() const { return mUpdateStatus; }
 
   nsresult Init(nsICryptoHash* aHasher);
 
   void SetCurrentTable(const nsACString& aTable);
 
   nsresult Begin();
-  nsresult AppendStream(const nsACString& aData);
+  virtual nsresult AppendStream(const nsACString& aData);
+
+  // Notify that the inbound data is ready for parsing if progressive
+  // parsing is not supported, for example in V4.
+  virtual void End();
 
   // 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.
@@ -68,16 +73,21 @@ private:
   // 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.
@@ -95,25 +105,45 @@ private:
     uint32_t hashSize;
     uint32_t length;
     void Clear() { num = 0; hashSize = 0; length = 0; }
   };
   ChunkState mChunkState;
 
   nsCOMPtr<nsICryptoHash> mCryptoHash;
 
-  nsresult mUpdateStatus;
-  nsCString mPending;
-
   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;
 };
 
+// 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 nsresult AppendStream(const nsACString& aData) override;
+  virtual void End() override;
+
+private:
+  virtual ~ProtocolParserProtobuf();
+
+  // For parsing update info.
+  nsresult ProcessOneResponse(const ListUpdateResponse& aResponse);
+  nsresult ProcessAdditionOrRemoval(const ThreatEntrySetList& aUpdate,
+                                    bool aIsAddition);
+  nsresult ProcessRawAddition(const ThreatEntrySet& aAddition);
+  nsresult ProcessRawRemoval(const ThreatEntrySet& aRemoval);
+};
+
 } // namespace safebrowsing
 } // namespace mozilla
 
 #endif
--- a/toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
+++ b/toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
@@ -436,17 +436,38 @@ nsUrlClassifierDBServiceWorker::BeginStr
 
   NS_ENSURE_STATE(mUpdateObserver);
   NS_ENSURE_STATE(!mInStream);
 
   mInStream = true;
 
   NS_ASSERTION(!mProtocolParser, "Should not have a protocol parser.");
 
-  mProtocolParser = new ProtocolParser();
+  // Check if we should use protobuf to parse the update.
+  bool useProtobuf = false;
+  for (size_t i = 0; i < mUpdateTables.Length(); i++) {
+    bool isCurProtobuf =
+      StringEndsWith(mUpdateTables[i], NS_LITERAL_CSTRING("-proto"));
+
+    if (0 == i) {
+      // Use the first table name to decice if all the subsequent tables
+      // should be '-proto'.
+      useProtobuf = isCurProtobuf;
+      continue;
+    }
+
+    if (useProtobuf != isCurProtobuf) {
+      NS_WARNING("Cannot mix 'proto' tables with other types "
+                 "within the same provider.");
+      break;
+    }
+  }
+
+  mProtocolParser = (useProtobuf ? new ProtocolParserProtobuf()
+                                 : new ProtocolParser());
   if (!mProtocolParser)
     return NS_ERROR_OUT_OF_MEMORY;
 
   mProtocolParser->Init(mCryptoHash);
 
   if (!table.IsEmpty()) {
     mProtocolParser->SetCurrentTable(table);
   }
@@ -507,16 +528,18 @@ nsUrlClassifierDBServiceWorker::FinishSt
     return NS_ERROR_NOT_INITIALIZED;
   }
 
   NS_ENSURE_STATE(mInStream);
   NS_ENSURE_STATE(mUpdateObserver);
 
   mInStream = false;
 
+  mProtocolParser->End();
+
   if (NS_SUCCEEDED(mProtocolParser->Status())) {
     if (mProtocolParser->UpdateWait()) {
       mUpdateWait = mProtocolParser->UpdateWait();
     }
     // XXX: Only allow forwards from the initial update?
     const nsTArray<ProtocolParser::ForwardedUpdate> &forwards =
       mProtocolParser->Forwards();
     for (uint32_t i = 0; i < forwards.Length(); i++) {
--- a/toolkit/components/url-classifier/nsUrlClassifierUtils.h
+++ b/toolkit/components/url-classifier/nsUrlClassifierUtils.h
@@ -39,17 +39,17 @@ private:
     // Store the 256 bits in an 8 byte array.
     uint32_t mMap[8];
   };
 
 
 public:
   nsUrlClassifierUtils();
 
-  NS_DECL_ISUPPORTS
+  NS_DECL_THREADSAFE_ISUPPORTS
   NS_DECL_NSIURLCLASSIFIERUTILS
 
   nsresult Init();
 
   nsresult CanonicalizeHostname(const nsACString & hostname,
                                 nsACString & _retval);
   nsresult CanonicalizePath(const nsACString & url, nsACString & _retval);