De-dup PrefixSet code in SafeBrowsingDatabaseManager.

chrome/browser/safe_browsing/safe_browsing_database.cc

 // Copyright (c) 2012 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 "chrome/browser/safe_browsing/safe_browsing_database.h"
 
 #include <algorithm>
 #include <iterator>
 
 #include "base/bind.h"
@@ skipping 585 matching lines @@
   if (extension_blacklist_store_.get()) {
     extension_blacklist_store_->Init(
         ExtensionBlacklistDBFilename(filename_base_),
         base::Bind(&SafeBrowsingDatabaseNew::HandleCorruptDatabase,
                    base::Unretained(this)));
   }
 
   if (side_effect_free_whitelist_store_.get()) {
     const base::FilePath side_effect_free_whitelist_filename =
         SideEffectFreeWhitelistDBFilename(filename_base_);
-    const base::FilePath side_effect_free_whitelist_prefix_set_filename =
-        PrefixSetForFilename(side_effect_free_whitelist_filename);
     side_effect_free_whitelist_store_->Init(
         side_effect_free_whitelist_filename,
         base::Bind(&SafeBrowsingDatabaseNew::HandleCorruptDatabase,
                    base::Unretained(this)));
 
-    // Only use the prefix set if database is present and non-empty.
-    if (GetFileSizeOrZero(side_effect_free_whitelist_filename)) {
-      const base::TimeTicks before = base::TimeTicks::Now();
-      side_effect_free_whitelist_prefix_set_ =
-          safe_browsing::PrefixSet::LoadFile(
-              side_effect_free_whitelist_prefix_set_filename);
-      UMA_HISTOGRAM_TIMES("SB2.SideEffectFreeWhitelistPrefixSetLoad",
-                          base::TimeTicks::Now() - before);
-      if (!side_effect_free_whitelist_prefix_set_.get())
-        RecordFailure(FAILURE_SIDE_EFFECT_FREE_WHITELIST_PREFIX_SET_READ);
-    }
+    LoadPrefixSet(side_effect_free_whitelist_filename,
+                  &side_effect_free_whitelist_prefix_set_,
+                  FAILURE_SIDE_EFFECT_FREE_WHITELIST_PREFIX_SET_READ);
   } else {
     // Delete any files of the side-effect free sidelist that may be around
     // from when it was previously enabled.
     SafeBrowsingStoreFile::DeleteStore(
         SideEffectFreeWhitelistDBFilename(filename_base_));
     base::DeleteFile(
         PrefixSetForFilename(SideEffectFreeWhitelistDBFilename(filename_base_)),
         false);
   }
 
@@ skipping 571 matching lines @@
         download_store_.get(),
         FAILURE_DOWNLOAD_DATABASE_UPDATE_FINISH);
     UMA_HISTOGRAM_COUNTS("SB2.DownloadDatabaseKilobytes",
                          static_cast<int>(size_bytes / 1024));
   }
 
   UpdatePrefixSetUrlStore(BrowseDBFilename(filename_base_),
                           browse_store_.get(),
                           &browse_prefix_set_,
                           FAILURE_BROWSE_DATABASE_UPDATE_FINISH,
-                          FAILURE_BROWSE_PREFIX_SET_WRITE);
+                          FAILURE_BROWSE_PREFIX_SET_WRITE,
+                          true);
 
   UpdateWhitelistStore(CsdWhitelistDBFilename(filename_base_),
                        csd_whitelist_store_.get(),
                        &csd_whitelist_);
   UpdateWhitelistStore(DownloadWhitelistDBFilename(filename_base_),
                        download_whitelist_store_.get(),
                        &download_whitelist_);
 
   if (extension_blacklist_store_) {
     int64 size_bytes = UpdateHashPrefixStore(
         ExtensionBlacklistDBFilename(filename_base_),
         extension_blacklist_store_.get(),
         FAILURE_EXTENSION_BLACKLIST_UPDATE_FINISH);
     UMA_HISTOGRAM_COUNTS("SB2.ExtensionBlacklistKilobytes",
                          static_cast<int>(size_bytes / 1024));
   }
 
