Merge the hash prefixes from the old store and the additions in the partial

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"
 #include "components/safe_browsing_db/v4_store.h"
 #include "components/safe_browsing_db/v4_store.pb.h"
 
 namespace safe_browsing {
 
 namespace {
 const uint32_t kFileMagic = 0x600D71FE;
 
 const uint32_t kFileVersion = 9;
 
+// The minimum expected size of a hash-prefix.

[Nathan Parker, 2016/07/12 20:35:05]

...size in bytes

[vakh (use Gerrit instead), 2016/07/12 21:57:02]

Done.

+const uint32_t kMinHashPrefixLength = 4;
+
+// The maximum expected size of a hash-prefix. This represents a full SHA256
+// hash.
+const uint32_t kMaxHashPrefixLength = 32;
+
 void RecordStoreReadResult(StoreReadResult result) {
   UMA_HISTOGRAM_ENUMERATION("SafeBrowsing.V4StoreReadResult", result,
                             STORE_READ_RESULT_MAX);
 }
 
 void RecordStoreWriteResult(StoreWriteResult result) {
   UMA_HISTOGRAM_ENUMERATION("SafeBrowsing.V4StoreWriteResult", result,
                             STORE_WRITE_RESULT_MAX);
 }
 
+void RecordMergeResult(MergeResult result) {
+  UMA_HISTOGRAM_ENUMERATION("SafeBrowsing.V4MergeResult", result,
+                            MERGE_RESULT_MAX);
+}
+
 // 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
 
 std::ostream& operator<<(std::ostream& os, const V4Store& store) {
   os << store.DebugString();
@@ skipping 35 matching lines @@
   state_ = "";
   return true;
 }
 
 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.
   // 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.
-  if (response->has_response_type() &&
-      response->response_type() == ListUpdateResponse::FULL_UPDATE) {
+  if (response->response_type() == ListUpdateResponse::PARTIAL_UPDATE) {
+    for (const auto& removal : response->removals()) {
+      // TODO(vakh): Allow other compression types.
+      // See: https://bugs.chromium.org/p/chromium/issues/detail?id=624567
+      DCHECK_EQ(RAW, removal.compression_type());
+    }
+
+    HashPrefixMap additions_map;
+    UpdateHashPrefixMapFromAdditions(response->additions(), &additions_map);
+
+    new_store->MergeUpdate(hash_prefix_map_, additions_map);
+
+    // TODO(vakh): Generate the updated ListUpdateResponse to write to disk.
+  } else if (response->response_type() == ListUpdateResponse::FULL_UPDATE) {
     StoreWriteResult result = new_store->WriteToDisk(std::move(response));
     RecordStoreWriteResult(result);
+  } else {
+    NOTREACHED() << "Unexpected response type: " << response->response_type();
   }
 
   // new_store is done updating, pass it to the callback.
   callback_task_runner->PostTask(
       FROM_HERE, base::Bind(callback, base::Passed(&new_store)));
 }
 
+// static
+void V4Store::UpdateHashPrefixMapFromAdditions(
+    const ::google::protobuf::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();
+    RecordMergeResult(AddUnlumpedHashes(
+        prefix_size, addition.raw_hashes().raw_hashes(), additions_map));
+  }
+}
+
+// static
+MergeResult V4Store::AddUnlumpedHashes(PrefixSize prefix_size,
+                                       const std::string& lumped_hashes,
+                                       HashPrefixMap* additions_map) {
+  if (prefix_size < kMinHashPrefixLength) {
+    NOTREACHED();
+    return PREFIX_SIZE_TOO_SMALL_FAILURE;

[Nathan Parker, 2016/07/12 20:35:06]

You don't need any of the else's here.

[vakh (use Gerrit instead), 2016/07/12 21:57:02]

Done.

+  } else if (prefix_size > kMaxHashPrefixLength) {
+    NOTREACHED();
+    return PREFIX_SIZE_TOO_LARGE_FAILURE;
+  } else {
+    if (lumped_hashes.size() % prefix_size != 0) {
+      return ADDITIONS_SIZE_UNEXPECTED_FAILURE;
+    } else {
+      (*additions_map)[prefix_size] = lumped_hashes;

[Nathan Parker, 2016/07/12 20:35:06]

To consider, or add a todo: This copies the string unnecessarily.  If you
guaranteed that the proto's lived at least as long as the update process, you
could store a pointed to the proto's string.

[Nathan Parker, 2016/07/12 22:47:12]

(Missed this comment) Though, I'm not sure how to implement this easily, since
it'd require HashPrefixMap to not own the data in some cases (additions) and own
it in others (hash_prefix_map_).  Maybe just a TODO for now.  This probably only
really matters at startup -- For a 2MB list loaded from disk, this'll use a
extra 2MB RAM.

[vakh (use Gerrit instead), 2016/07/13 00:29:17]

Done.

+    }
+  }
+  return MERGE_SUCCESS;
+}
+
+// static
+bool V4Store::GetNextSmallestPrefixSize(const HashPrefixMap& hash_prefix_map,
+                                        const CounterMap& counter_map,
+                                        PrefixSize* smallest_prefix_size) {
+  HashPrefix smallest_prefix;

[Nathan Parker, 2016/07/12 20:35:06]

You could use a StringPiece here and within the loop so it doesn't actually copy
the substring.

[vakh (use Gerrit instead), 2016/07/12 21:57:02]

That would break the consistency of hash prefixes being HashPrefixes and expose
the implementation detail about a hash prefix being a string.
These strings are fairly small so copying them isn't worth the sacrifice in
readability.

[Nathan Parker, 2016/07/12 22:47:12]

They are small, but there is a dynamic allocation everytime one gets copied. 
Actually, just pulling it out of the loop would fix most of that, since then it
could reuse the same buffer over and over.

[vakh (use Gerrit instead), 2016/07/13 00:29:17]

Done.

+  for (const auto& counter_pair : counter_map) {
+    PrefixSize prefix_size = counter_pair.first;
+    size_t index = counter_pair.second;
+    size_t sized_index = prefix_size * index;
+
+    const HashPrefixes& hash_prefixes = hash_prefix_map.at(prefix_size);
+    if (sized_index < hash_prefixes.size()) {
+      HashPrefix this_prefix = hash_prefixes.substr(sized_index, prefix_size);
+      if (smallest_prefix.empty() || smallest_prefix > this_prefix) {
+        smallest_prefix = this_prefix;
+        *smallest_prefix_size = prefix_size;
+      }
+    }

[Nathan Parker, 2016/07/12 20:35:06]

else dcheck

[vakh (use Gerrit instead), 2016/07/12 21:57:02]

That's a valid case. For instance, when one map has been merged entirely. In
that case, all sized_index values would be outside the valid range for all
corresponding HashPrefixes.

+  }
+  return !smallest_prefix.empty();
+}
+
+// static
+void V4Store::InitializeCounterMap(const HashPrefixMap& hash_prefix_map,
+                                   CounterMap* counter_map) {
+  for (const auto& map_pair : hash_prefix_map) {
+    (*counter_map)[map_pair.first] = 0;
+  }
+}
+
+// static
+HashPrefix V4Store::GetNextUnmergedPrefixForSize(
+    PrefixSize prefix_size,
+    const HashPrefixMap& hash_prefix_map,
+    const CounterMap& counter_map) {
+  HashPrefixes hash_prefixes = hash_prefix_map.at(prefix_size);

[Nathan Parker, 2016/07/12 20:35:06]

This copies the whole string.  You could use const HashPrefixes&.

[vakh (use Gerrit instead), 2016/07/12 21:57:03]

Done.

+  size_t index_within_list = counter_map.at(prefix_size);
+  size_t sized_index = prefix_size * index_within_list;
+  return hash_prefixes.substr(sized_index, prefix_size);
+}
+
+// static
+void V4Store::ReserveSpaceInPrefixMap(const HashPrefixMap& other_prefixes_map,
+                                      HashPrefixMap* prefix_map_to_update) {
+  for (const auto& pair : other_prefixes_map) {

[Nathan Parker, 2016/07/12 20:35:06]

Add a comment about how this gets close to ideal, but will leave a bit of space
due to deletions.

[vakh (use Gerrit instead), 2016/07/12 21:57:02]

Updated the comment in header file (consistent with the other functions).

+    PrefixSize prefix_size = pair.first;
+    size_t prefix_length_to_add = pair.second.length();
+
+    HashPrefixes existing_prefixes = (*prefix_map_to_update)[prefix_size];

[Nathan Parker, 2016/07/12 20:35:06]

const HashPrefixes&.  Otherwise it makes a copy.

[vakh (use Gerrit instead), 2016/07/12 21:57:02]

Done.

+    size_t existing_capacity = existing_prefixes.capacity();
+
+    (*prefix_map_to_update)[prefix_size].reserve(existing_capacity +
+                                                 prefix_length_to_add);
+  }
+}
+
+void V4Store::MergeUpdate(const HashPrefixMap& old_prefixes_map,
+                          const HashPrefixMap& additions_map) {
+  ReserveSpaceInPrefixMap(old_prefixes_map, &hash_prefix_map_);

[Nathan Parker, 2016/07/12 20:35:06]

Is the assumption that hash_prefix_map_ is empty here?  Otherwise, it'll add
old+additions to current.  Maybe making this one call (that takes both as args)
could enforce that, by not using the existing capacity.

[vakh (use Gerrit instead), 2016/07/12 21:57:02]

Good map. However, for code reuse purposes, I think it is better to have a
method that method as-is.
Clearing out the hash_prefix_map_ explicitly and added a DCHECK.

+  ReserveSpaceInPrefixMap(additions_map, &hash_prefix_map_);
+
+  PrefixSize next_size_for_old;
+  CounterMap old_counter_map;
+  InitializeCounterMap(old_prefixes_map, &old_counter_map);
+  bool found_in_old = GetNextSmallestPrefixSize(
+      old_prefixes_map, old_counter_map, &next_size_for_old);
+
+  PrefixSize next_size_for_additions;
+  CounterMap additions_counter_map;
+  InitializeCounterMap(additions_map, &additions_counter_map);
+  bool found_in_additions = GetNextSmallestPrefixSize(
+      additions_map, additions_counter_map, &next_size_for_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.
+  while (found_in_old || found_in_additions) {
+    // Both maps still have elements that need to be merged into the new store.
+    if (found_in_old && found_in_additions) {
+      // Get the lexographically smallest hash prefix from the old store.
+      HashPrefix next_smallest_prefix_old = GetNextUnmergedPrefixForSize(
+          next_size_for_old, old_prefixes_map, old_counter_map);
+      // Get the lexographically smallest hash prefix from the additions in the
+      // latest update from the server.
+      HashPrefix next_smallest_prefix_additions = GetNextUnmergedPrefixForSize(
+          next_size_for_additions, additions_map, additions_counter_map);
+
+      // If the smallest unmerged hash prefix in old store is smaller than that
+      // in the update, add that to the new store.
+      if (next_smallest_prefix_old < next_smallest_prefix_additions) {
+        hash_prefix_map_[next_size_for_old] += next_smallest_prefix_old;
+
+        // Update the counter map, which means that we have merged one hash
+        // prefix of size |next_size_for_old| from the old store.
+        old_counter_map[next_size_for_old]++;
+
+        // Find the next smallest unmerged element in the old store's map.
+        found_in_old = GetNextSmallestPrefixSize(
+            old_prefixes_map, old_counter_map, &next_size_for_old);
+      } else {
+        // This means that the smallest unmerged hash prefix in the update was
+        // smaller than that in the old store, so add that hash prefix (from the
+        // update) into the new store.
+        hash_prefix_map_[next_size_for_additions] +=
+            next_smallest_prefix_additions;
+
+        // Update the counter map, which means that we have merged one hash
+        // prefix of size |next_size_for_additions| from the update.
+        additions_counter_map[next_size_for_additions]++;
+
+        // Find the next smallest unmerged element in the additions map.
+        found_in_additions = GetNextSmallestPrefixSize(
+            additions_map, additions_counter_map, &next_size_for_additions);
+      }
+    } else {
+      // At least one of the maps has become empty.
+      if (!found_in_old) {
+        // The map for the old store has been completely merged. Now only merge
+        // the hash prefixes from the additions map.
+        HashPrefix next_smallest_prefix_additions =
+            GetNextUnmergedPrefixForSize(next_size_for_additions, additions_map,
+                                         additions_counter_map);
+        hash_prefix_map_[next_size_for_additions] +=
+            next_smallest_prefix_additions;
+        additions_counter_map[next_size_for_additions]++;
+        found_in_additions = GetNextSmallestPrefixSize(
+            additions_map, additions_counter_map, &next_size_for_additions);
+      } else {
+        // The additions map has been completely merged. Now only merge
+        // the hash prefixes from the old store map.
+        HashPrefix next_smallest_prefix_old = GetNextUnmergedPrefixForSize(
+            next_size_for_old, old_prefixes_map, old_counter_map);
+        hash_prefix_map_[next_size_for_old] += next_smallest_prefix_old;
+        old_counter_map[next_size_for_old]++;
+        found_in_old = GetNextSmallestPrefixSize(
+            old_prefixes_map, old_counter_map, &next_size_for_old);
+      }
+    }
+  }
+}
+
 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;
   }
 
   if (contents.empty()) {
@@ skipping 60 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