Rice decode removals and additions and use them to update the store.

components/safe_browsing_db/v4_store.cc

Index: components/safe_browsing_db/v4_store.cc
diff --git a/components/safe_browsing_db/v4_store.cc b/components/safe_browsing_db/v4_store.cc
index 4e731ae797c57a5929fd81b6dacd1b962487dbea..aef3b6ea179bdc76f5da5d2ca659ced43140e242 100644
--- a/components/safe_browsing_db/v4_store.cc
+++ b/components/safe_browsing_db/v4_store.cc
@@ -5,9 +5,11 @@
 #include "base/base64.h"
 #include "base/bind.h"
 #include "base/files/file_util.h"
+#include "base/memory/ptr_util.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/metrics/sparse_histogram.h"
 #include "base/strings/stringprintf.h"
+#include "components/safe_browsing_db/v4_rice.h"
 #include "components/safe_browsing_db/v4_store.h"
 #include "components/safe_browsing_db/v4_store.pb.h"
 
@@ -40,6 +42,16 @@ void RecordApplyUpdateResult(ApplyUpdateResult result) {
                             APPLY_UPDATE_RESULT_MAX);
 }
 
+void RecordDecodeAdditionsResult(V4DecodeResult result) {
+  UMA_HISTOGRAM_ENUMERATION("SafeBrowsing.V4DecodeAdditionsResult", result,
+                            DECODE_RESULT_MAX);
+}
+
+void RecordDecodeRemovalsResult(V4DecodeResult result) {
+  UMA_HISTOGRAM_ENUMERATION("SafeBrowsing.V4DecodeRemovalsResult", result,

[Nathan Parker, 2016/08/02 23:43:30]

It'd be good to get some size/time metrics in here as well, in another CL.  Size
of updates, time to read updates, time to unpack/write updates, etc.

[vakh (use Gerrit instead), 2016/08/03 01:44:49]

The update size is currently being logged as:
SafeBrowsing.V4UpdateResponseSizeKB

Added a TODO for measuring and recording the timings.

+                            DECODE_RESULT_MAX);
+}
+
 void RecordApplyUpdateResultWhenReadingFromDisk(ApplyUpdateResult result) {
   UMA_HISTOGRAM_ENUMERATION(
       "SafeBrowsing.V4ApplyUpdateResultWhenReadingFromDisk", result,
@@ -112,24 +124,47 @@ void V4Store::ApplyUpdate(
   // response_type as FULL_UPDATE, and write that to disk.
   // 3. Remove this if condition after completing 1. and 2.
   HashPrefixMap hash_prefix_map;
-  ApplyUpdateResult apply_update_result;
+  ApplyUpdateResult apply_update_result = APPLY_UPDATE_SUCCESS;
   if (response->response_type() == ListUpdateResponse::PARTIAL_UPDATE) {
     const RepeatedField<int32>* raw_removals = nullptr;
+    std::unique_ptr<RepeatedField<int32>> rice_removals;

[Nathan Parker, 2016/08/02 23:43:30]

This can just be a RepeatedField<int32>

[vakh (use Gerrit instead), 2016/08/03 01:44:49]

Done.

     size_t removals_size = response->removals_size();
     DCHECK_LE(removals_size, 1u);
     if (removals_size == 1) {
       const ThreatEntrySet& removal = response->removals().Get(0);
-      // TODO(vakh): Allow other compression types.
-      // See: https://bugs.chromium.org/p/chromium/issues/detail?id=624567
-      DCHECK_EQ(RAW, removal.compression_type());
-      raw_removals = &removal.raw_indices().indices();
+      const CompressionType compression_type = removal.compression_type();

[Nathan Parker, 2016/08/02 23:43:31]

nit: Not sure this temporary var helps much.  Removal.compression_type() is
pretty easy to read.

[vakh (use Gerrit instead), 2016/08/03 06:54:52]

Done.

+      if (compression_type == RAW) {
+        raw_removals = &removal.raw_indices().indices();
+      } else if (compression_type == RICE) {
+        DCHECK(removal.has_rice_indices());
+
+        rice_removals = base::MakeUnique<RepeatedField<int32>>();

[Nathan Parker, 2016/08/02 23:43:30]

This just re-inits it to what it was before, ya?

[vakh (use Gerrit instead), 2016/08/03 06:54:52]

Done.

+        const RiceDeltaEncoding& rice_indices = removal.rice_indices();
+        V4DecodeResult decode_result = V4RiceDecoder::DecodeIntegers(
+            rice_indices.first_value(), rice_indices.rice_parameter(),
+            rice_indices.num_entries(), rice_indices.encoded_data(),
+            rice_removals.get());
+        RecordDecodeRemovalsResult(decode_result);
+        if (decode_result != DECODE_SUCCESS) {
+          apply_update_result = RICE_DECODING_FAILURE;
+        } else {
+          raw_removals = rice_removals.get();
+        }
+      } else {
+        apply_update_result = UNEXPECTED_COMPRESSION_TYPE_REMOVALS_FAILURE;
+        NOTREACHED() << "Unexpected compression_type type: "
+                     << compression_type;
+      }
     }
 
-    apply_update_result = UpdateHashPrefixMapFromAdditions(
-        response->additions(), &hash_prefix_map);
+    // If getting the list of removals did not fail, keep processing the update.
     if (apply_update_result == APPLY_UPDATE_SUCCESS) {
-      apply_update_result = new_store->MergeUpdate(
-          hash_prefix_map_, hash_prefix_map, raw_removals);
+      apply_update_result = UpdateHashPrefixMapFromAdditions(
+          response->additions(), &hash_prefix_map);
+      if (apply_update_result == APPLY_UPDATE_SUCCESS) {
+        apply_update_result = new_store->MergeUpdate(
+            hash_prefix_map_, hash_prefix_map, raw_removals);
+      }
     }
     // TODO(vakh): Generate the updated ListUpdateResponse to write to disk.
   } else if (response->response_type() == ListUpdateResponse::FULL_UPDATE) {
@@ -160,20 +195,45 @@ void V4Store::ApplyUpdate(
 ApplyUpdateResult V4Store::UpdateHashPrefixMapFromAdditions(
     const RepeatedPtrField<ThreatEntrySet>& additions,
     HashPrefixMap* additions_map) {
+  ApplyUpdateResult apply_update_result;

[Nathan Parker, 2016/08/02 23:43:30]

This could go inside the for loop.  And actually you don't need it at all. 
Since you want to return early as soon as there's an error, you could just do
that at each point.  That'd reduce some of the if/else nesting.

[vakh (use Gerrit instead), 2016/08/03 01:44:48]

Done. Need to check return values from AddUnlumpedHashes.
Removed from other palces.

[Nathan Parker, 2016/08/03 21:16:37]

I don't see this change -- maybe I'm looking at the wrong version?

[vakh (use Gerrit instead), 2016/08/03 22:44:13]

apply_update_result can't be removed from this method.
This method can't directly return a ApplyUpdateResult value on failure, since it
needs to schedule a callback on the caller's thread.

[Nathan Parker, 2016/08/04 00:01:09]

I'm confused -- I don't see any callbacks in
V4Store::UpdateHashPrefixMapFromAdditions().

   for (const auto& addition : additions) {
-    // TODO(vakh): Allow other compression types.
-    // See: https://bugs.chromium.org/p/chromium/issues/detail?id=624567
-    DCHECK_EQ(RAW, addition.compression_type());
-
-    DCHECK(addition.has_raw_hashes());
-    DCHECK(addition.raw_hashes().has_raw_hashes());
+    const CompressionType compression_type = addition.compression_type();
+    if (compression_type == RAW) {
+      DCHECK(addition.has_raw_hashes());
+      DCHECK(addition.raw_hashes().has_raw_hashes());
+
+      PrefixSize prefix_size = addition.raw_hashes().prefix_size();

[Nathan Parker, 2016/08/02 23:43:30]

nit: I think this temporary doesn't add to the readability, esp since the member
name is already descriptive.

[vakh (use Gerrit instead), 2016/08/03 01:44:48]

Done.

+      apply_update_result = AddUnlumpedHashes(
+          prefix_size, addition.raw_hashes().raw_hashes(), additions_map);
+    } else if (compression_type == RICE) {
+      DCHECK(addition.has_rice_hashes());
+
+      const RiceDeltaEncoding& rice_hashes = addition.rice_hashes();
+      std::string raw_hashes_for_rice_encoded_hashes;

[Nathan Parker, 2016/08/02 23:43:30]

nit: Since this block of code is just for Rice encoding, you could use a shorter
var name (raw_hashes?).

[vakh (use Gerrit instead), 2016/08/03 01:44:49]

Done.

+      V4DecodeResult decode_result = V4RiceDecoder::DecodeBytes(
+          rice_hashes.first_value(), rice_hashes.rice_parameter(),
+          rice_hashes.num_entries(), rice_hashes.encoded_data(),
+          &raw_hashes_for_rice_encoded_hashes);
+      RecordDecodeAdditionsResult(decode_result);
+      if (decode_result != DECODE_SUCCESS) {
+        apply_update_result = RICE_DECODING_FAILURE;
+      } else {
+        // Rice-Golomb encoding is used to send compressed compressed 4-byte
+        // hash prefixes. Hash prefixes longer than 4 bytes will not be
+        // compressed, and will be served in raw format instead.
+        // Source: https://developers.google.com/safe-browsing/v4/compression
+        const PrefixSize kPrefixSize = 4;
+        apply_update_result = AddUnlumpedHashes(
+            kPrefixSize, raw_hashes_for_rice_encoded_hashes, additions_map);
+      }
+    } else {
+      apply_update_result = UNEXPECTED_COMPRESSION_TYPE_ADDITIONS_FAILURE;
+      NOTREACHED() << "Unexpected compression_type type: " << compression_type;
+    }
 
-    PrefixSize prefix_size = addition.raw_hashes().prefix_size();
-    ApplyUpdateResult result = AddUnlumpedHashes(
-        prefix_size, addition.raw_hashes().raw_hashes(), additions_map);
-    if (result != APPLY_UPDATE_SUCCESS) {
+    if (apply_update_result != APPLY_UPDATE_SUCCESS) {
       // If there was an error in updating the map, discard the update entirely.
-      return result;
+      return apply_update_result;
     }
   }
   return APPLY_UPDATE_SUCCESS;
@@ -337,7 +397,7 @@ ApplyUpdateResult V4Store::MergeUpdate(
 
   return (!raw_removals || removals_iter == raw_removals->end())
              ? APPLY_UPDATE_SUCCESS
-             : REMOVALS_INDEX_TOO_LARGE;
+             : REMOVALS_INDEX_TOO_LARGE_FAILURE;
 }
 
 StoreReadResult V4Store::ReadFromDisk() {
@@ -418,7 +478,6 @@ StoreWriteResult V4Store::WriteToDisk(
   DCHECK_EQ(file_format_string.size(), written);
 
   if (!base::Move(new_filename, store_path_)) {
-    DVLOG(1) << "store_path_: " << store_path_.value();
     return UNABLE_TO_RENAME_FAILURE;
   }