Bug 1274112 - Part 2: Use protobuf API to parse v4 update response. r=francois

MozReview-Commit-ID: 3sjR3Feq4ua

--HG--
extra : rebase_source : 0c3ef1b0b0d564190c8e2e410a1201216cdbf53c

toolkit/components/url-classifier/ProtocolParser.cpp

@@ -11,6 +11,7 @@
 #include "prprf.h"
 
 #include "nsUrlClassifierUtils.h"
+#include "nsPrintfCString.h"
 
 // MOZ_LOG=UrlClassifierProtocolParser:5
 mozilla::LazyLogModule gUrlClassifierProtocolParserLog("UrlClassifierProtocolParser");
@@ -61,8 +62,8 @@ ParseChunkRange(nsACString::const_iterator& aBegin,
 }
 
 ProtocolParser::ProtocolParser()
-    : mState(PROTOCOL_STATE_CONTROL)
-  , mUpdateStatus(NS_OK)
+  : mUpdateStatus(NS_OK)
+  , mState(PROTOCOL_STATE_CONTROL)
   , mUpdateWait(0)
   , mResetRequested(false)
   , mTableUpdate(nullptr)
@@ -114,6 +115,12 @@ ProtocolParser::AppendStream(const nsACString& aData)
   return NS_OK;
 }
 
+void
+ProtocolParser::End()
+{
+  // Inbound data has already been processed in every AppendStream() call.
+}
+
 nsresult
 ProtocolParser::ProcessControl(bool* aDone)
 {
@@ -694,5 +701,178 @@ ProtocolParser::GetTableUpdate(const nsACString& aTable)
   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)) {
+    PARSER_LOG((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

toolkit/components/url-classifier/ProtocolParser.h

@@ -8,12 +8,13 @@
 
 #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:
@@ -23,7 +24,7 @@ public:
   };
 
   ProtocolParser();
-  ~ProtocolParser();
+  virtual ~ProtocolParser();
 
   nsresult Status() const { return mUpdateStatus; }
 
@@ -32,7 +33,11 @@ public:
   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.
@@ -73,6 +78,11 @@ private:
 
   void CleanupUpdates();
 
+protected:
+  nsCString mPending;
+  nsresult mUpdateStatus;
+
+private:
   enum ParserState {
     PROTOCOL_STATE_CONTROL,
     PROTOCOL_STATE_CHUNK
@@ -100,9 +110,6 @@ private:
 
   nsCOMPtr<nsICryptoHash> mCryptoHash;
 
-  nsresult mUpdateStatus;
-  nsCString mPending;
-
   uint32_t mUpdateWait;
   bool mResetRequested;
 
@@ -113,6 +120,29 @@ private:
   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
 

toolkit/components/url-classifier/nsUrlClassifierDBService.cpp

@@ -441,7 +441,28 @@ nsUrlClassifierDBServiceWorker::BeginStream(const nsACString &table)
 
   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;
 
@@ -512,6 +533,8 @@ nsUrlClassifierDBServiceWorker::FinishStream()
 
   mInStream = false;
 
+  mProtocolParser->End();
+
   if (NS_SUCCEEDED(mProtocolParser->Status())) {
     if (mProtocolParser->UpdateWait()) {
       mUpdateWait = mProtocolParser->UpdateWait();

toolkit/components/url-classifier/nsUrlClassifierUtils.h

@@ -44,7 +44,7 @@ private:
 public:
   nsUrlClassifierUtils();
 
-  NS_DECL_ISUPPORTS
+  NS_DECL_THREADSAFE_ISUPPORTS
   NS_DECL_NSIURLCLASSIFIERUTILS
 
   nsresult Init();