PVer4: Test checksum on startup outside the hotpath of DB load

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 "content/public/browser/browser_thread.h"
 #include "crypto/secure_hash.h"
 #include "crypto/sha2.h"
 
 using base::TimeTicks;
 
 namespace safe_browsing {
 
 namespace {
 
 const uint32_t kFileMagic = 0x600D71FE;
 
 const uint32_t kFileVersion = 9;
 
 std::string GetUmaSuffixForStore(const base::FilePath& file_path) {
   return base::StringPrintf(
       ".%" PRIsFP, file_path.BaseName().RemoveExtension().value().c_str());
 }
 
 void RecordTimeWithAndWithoutStore(const std::string& metric,
                                    base::TimeDelta time,
                                    const base::FilePath& file_path) {
-  std::string suffix = GetUmaSuffixForStore(file_path);
-
   // The histograms below are a modified expansion of the
   // UMA_HISTOGRAM_LONG_TIMES macro adapted to allow for a dynamically suffixed
   // histogram name.
   // Note: The factory creates and owns the histogram.
   const int kBucketCount = 100;
   base::HistogramBase* histogram = base::Histogram::FactoryTimeGet(
       metric, base::TimeDelta::FromMilliseconds(1),
       base::TimeDelta::FromMinutes(1), kBucketCount,
       base::HistogramBase::kUmaTargetedHistogramFlag);
   if (histogram) {
     histogram->AddTime(time);
   }
 
+  std::string suffix = GetUmaSuffixForStore(file_path);
   base::HistogramBase* histogram_suffix = base::Histogram::FactoryTimeGet(
       metric + suffix, base::TimeDelta::FromMilliseconds(1),
       base::TimeDelta::FromMinutes(1), kBucketCount,
       base::HistogramBase::kUmaTargetedHistogramFlag);
   if (histogram_suffix) {
     histogram_suffix->AddTime(time);
   }
 }
 
 void RecordAddUnlumpedHashesTime(base::TimeDelta time) {
@@ skipping 107 matching lines @@
 V4Store::~V4Store() {}
 
 std::string V4Store::DebugString() const {
   std::string state_base64;
   base::Base64Encode(state_, &state_base64);
 
   return base::StringPrintf("path: %" PRIsFP "; state: %s",
                             store_path_.value().c_str(), state_base64.c_str());
 }
 
-bool V4Store::Reset() {
-  // TODO(vakh): Implement skeleton.
+void V4Store::Reset() {
+  expected_checksum_.clear();
+  hash_prefix_map_.clear();
   state_ = "";
-  return true;
 }
 
 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());
 
   TimeTicks before = TimeTicks::Now();
   ApplyUpdateResult result = ProcessUpdate(hash_prefix_map_old, response);
   if (result == APPLY_UPDATE_SUCCESS) {
     RecordProcessPartialUpdateTime(TimeTicks::Now() - before, store_path_);
     // TODO(vakh): Create a ListUpdateResponse containing RICE encoded
     // hash prefixes and response_type as FULL_UPDATE, and write that to disk.
   }
   return result;
 }
 
 ApplyUpdateResult V4Store::ProcessFullUpdateAndWriteToDisk(
     std::unique_ptr<ListUpdateResponse> response) {
   TimeTicks before = TimeTicks::Now();
   ApplyUpdateResult result = ProcessFullUpdate(response);
   if (result == APPLY_UPDATE_SUCCESS) {
+    RecordProcessFullUpdateTime(TimeTicks::Now() - before, store_path_);
     RecordStoreWriteResult(WriteToDisk(std::move(response)));
-    RecordProcessFullUpdateTime(TimeTicks::Now() - before, store_path_);
   }
   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
@@ skipping 87 matching lines @@
   } else {
     DVLOG(1) << "Failure: ApplyUpdate: 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,
     HashPrefixMap* additions_map) {
   for (const auto& addition : additions) {
     ApplyUpdateResult apply_update_result = APPLY_UPDATE_SUCCESS;
     const CompressionType compression_type = addition.compression_type();
     if (compression_type == RAW) {
       DCHECK(addition.has_raw_hashes());
       DCHECK(addition.raw_hashes().has_raw_hashes());
 
@@ skipping 117 matching lines @@
 
     (*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,
                                        const std::string& expected_checksum) {
+  DCHECK(task_runner_->RunsTasksOnCurrentThread());
   DCHECK(hash_prefix_map_.empty());
+
+  bool calculate_checksum = !expected_checksum.empty();
+  if (old_prefixes_map.empty()) {
+    // If the old map is empty, which it is at startup, then just copy over the
+    // additions map.
+    DCHECK(!raw_removals);
+    hash_prefix_map_ = additions_map;
+
+    if (calculate_checksum) {
+      // Calculate the checksum asynchronously later and if it doesn't match,
+      // reset the store.
+      expected_checksum_ = expected_checksum;
+    }
+
+    return APPLY_UPDATE_SUCCESS;
+  }
+
   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) {
     // 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.
+      // prefix of size |next_smallest_prefix_size| 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(base::string_as_array(&next_smallest_prefix_old),
                                next_smallest_prefix_size);
@@ skipping 38 matching lines @@
   }
 
   if (calculate_checksum) {
     std::string checksum(crypto::kSHA256Length, 0);
     checksum_ctx->Finish(base::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) << "Failure: Checksum mismatch: calculated: " << checksum_base64
-               << " expected: " << expected_checksum_base64;
+               << "; expected: " << expected_checksum_base64
+               << "; store: " << *this;
+      ;
       return CHECKSUM_MISMATCH_FAILURE;
     }
   }
 
   return APPLY_UPDATE_SUCCESS;
 }
 
 StoreReadResult V4Store::ReadFromDisk() {
   DCHECK(task_runner_->RunsTasksOnCurrentThread());
 
@@ skipping 69 matching lines @@
   DCHECK_EQ(file_format_string.size(), written);
 
   if (!base::Move(new_filename, store_path_)) {
     return UNABLE_TO_RENAME_FAILURE;
   }
 
   return WRITE_SUCCESS;
 }
 
 HashPrefix V4Store::GetMatchingHashPrefix(const FullHash& full_hash) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
   // It should never be the case that more than one hash prefixes match a given
   // full hash. However, if that happens, this method returns any one of them.
   // It does not guarantee which one of those will be returned.
   DCHECK_EQ(32u, full_hash.size());
   for (const auto& pair : hash_prefix_map_) {
     const PrefixSize& prefix_size = pair.first;
     const HashPrefixes& hash_prefixes = pair.second;
     HashPrefix hash_prefix = full_hash.substr(0, prefix_size);
     if (HashPrefixMatches(hash_prefix, hash_prefixes.begin(),
                           hash_prefixes.end())) {
@@ skipping 19 matching lines @@
   int result = hash_prefix.compare(mid_prefix);
   if (result == 0) {
     return true;
   } else if (result < 0) {
     return HashPrefixMatches(hash_prefix, begin, mid);
   } else {
     return HashPrefixMatches(hash_prefix, mid + prefix_size, end);
   }
 }
 
+bool V4Store::VerifyChecksum() {
+  DCHECK(task_runner_->RunsTasksOnCurrentThread());
+
+  if (expected_checksum_.empty()) {
+    // If the |expected_checksum_| is empty, the file (or hash_prefix_map_)
+    // should also be empty.
+    return hash_prefix_map_.empty();
+  }
+
+  IteratorMap iterator_map;
+  HashPrefix next_smallest_prefix;
+  InitializeIteratorMap(hash_prefix_map_, &iterator_map);
+  bool has_unmerged = GetNextSmallestUnmergedPrefix(
+      hash_prefix_map_, iterator_map, &next_smallest_prefix);
+
+  std::unique_ptr<crypto::SecureHash> checksum_ctx(
+      crypto::SecureHash::Create(crypto::SecureHash::SHA256));
+  while (has_unmerged) {
+    PrefixSize next_smallest_prefix_size;
+    next_smallest_prefix_size = next_smallest_prefix.size();
+
+    // Update the iterator map, which means that we have read one hash
+    // prefix of size |next_smallest_prefix_size| from hash_prefix_map_.
+    iterator_map[next_smallest_prefix_size] += next_smallest_prefix_size;
+
+    checksum_ctx->Update(base::string_as_array(&next_smallest_prefix),
+                         next_smallest_prefix_size);
+
+    // Find the next smallest unmerged element in the map.
+    has_unmerged = GetNextSmallestUnmergedPrefix(hash_prefix_map_, iterator_map,
+                                                 &next_smallest_prefix);
+  }
+
+  std::string checksum(crypto::kSHA256Length, 0);
+  checksum_ctx->Finish(base::string_as_array(&checksum), checksum.size());
+  if (checksum == expected_checksum_) {
+    expected_checksum_.clear();
+    return true;
+  }
+
+  std::string checksum_base64, expected_checksum_base64;
+  base::Base64Encode(checksum, &checksum_base64);
+  base::Base64Encode(expected_checksum_, &expected_checksum_base64);
+  DVLOG(1) << "Failure: Checksum mismatch: calculated: " << checksum_base64
+           << "; expected: " << expected_checksum_base64
+           << "; store: " << *this;
+  RecordApplyUpdateResultWhenReadingFromDisk(CHECKSUM_MISMATCH_FAILURE);
+
+  return false;
+}
+
 }  // namespace safe_browsing