PVer4: Drop hash prefixes based on the removals field in the response

components/safe_browsing_db/v4_store.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/base64.h"
 #include "base/bind.h"
 #include "base/files/file_util.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/metrics/sparse_histogram.h"
 #include "base/strings/stringprintf.h"
@@ skipping 36 matching lines @@
 }
 
 // Returns the name of the temporary file used to buffer data for
 // |filename|.  Exported for unit tests.
 const base::FilePath TemporaryFileForFilename(const base::FilePath& filename) {
   return base::FilePath(filename.value() + FILE_PATH_LITERAL("_new"));
 }
 
 }  // namespace
 
+using ::google::protobuf::RepeatedField;
+using ::google::protobuf::RepeatedPtrField;
+using ::google::protobuf::int32;
+
 std::ostream& operator<<(std::ostream& os, const V4Store& store) {
   os << store.DebugString();
   return os;
 }
 
 V4Store* V4StoreFactory::CreateV4Store(
     const scoped_refptr<base::SequencedTaskRunner>& task_runner,
     const base::FilePath& store_path) {
   V4Store* new_store = new V4Store(task_runner, store_path);
   new_store->Initialize();
@@ skipping 29 matching lines @@
 }
 
 void V4Store::ApplyUpdate(
     std::unique_ptr<ListUpdateResponse> response,
     const scoped_refptr<base::SingleThreadTaskRunner>& callback_task_runner,
     UpdatedStoreReadyCallback callback) {
   std::unique_ptr<V4Store> new_store(
       new V4Store(this->task_runner_, this->store_path_));
   new_store->state_ = response->new_client_state();
   // TODO(vakh):
-  // 1. Merge the old store and the new update in new_store.
+  // 1. Done: Merge the old store and the new update in new_store.
   // 2. Create a ListUpdateResponse containing RICE encoded hash-prefixes and
   // 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;
   if (response->response_type() == ListUpdateResponse::PARTIAL_UPDATE) {
-    for (const auto& removal : response->removals()) {
+    const RepeatedField<int32>* raw_removals = nullptr;
+    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();
     }
 
     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);
+      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) {
     apply_update_result = UpdateHashPrefixMapFromAdditions(
         response->additions(), &hash_prefix_map);
     if (apply_update_result == APPLY_UPDATE_SUCCESS) {
       new_store->hash_prefix_map_ = hash_prefix_map;
       RecordStoreWriteResult(new_store->WriteToDisk(std::move(response)));
     }
   } else {
     apply_update_result = UNEXPECTED_RESPONSE_TYPE_FAILURE;
     NOTREACHED() << "Unexpected response type: " << response->response_type();
   }
 
   if (apply_update_result == APPLY_UPDATE_SUCCESS) {
     // new_store is done updating, pass it to the callback.
     callback_task_runner->PostTask(
         FROM_HERE, base::Bind(callback, base::Passed(&new_store)));
   } else {
     // new_store failed updating. Pass a nullptr to the callback.
     callback_task_runner->PostTask(FROM_HERE, base::Bind(callback, nullptr));
   }
 
   RecordApplyUpdateResult(apply_update_result);
 }
 
 // static
 ApplyUpdateResult V4Store::UpdateHashPrefixMapFromAdditions(
-    const ::google::protobuf::RepeatedPtrField<ThreatEntrySet>& additions,
+    const RepeatedPtrField<ThreatEntrySet>& additions,
     HashPrefixMap* additions_map) {
   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());
 
     PrefixSize prefix_size = addition.raw_hashes().prefix_size();
@@ skipping 68 matching lines @@
 
     const HashPrefixes& existing_prefixes =
         (*prefix_map_to_update)[prefix_size];
     size_t existing_capacity = existing_prefixes.capacity();
 
     (*prefix_map_to_update)[prefix_size].reserve(existing_capacity +
                                                  prefix_length_to_add);
   }
 }
 
