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/stl_util.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/secure_hash.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(
-    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;
-    RecordStoreWriteResult(new_store->WriteToDisk(std::move(response)));
+ApplyUpdateResult V4Store::ProcessPartialUpdateAndWriteToDisk(
+    const HashPrefixMap& hash_prefix_map_old,
+    std::unique_ptr<ListUpdateResponse> response) {
+  DCHECK(response->has_response_type());
+  DCHECK_EQ(ListUpdateResponse::PARTIAL_UPDATE, response->response_type());
+
+  ApplyUpdateResult result = ProcessUpdate(hash_prefix_map_old, response);
+  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 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.
+ApplyUpdateResult V4Store::ProcessFullUpdateAndWriteToDisk(
+    std::unique_ptr<ListUpdateResponse> response) {
+  ApplyUpdateResult result = ProcessFullUpdate(response);
+  if (result == APPLY_UPDATE_SUCCESS) {
+    RecordStoreWriteResult(WriteToDisk(std::move(response)));
+  }
+  return result;
+}
 
+ApplyUpdateResult V4Store::ProcessFullUpdate(
+    const std::unique_ptr<ListUpdateResponse>& response) {
+  DCHECK(response->has_response_type());
+  DCHECK_EQ(ListUpdateResponse::FULL_UPDATE, response->response_type());
+  // TODO(vakh): For a full update, we don't need to process the update in
+  // lexographical order to store it, but we do need to do that for calculating

[noé, 2016/08/11 08:15:25]

I believe the right spelling is lexicographical :). Here and below.

[vakh (use Gerrit instead), 2016/08/11 08:31:20]

Thanks. I'll fix that in the next CL.

+  // checksum. It might save some CPU cycles to store the full update as-is and
+  // walk the list of hash prefixes in lexographical order only for checksum
+  // calculation.
+  return ProcessUpdate(HashPrefixMap(), response);
+}
+
+ApplyUpdateResult V4Store::ProcessUpdate(
+    const HashPrefixMap& hash_prefix_map_old,
+    const std::unique_ptr<ListUpdateResponse>& response) {
   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 = MergeUpdate(hash_prefix_map_old, hash_prefix_map,
+                                    raw_removals, expected_checksum);
+  if (apply_update_result != APPLY_UPDATE_SUCCESS) {
+    return apply_update_result;
+  }
+
+  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 = new_store->ProcessPartialUpdateAndWriteToDisk(
+        hash_prefix_map_, std::move(response));
   } else if (response->response_type() == ListUpdateResponse::FULL_UPDATE) {
-    apply_update_result = ProcessFullUpdate(std::move(response), new_store);
+    apply_update_result =
+        new_store->ProcessFullUpdateAndWriteToDisk(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 {
+    DVLOG(1) << "ApplyUpdate failed: reason: " << apply_update_result
+             << "; store: " << *this;
     // 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.
 
+  bool calculate_checksum = !expected_checksum.empty();
+  std::unique_ptr<crypto::SecureHash> checksum_ctx(
+      crypto::SecureHash::Create(crypto::SecureHash::SHA256));
+
   // 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) {

[noé, 2016/08/11 08:15:25]

This is quite complicated code. Have you considered putting all the prefixes
into a set<string> and then applying additions and removals. At the end
converting it back to a HashPrefixMap?

[vakh (use Gerrit instead), 2016/08/11 08:31:20]

Yes, it would simplify the code but at the cost of increasing time complexity.
Using a set would cost O(N*log(N)) time, whereas the current implementation
costs O(N).

[vakh (use Gerrit instead), 2016/08/11 17:05:59]

Sorry, I take that back. I think the time complexity is probably the same.

However, creating a set means creating a 32 byte string object for each of the
hash prefixes, most of which are 4-bytes long. This would end up taking 8 times
more memory, even though it would be for a small duration. That's what I was
trying to avoid.

     // 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();
 
       // 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) {
+          checksum_ctx->Update(string_as_array(&next_smallest_prefix_old),
+                               next_smallest_prefix_size);
+        }
       } 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) {
+        checksum_ctx->Update(string_as_array(&next_smallest_prefix_additions),
+                             next_smallest_prefix_size);
+      }
+
       // 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 (!raw_removals || removals_iter == raw_removals->end())
-             ? APPLY_UPDATE_SUCCESS
-             : REMOVALS_INDEX_TOO_LARGE_FAILURE;
+  if (raw_removals && removals_iter != raw_removals->end()) {
+    return REMOVALS_INDEX_TOO_LARGE_FAILURE;
+  }
+
+  if (calculate_checksum) {
+    std::string checksum(crypto::kSHA256Length, 0);
+    checksum_ctx->Finish(string_as_array(&checksum), checksum.size());
+    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) << "Checksum failed: calculated: " << checksum_base64
+               << "expected: " << expected_checksum_base64;
+      return CHECKSUM_MISMATCH_FAILURE;
+    }
+  }
+
+  return APPLY_UPDATE_SUCCESS;
 }
 
 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 18 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);
   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