PVer4: Verify checksum for downloaded updates

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/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"
+#include "crypto/sha2.h"
 
 namespace safe_browsing {
 
 namespace {
 const uint32_t kFileMagic = 0x600D71FE;
 
 const uint32_t kFileVersion = 9;
 
 // The minimum expected size (in bytes) of a hash-prefix.
 const uint32_t kMinHashPrefixLength = 4;
@@ skipping 81 matching lines @@
   return base::StringPrintf("path: %" PRIsFP "; state: %s",
                             store_path_.value().c_str(), state_base64.c_str());
 }
 
 bool V4Store::Reset() {
   // TODO(vakh): Implement skeleton.
   state_ = "";
   return true;
 }
 
-ApplyUpdateResult V4Store::ProcessFullUpdate(
+// static
+ApplyUpdateResult V4Store::ProcessPartialUpdateAndWriteToDisk(
+    const HashPrefixMap& hash_prefix_map_old,
     std::unique_ptr<ListUpdateResponse> response,
     const std::unique_ptr<V4Store>& new_store) {
-  HashPrefixMap hash_prefix_map;
-  ApplyUpdateResult apply_update_result =
-      UpdateHashPrefixMapFromAdditions(response->additions(), &hash_prefix_map);
-  if (apply_update_result == APPLY_UPDATE_SUCCESS) {
-    new_store->hash_prefix_map_ = hash_prefix_map;
+  DCHECK(response->has_response_type());
+  DCHECK_EQ(ListUpdateResponse::PARTIAL_UPDATE, response->response_type());
+
+  ApplyUpdateResult result =
+      ProcessUpdate(hash_prefix_map_old, response, new_store.get());

[Nathan Parker, 2016/08/08 21:17:19]

This change removes the nicely optimized short-cut you had for full-updates. 
Before, we'd just copy the whole thing in one motion. Now we have to iterate
through the whole map and append hashes one at a time in lex-order.  It looks
like you do need to iterate in lex order to compute the checksum, but we could
save some copying time by copying in large blocks rather than many small ones.

I think that'd save some CPU, but it shouldn't change memory usage.  So I'd be
OK with a TODO here suggesting we could optimize the full updates if needed but
leave the work till we have some measurements.

[vakh (use Gerrit instead), 2016/08/08 22:46:47]

Done.

+  if (result == APPLY_UPDATE_SUCCESS) {
+    // TODO(vakh): Create a ListUpdateResponse containing RICE encoded
+    // hash prefixes and response_type as FULL_UPDATE, and write that to disk.
+  }
+  return result;
+}
+
+// static
+ApplyUpdateResult V4Store::ProcessFullUpdateAndWriteToDisk(
+    std::unique_ptr<ListUpdateResponse> response,
+    const std::unique_ptr<V4Store>& new_store) {
+  ApplyUpdateResult result = ProcessFullUpdate(response, new_store.get());
+  if (result == APPLY_UPDATE_SUCCESS) {
     RecordStoreWriteResult(new_store->WriteToDisk(std::move(response)));
   }
-  return apply_update_result;
+  return result;
 }
 
-ApplyUpdateResult V4Store::ProcessPartialUpdate(
-    std::unique_ptr<ListUpdateResponse> response,
-    const std::unique_ptr<V4Store>& new_store) {
-  // TODO(vakh):
-  // 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.
+// static
+ApplyUpdateResult V4Store::ProcessFullUpdate(
+    const std::unique_ptr<ListUpdateResponse>& response,
+    V4Store* new_store) {
+  DCHECK(response->has_response_type());
+  DCHECK_EQ(ListUpdateResponse::FULL_UPDATE, response->response_type());
+  return ProcessUpdate(HashPrefixMap(), response, new_store);
+}
 
+// static
+ApplyUpdateResult V4Store::ProcessUpdate(
+    const HashPrefixMap& hash_prefix_map_old,
+    const std::unique_ptr<ListUpdateResponse>& response,
+    V4Store* new_store) {
   const RepeatedField<int32>* raw_removals = nullptr;
   RepeatedField<int32> rice_removals;
   size_t removals_size = response->removals_size();
   DCHECK_LE(removals_size, 1u);
   if (removals_size == 1) {
     const ThreatEntrySet& removal = response->removals().Get(0);
     const CompressionType compression_type = removal.compression_type();
     if (compression_type == RAW) {
       raw_removals = &removal.raw_indices().indices();
     } else if (compression_type == RICE) {
@@ skipping 12 matching lines @@
       }
     } else {
       NOTREACHED() << "Unexpected compression_type type: " << compression_type;
       return UNEXPECTED_COMPRESSION_TYPE_REMOVALS_FAILURE;
     }
   }
 
   HashPrefixMap hash_prefix_map;
   ApplyUpdateResult apply_update_result =
       UpdateHashPrefixMapFromAdditions(response->additions(), &hash_prefix_map);
+  if (apply_update_result != APPLY_UPDATE_SUCCESS) {
+    return apply_update_result;
+  }
 
-  if (apply_update_result == APPLY_UPDATE_SUCCESS) {
-    apply_update_result =
-        new_store->MergeUpdate(hash_prefix_map_, hash_prefix_map, raw_removals);
+  std::string expected_checksum;
+  if (response->has_checksum() && response->checksum().has_sha256()) {
+    expected_checksum = response->checksum().sha256();
   }
-  return apply_update_result;
+
+  apply_update_result = new_store->MergeUpdate(
+      hash_prefix_map_old, hash_prefix_map, raw_removals, expected_checksum);
+  if (apply_update_result != APPLY_UPDATE_SUCCESS) {
+    return apply_update_result;
+  }
+
+  new_store->state_ = response->new_client_state();
+  return APPLY_UPDATE_SUCCESS;
 }
 
 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();
 
   ApplyUpdateResult apply_update_result;
   if (response->response_type() == ListUpdateResponse::PARTIAL_UPDATE) {
-    apply_update_result = ProcessPartialUpdate(std::move(response), new_store);
+    apply_update_result = ProcessPartialUpdateAndWriteToDisk(
+        hash_prefix_map_, std::move(response), new_store);
   } else if (response->response_type() == ListUpdateResponse::FULL_UPDATE) {
-    apply_update_result = ProcessFullUpdate(std::move(response), new_store);
+    apply_update_result =
+        ProcessFullUpdateAndWriteToDisk(std::move(response), new_store);
   } 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 {
+    DVLOG(1) << "ApplyUpdate failed: " << *this;

[Nathan Parker, 2016/08/08 21:17:19]

Make it a single DVLOG.

[vakh (use Gerrit instead), 2016/08/08 22:46:47]

Done.

+    DVLOG(1) << "reason: " << apply_update_result;
     // 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 RepeatedPtrField<ThreatEntrySet>& additions,
@@ skipping 103 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,
-    const RepeatedField<int32>* raw_removals) {
+ApplyUpdateResult V4Store::MergeUpdate(const HashPrefixMap& old_prefixes_map,
+                                       const HashPrefixMap& additions_map,
+                                       const RepeatedField<int32>* raw_removals,
+                                       const std::string& expected_checksum) {
   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.
 
+  // |all_prefixes_concatenated| stores the concatenated list of hash prefixes
+  // in lexographically sorted order. It is used to calculate the |checksum| at
+  // the end. This checksum is matched against the expected checksum sent by
+  // the server.
+  HashPrefixes all_prefixes_concatenated;
+  bool calculate_checksum = !expected_checksum.empty();
+
   // 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 &&
@@ skipping 13 matching lines @@
       next_smallest_prefix_size = next_smallest_prefix_old.size();
 
       // 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;
+
+        if (calculate_checksum) {
+          all_prefixes_concatenated += next_smallest_prefix_old;

[Nathan Parker, 2016/08/08 21:17:19]

Rather than building up another whole copy of the hashes in memory, you should
be able to compute the SHA256 incrementally.

[vakh (use Gerrit instead), 2016/08/08 22:46:47]

Done.

+        }
       } 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;
 
+      if (calculate_checksum) {
+        all_prefixes_concatenated += 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);
     }
   }
 
+  if (calculate_checksum) {
+    std::string checksum = crypto::SHA256HashString(all_prefixes_concatenated);
+    if (checksum != expected_checksum) {
+      std::string checksum_base64, expected_checksum_base64;
+      base::Base64Encode(checksum, &checksum_base64);
+      base::Base64Encode(expected_checksum, &expected_checksum_base64);
+      DVLOG(1) << "calculated checksum: " << checksum_base64;

[Nathan Parker, 2016/08/08 21:17:19]

How about making it a single DVLOG, with an error message ("Checksum failed...")

[vakh (use Gerrit instead), 2016/08/08 22:46:47]

Done.

+      DVLOG(1) << "expected checksum: " << expected_checksum_base64;
+      return CHECKSUM_MISMATCH_FAILURE;
+    }
+  }
+
   return (!raw_removals || removals_iter == raw_removals->end())

[Nathan Parker, 2016/08/08 21:17:19]

It might be useful to check (!raw_removals || removals_iter ==
raw_removals->end()) BEFORE checking the checksum, since if the removals weren't
processed correctly it'll just show up as a checksum failure otherwise.

[vakh (use Gerrit instead), 2016/08/08 22:46:47]

Done.

              ? APPLY_UPDATE_SUCCESS
              : REMOVALS_INDEX_TOO_LARGE_FAILURE;
 }
 
 StoreReadResult V4Store::ReadFromDisk() {
   DCHECK(task_runner_->RunsTasksOnCurrentThread());
 
   std::string contents;
   bool read_success = base::ReadFileToString(store_path_, &contents);
   if (!read_success) {
@@ skipping 20 matching lines @@
   if (file_format.version_number() != kFileVersion) {
     DVLOG(1) << "File version incompatible: " << file_format.version_number()
              << "; expected: " << kFileVersion;
     return FILE_VERSION_INCOMPATIBLE_FAILURE;
   }
 
   if (!file_format.has_list_update_response()) {
     return HASH_PREFIX_INFO_MISSING_FAILURE;
   }
 
-  const ListUpdateResponse& response = file_format.list_update_response();
-  ApplyUpdateResult apply_update_result = UpdateHashPrefixMapFromAdditions(
-      response.additions(), &hash_prefix_map_);
+  std::unique_ptr<ListUpdateResponse> response(new ListUpdateResponse);
+  response->Swap(file_format.mutable_list_update_response());
+  ApplyUpdateResult apply_update_result = ProcessFullUpdate(response, this);
   RecordApplyUpdateResultWhenReadingFromDisk(apply_update_result);
   if (apply_update_result != APPLY_UPDATE_SUCCESS) {
     hash_prefix_map_.clear();
     return HASH_PREFIX_MAP_GENERATION_FAILURE;
   }
 
-  state_ = response.new_client_state();
   return READ_SUCCESS;
 }
 
 StoreWriteResult V4Store::WriteToDisk(
     std::unique_ptr<ListUpdateResponse> response) const {
   // Do not write partial updates to the disk.
   // After merging the updates, the ListUpdateResponse passed to this method
   // should be a FULL_UPDATE.
   if (!response->has_response_type() ||
       response->response_type() != ListUpdateResponse::FULL_UPDATE) {
@@ skipping 59 matching lines @@
   if (result == 0) {
     return true;
   } else if (result < 0) {
     return HashPrefixMatches(hash_prefix, begin, mid);
   } else {
     return HashPrefixMatches(hash_prefix, mid + prefix_size, end);
   }
 }
 
 }  // namespace safe_browsing