-ApplyUpdateResult V4Store::MergeUpdate(const HashPrefixMap& old_prefixes_map,
-                                       const HashPrefixMap& additions_map) {
+ApplyUpdateResult V4Store::MergeUpdate(
+    const HashPrefixMap& old_prefixes_map,
+    const HashPrefixMap& additions_map,
+    const RepeatedField<int32>* raw_removals) {
   DCHECK(hash_prefix_map_.empty());
   hash_prefix_map_.clear();
   ReserveSpaceInPrefixMap(old_prefixes_map, &hash_prefix_map_);
   ReserveSpaceInPrefixMap(additions_map, &hash_prefix_map_);
 
   IteratorMap old_iterator_map;
   HashPrefix next_smallest_prefix_old;
   InitializeIteratorMap(old_prefixes_map, &old_iterator_map);
   bool old_has_unmerged = GetNextSmallestUnmergedPrefix(
       old_prefixes_map, old_iterator_map, &next_smallest_prefix_old);
 
   IteratorMap additions_iterator_map;
   HashPrefix next_smallest_prefix_additions;
   InitializeIteratorMap(additions_map, &additions_iterator_map);
   bool additions_has_unmerged = GetNextSmallestUnmergedPrefix(
       additions_map, additions_iterator_map, &next_smallest_prefix_additions);
 
   // Classical merge sort.
   // The two constructs to merge are maps: old_prefixes_map, additions_map.
   // At least one of the maps still has elements that need to be merged into the
   // new store.
+
+  // Keep track of the number of elements picked from the old map. This is used
+  // to determine which elements to drop based on the raw_removals. Note that
+  // picked is not the same as merged. A picked element isn't merged if its
+  // index is on the raw_removals list.
+  int total_picked_from_old = 0;
+  const int* removals_iter = raw_removals ? raw_removals->begin() : nullptr;
   while (old_has_unmerged || additions_has_unmerged) {
     // If the same hash prefix appears in the existing store and the additions
     // list, something is clearly wrong. Discard the update.
     if (old_has_unmerged && additions_has_unmerged &&
         next_smallest_prefix_old == next_smallest_prefix_additions) {
       return ADDITIONS_HAS_EXISTING_PREFIX_FAILURE;
     }
 
     // Select which map to pick the next hash prefix from to keep the result in
     // lexographically sorted order.
     bool pick_from_old =
         old_has_unmerged &&
         (!additions_has_unmerged ||
          (next_smallest_prefix_old < next_smallest_prefix_additions));
 
     PrefixSize next_smallest_prefix_size;
     if (pick_from_old) {
       next_smallest_prefix_size = next_smallest_prefix_old.size();
 
-      // Append the smallest hash to the appropriate list.
-      hash_prefix_map_[next_smallest_prefix_size] += next_smallest_prefix_old;
-
       // Update the iterator map, which means that we have merged one hash
       // prefix of size |next_size_for_old| from the old store.
       old_iterator_map[next_smallest_prefix_size] += next_smallest_prefix_size;
 
+      if (!raw_removals || removals_iter == raw_removals->end() ||
+          *removals_iter != total_picked_from_old) {
+        // Append the smallest hash to the appropriate list.
+        hash_prefix_map_[next_smallest_prefix_size] += next_smallest_prefix_old;
+      } else {
+        // Element not added to new map. Move the removals iterator forward.
+        removals_iter++;
+      }
+
+      total_picked_from_old++;
+
       // Find the next smallest unmerged element in the old store's map.
       old_has_unmerged = GetNextSmallestUnmergedPrefix(
           old_prefixes_map, old_iterator_map, &next_smallest_prefix_old);
     } else {
       next_smallest_prefix_size = next_smallest_prefix_additions.size();
 
       // Append the smallest hash to the appropriate list.
       hash_prefix_map_[next_smallest_prefix_size] +=
           next_smallest_prefix_additions;
 
       // Update the iterator map, which means that we have merged one hash
       // prefix of size |next_smallest_prefix_size| from the update.
       additions_iterator_map[next_smallest_prefix_size] +=
           next_smallest_prefix_size;
 
       // Find the next smallest unmerged element in the additions map.
       additions_has_unmerged =
           GetNextSmallestUnmergedPrefix(additions_map, additions_iterator_map,
                                         &next_smallest_prefix_additions);
     }
   }
 
-  return APPLY_UPDATE_SUCCESS;
+  return (!raw_removals || removals_iter == raw_removals->end())
+             ? APPLY_UPDATE_SUCCESS
+             : REMOVALS_INDEX_TOO_LARGE;
 }
 
 StoreReadResult V4Store::ReadFromDisk() {
   DCHECK(task_runner_->RunsTasksOnCurrentThread());
 
   std::string contents;
   bool read_success = base::ReadFileToString(store_path_, &contents);
   if (!read_success) {
     return FILE_UNREADABLE_FAILURE;
   }
@@ skipping 68 matching lines @@
 
   if (!base::Move(new_filename, store_path_)) {
     DVLOG(1) << "store_path_: " << store_path_.value();
     return UNABLE_TO_RENAME_FAILURE;
   }
 
   return WRITE_SUCCESS;
 }
 
 }  // namespace safe_browsing