-  if (side_effect_free_whitelist_store_)
-    UpdateSideEffectFreeWhitelistStore();
+  if (side_effect_free_whitelist_store_) {
+    UpdatePrefixSetUrlStore(SideEffectFreeWhitelistDBFilename(filename_base_),
+                            side_effect_free_whitelist_store_.get(),
+                            &side_effect_free_whitelist_prefix_set_,
+                            FAILURE_SIDE_EFFECT_FREE_WHITELIST_UPDATE_FINISH,
+                            FAILURE_SIDE_EFFECT_FREE_WHITELIST_PREFIX_SET_WRITE,
+                            false);
+  }
 
   if (ip_blacklist_store_)
     UpdateIpBlacklistStore();
 
   if (unwanted_software_store_) {
     UpdatePrefixSetUrlStore(UnwantedSoftwareDBFilename(filename_base_),
                             unwanted_software_store_.get(),
                             &unwanted_software_prefix_set_,
                             FAILURE_UNWANTED_SOFTWARE_DATABASE_UPDATE_FINISH,
-                            FAILURE_UNWANTED_SOFTWARE_PREFIX_SET_WRITE);
+                            FAILURE_UNWANTED_SOFTWARE_PREFIX_SET_WRITE,
+                            true);
   }
 }
 
 void SafeBrowsingDatabaseNew::UpdateWhitelistStore(
     const base::FilePath& store_filename,
     SafeBrowsingStore* store,
     SBWhitelist* whitelist) {
   DCHECK(thread_checker_.CalledOnValidThread());
 
   if (!store)
@@ skipping 35 matching lines @@
 #endif
 
   return GetFileSizeOrZero(store_filename);
 }
 
 void SafeBrowsingDatabaseNew::UpdatePrefixSetUrlStore(
     const base::FilePath& db_filename,
     SafeBrowsingStore* url_store,
     scoped_ptr<const safe_browsing::PrefixSet>* prefix_set,
     FailureType finish_failure_type,
-    FailureType write_failure_type) {
+    FailureType write_failure_type,
+    bool store_full_hashes_in_prefix_set) {
   DCHECK(thread_checker_.CalledOnValidThread());
+  DCHECK(url_store);
+  DCHECK(prefix_set);
 
   // Measure the amount of IO during the filter build.
   base::IoCounters io_before, io_after;
   base::ProcessHandle handle = base::GetCurrentProcessHandle();
   scoped_ptr<base::ProcessMetrics> metric(
 #if !defined(OS_MACOSX)
       base::ProcessMetrics::CreateProcessMetrics(handle)
 #else
       // Getting stats only for the current process is enough, so NULL is fine.
       base::ProcessMetrics::CreateProcessMetrics(handle, NULL)
@@ skipping 10 matching lines @@
   // TODO(shess): Perhaps refactor to let builder accumulate full hashes on the
   // fly?  Other clients use the SBAddFullHash vector, but AFAICT they only use
   // the SBFullHash portion.  It would need an accessor on PrefixSet.
   safe_browsing::PrefixSetBuilder builder;
   std::vector<SBAddFullHash> add_full_hashes;
   if (!url_store->FinishUpdate(&builder, &add_full_hashes)) {
     RecordFailure(finish_failure_type);
     return;
   }
 
-  std::vector<SBFullHash> full_hash_results;
-  for (size_t i = 0; i < add_full_hashes.size(); ++i) {
-    full_hash_results.push_back(add_full_hashes[i].full_hash);
+  scoped_ptr<const safe_browsing::PrefixSet> new_prefix_set;
+  if (store_full_hashes_in_prefix_set) {
+    std::vector<SBFullHash> full_hash_results;
+    for (size_t i = 0; i < add_full_hashes.size(); ++i) {
+      full_hash_results.push_back(add_full_hashes[i].full_hash);
+    }
+
+    new_prefix_set = builder.GetPrefixSet(full_hash_results);
+  } else {
+    // TODO(gab): Ensure that stores which do not want full hashes just don't
+    // have full hashes in the first place and remove
+    // |store_full_hashes_in_prefix_set| and the code specialization incurred
+    // here.
+    new_prefix_set = builder.GetPrefixSetNoHashes();
   }
 
-  scoped_ptr<const safe_browsing::PrefixSet> new_prefix_set(
-      builder.GetPrefixSet(full_hash_results));
-
   // Swap in the newly built filter.
   {
     base::AutoLock locked(lookup_lock_);
     prefix_set->swap(new_prefix_set);
   }
 
   UMA_HISTOGRAM_LONG_TIMES("SB2.BuildFilter", base::TimeTicks::Now() - before);
 
   // Persist the prefix set to disk. Note: there is no need to lock since the
   // only write to |*prefix_set| is on this thread (in the swap() above).
   WritePrefixSet(db_filename, prefix_set->get(), write_failure_type);
 
   // Gather statistics.
   if (got_counters && metric->GetIOCounters(&io_after)) {
     UMA_HISTOGRAM_COUNTS("SB2.BuildReadKilobytes",
                          static_cast<int>(io_after.ReadTransferCount -
                                           io_before.ReadTransferCount) / 1024);
     UMA_HISTOGRAM_COUNTS("SB2.BuildWriteKilobytes",
                          static_cast<int>(io_after.WriteTransferCount -
                                           io_before.WriteTransferCount) / 1024);
     UMA_HISTOGRAM_COUNTS("SB2.BuildReadOperations",
                          static_cast<int>(io_after.ReadOperationCount -
                                           io_before.ReadOperationCount));
     UMA_HISTOGRAM_COUNTS("SB2.BuildWriteOperations",
                          static_cast<int>(io_after.WriteOperationCount -
                                           io_before.WriteOperationCount));
   }
 
   const int64 file_size = GetFileSizeOrZero(db_filename);
-  UMA_HISTOGRAM_COUNTS("SB2.BrowseDatabaseKilobytes",
+  UMA_HISTOGRAM_COUNTS("SB2.DatabaseKilobytes",
                        static_cast<int>(file_size / 1024));
 
 #if defined(OS_MACOSX)
   base::mac::SetFileBackupExclusion(db_filename);
 #endif
 }
 
-void SafeBrowsingDatabaseNew::UpdateSideEffectFreeWhitelistStore() {
-  DCHECK(thread_checker_.CalledOnValidThread());
-
-  safe_browsing::PrefixSetBuilder builder;
-  std::vector<SBAddFullHash> add_full_hashes_result;
-
-  if (!side_effect_free_whitelist_store_->FinishUpdate(
-          &builder, &add_full_hashes_result)) {
-    RecordFailure(FAILURE_SIDE_EFFECT_FREE_WHITELIST_UPDATE_FINISH);
-    return;
-  }
-  scoped_ptr<const safe_browsing::PrefixSet> new_prefix_set(
-      builder.GetPrefixSetNoHashes());
-
-  // Swap in the newly built prefix set.
-  {
-    base::AutoLock locked(lookup_lock_);
-    side_effect_free_whitelist_prefix_set_.swap(new_prefix_set);
-  }
-
-  const base::FilePath side_effect_free_whitelist_filename =
-      SideEffectFreeWhitelistDBFilename(filename_base_);
-  const base::FilePath side_effect_free_whitelist_prefix_set_filename =
-      PrefixSetForFilename(side_effect_free_whitelist_filename);
-  const base::TimeTicks before = base::TimeTicks::Now();
-  const bool write_ok = side_effect_free_whitelist_prefix_set_->WriteFile(
-      side_effect_free_whitelist_prefix_set_filename);
-  UMA_HISTOGRAM_TIMES("SB2.SideEffectFreePrefixSetWrite",
-                      base::TimeTicks::Now() - before);
-
-  if (!write_ok)
-    RecordFailure(FAILURE_SIDE_EFFECT_FREE_WHITELIST_PREFIX_SET_WRITE);
-
-  // Gather statistics.
-  int64 file_size = GetFileSizeOrZero(
-      side_effect_free_whitelist_prefix_set_filename);
-  UMA_HISTOGRAM_COUNTS("SB2.SideEffectFreeWhitelistPrefixSetKilobytes",
-                       static_cast<int>(file_size / 1024));
-  file_size = GetFileSizeOrZero(side_effect_free_whitelist_filename);
-  UMA_HISTOGRAM_COUNTS("SB2.SideEffectFreeWhitelistDatabaseKilobytes",
-                       static_cast<int>(file_size / 1024));
-
-#if defined(OS_MACOSX)
-  base::mac::SetFileBackupExclusion(side_effect_free_whitelist_filename);
-  base::mac::SetFileBackupExclusion(
-      side_effect_free_whitelist_prefix_set_filename);
-#endif
-}
-
 void SafeBrowsingDatabaseNew::UpdateIpBlacklistStore() {
   DCHECK(thread_checker_.CalledOnValidThread());
 
   // Note: prefixes will not be empty.  The current data store implementation
   // stores all full-length hashes as both full and prefix hashes.
   safe_browsing::PrefixSetBuilder builder;
   std::vector<SBAddFullHash> full_hashes;
   if (!ip_blacklist_store_->FinishUpdate(&builder, &full_hashes)) {
     RecordFailure(FAILURE_IP_BLACKLIST_UPDATE_FINISH);
     LoadIpBlacklist(std::vector<SBAddFullHash>());  // Clear the list.
@@ skipping 256 matching lines @@
   SBFullHash malware_kill_switch = SBFullHashForString(kMalwareIPKillSwitchUrl);
   std::vector<SBFullHash> full_hashes;
   full_hashes.push_back(malware_kill_switch);
   return ContainsWhitelistedHashes(csd_whitelist_, full_hashes);
 }
 
 bool SafeBrowsingDatabaseNew::IsCsdWhitelistKillSwitchOn() {
   base::AutoLock locked(lookup_lock_);
   return csd_whitelist_.second;
 }