Small modifications to safebrowsing code to make it simpler to add the extension

chrome/browser/safe_browsing/database_manager.cc

Index: chrome/browser/safe_browsing/database_manager.cc
diff --git a/chrome/browser/safe_browsing/database_manager.cc b/chrome/browser/safe_browsing/database_manager.cc
index 5c82eed240a38109d70a36411a4a7c63a110b822..553464dbe1e11cd496f095a5ca6e587817fe7014 100644
--- a/chrome/browser/safe_browsing/database_manager.cc
+++ b/chrome/browser/safe_browsing/database_manager.cc
@@ -35,18 +35,14 @@ using content::BrowserThread;
 
 namespace {
 
-// When download url check takes this long, client's callback will be called
-// without waiting for the result.
-const int64 kDownloadUrlCheckTimeoutMs = 10000;
-
-// Similar to kDownloadUrlCheckTimeoutMs, but for download hash checks.
-const int64 kDownloadHashCheckTimeoutMs = 10000;
+// Timeout for match checks, e.g. download URLs, hashes.
+const int kCheckTimeoutMs = 10000;
 
 // Records disposition information about the check.  |hit| should be
 // |true| if there were any prefix hits in |full_hashes|.
 void RecordGetHashCheckStatus(
     bool hit,
-    bool is_download,
+    safe_browsing_util::ListType check_type,
     const std::vector<SBFullHashResult>& full_hashes) {
   SafeBrowsingProtocolManager::ResultType result;
   if (full_hashes.empty()) {
@@ -56,37 +52,69 @@ void RecordGetHashCheckStatus(
   } else {
     result = SafeBrowsingProtocolManager::GET_HASH_FULL_HASH_MISS;
   }
+  bool is_download = check_type == safe_browsing_util::BINURL ||
+                     check_type == safe_browsing_util::BINHASH;
   SafeBrowsingProtocolManager::RecordGetHashResult(is_download, result);
 }
 
 }  // namespace
 
-SafeBrowsingDatabaseManager::SafeBrowsingCheck::SafeBrowsingCheck()
-    : full_hash(NULL),
-      client(NULL),
+SafeBrowsingDatabaseManager::SafeBrowsingCheck::SafeBrowsingCheck(
+    safe_browsing_util::ListType check_type)
+    : client(NULL),
       need_get_hash(false),
-      threat_type(SB_THREAT_TYPE_SAFE),
-      is_download(false),
+      check_type(check_type),
       timeout_factory_(NULL) {
 }
 
 SafeBrowsingDatabaseManager::SafeBrowsingCheck::~SafeBrowsingCheck() {}
 
+// Gets the max threat type from a map of entity -> threat type. If all
+// entities are safe then this will be SB_THREAT_TYPE_SAFE.
+//
+// This is really just a semi-arbitrary way to choose a single threat type from
+// many. Under normal circumstances there will only be one type, and SAFE
+// threats will be filtered out.
+template<typename T>
+SBThreatType GetMaxThreatType(const std::map<T, SBThreatType>& threats) {
+  SBThreatType max = SB_THREAT_TYPE_SAFE;
+  for (typename std::map<T, SBThreatType>::const_iterator it = threats.begin();
+       it != threats.end(); ++it) {
+    if (it->second > max)
+      max = it->second;
+  }
+  return max;
+}
+
 void SafeBrowsingDatabaseManager::Client::OnSafeBrowsingResult(
     const SafeBrowsingCheck& check) {
   if (!check.urls.empty()) {
-
-    DCHECK(!check.full_hash.get());
-    if (!check.is_download) {
-      DCHECK_EQ(1U, check.urls.size());
-      OnCheckBrowseUrlResult(check.urls[0], check.threat_type);
-    } else {
-      OnCheckDownloadUrlResult(check.urls, check.threat_type);
+    DCHECK(check.full_hashes.empty());
+    switch (check.check_type) {
+      case safe_browsing_util::MALWARE:
+      case safe_browsing_util::PHISH:
+        DCHECK_EQ(1U, check.urls.size());
+        OnCheckBrowseUrlResult(check.urls[0],
+                               GetMaxThreatType(check.url_threats));

[Scott Hess - ex-Googler, 2013/01/11 23:44:05]

If url_threats were a parallel vector pre-filled with SB_THREAT_TYPE_SAFE, then
this would just be url_threats[0].

[not at google - send to devlin, 2013/01/14 23:00:55]

Done.

+        break;
+      case safe_browsing_util::BINURL:
+        OnCheckDownloadUrlResult(check.urls,
+                                 GetMaxThreatType(check.url_threats));

[Scott Hess - ex-Googler, 2013/01/11 23:44:05]

And this would be *max_element(url_threats.begin(), url_threads.end()).

[not at google - send to devlin, 2013/01/14 23:00:55]

Done.

+        break;
+      default:
+        NOTREACHED();
+    }
+  } else if (!check.full_hashes.empty()) {
+    switch (check.check_type) {
+      case safe_browsing_util::BINHASH:
+        DCHECK_EQ(1u, check.full_hashes.size());
+        OnCheckDownloadHashResult(
+            safe_browsing_util::SBFullHashToString(check.full_hashes[0]),
+            GetMaxThreatType(check.full_hash_threats));

[Scott Hess - ex-Googler, 2013/01/11 23:44:05]

Likewise here, just a different vector.

[not at google - send to devlin, 2013/01/14 23:00:55]

Done.

+        break;
+      default:
+        NOTREACHED();
     }
-  } else if (check.full_hash.get()) {
-    OnCheckDownloadHashResult(
-        safe_browsing_util::SBFullHashToString(*check.full_hash),
-        check.threat_type);
   } else {
     NOTREACHED();
   }
@@ -103,8 +131,10 @@ SafeBrowsingDatabaseManager::SafeBrowsingDatabaseManager(
       update_in_progress_(false),
       database_update_in_progress_(false),
       closing_database_(false),
-      download_urlcheck_timeout_ms_(kDownloadUrlCheckTimeoutMs),
-      download_hashcheck_timeout_ms_(kDownloadHashCheckTimeoutMs) {
+      download_url_check_timeout_(
+          base::TimeDelta::FromMilliseconds(kCheckTimeoutMs)),
+      download_hash_check_timeout_(
+          base::TimeDelta::FromMilliseconds(kCheckTimeoutMs)) {

[Scott Hess - ex-Googler, 2013/01/11 23:44:05]

If they're going to be the same, do they need to be distinct settings?  AFAICT,
the only reason the instance variables are needed is so tests can reach in and
mess with them, but it looks like those tests won't care if these are a single
timeout or separate timeouts.

[not at google - send to devlin, 2013/01/14 23:00:55]

No idea, I assumed it was this way for some good reason but I'll change it.

   DCHECK(sb_service_ != NULL);
 
   CommandLine* cmdline = CommandLine::ForCurrentProcess();
@@ -121,7 +151,6 @@ SafeBrowsingDatabaseManager::SafeBrowsingDatabaseManager(
   // list right now.  This means that we need to be able to disable this list
   // for the SafeBrowsing test to pass.
   enable_download_whitelist_ = enable_csd_whitelist_;
-
 }
 
 SafeBrowsingDatabaseManager::~SafeBrowsingDatabaseManager() {
@@ -145,14 +174,14 @@ bool SafeBrowsingDatabaseManager::CheckDownloadUrl(
 
   // We need to check the database for url prefix, and later may fetch the url
   // from the safebrowsing backends. These need to be asynchronous.
-  SafeBrowsingCheck* check = new SafeBrowsingCheck();
+  SafeBrowsingCheck* check = new SafeBrowsingCheck(safe_browsing_util::BINURL);
   check->urls = url_chain;
-  StartDownloadCheck(
+  StartSafeBrowsingCheck(
       check,
       client,
       base::Bind(&SafeBrowsingDatabaseManager::CheckDownloadUrlOnSBThread, this,
                  check),
-      download_urlcheck_timeout_ms_);
+      download_url_check_timeout_);
   return false;
 }
 
@@ -166,15 +195,15 @@ bool SafeBrowsingDatabaseManager::CheckDownloadHash(
 
   // We need to check the database for url prefix, and later may fetch the url
   // from the safebrowsing backends. These need to be asynchronous.
-  SafeBrowsingCheck* check = new SafeBrowsingCheck();
-  check->full_hash.reset(new SBFullHash);
-  safe_browsing_util::StringToSBFullHash(full_hash, check->full_hash.get());
-  StartDownloadCheck(
+  SafeBrowsingCheck* check = new SafeBrowsingCheck(safe_browsing_util::BINHASH);
+  check->full_hashes.push_back(
+      safe_browsing_util::StringToSBFullHash(full_hash));
+  StartSafeBrowsingCheck(
       check,
       client,
       base::Bind(&SafeBrowsingDatabaseManager::CheckDownloadHashOnSBThread,this,
                  check),
-      download_hashcheck_timeout_ms_);
+      download_hash_check_timeout_);
   return false;
 }
 
@@ -219,6 +248,7 @@ bool SafeBrowsingDatabaseManager::CheckBrowseUrl(const GURL& url,
   const base::TimeTicks start = base::TimeTicks::Now();
   if (!MakeDatabaseAvailable()) {
     QueuedCheck check;
+    check.check_type = safe_browsing_util::MALWARE;  // or PHISH
     check.client = client;
     check.url = url;
     check.start = start;
@@ -241,11 +271,10 @@ bool SafeBrowsingDatabaseManager::CheckBrowseUrl(const GURL& url,
 
   // Needs to be asynchronous, since we could be in the constructor of a
   // ResourceDispatcherHost event handler which can't pause there.
-  SafeBrowsingCheck* check = new SafeBrowsingCheck();
+  SafeBrowsingCheck* check =
+      new SafeBrowsingCheck(safe_browsing_util::MALWARE);  // or PHISH
   check->urls.push_back(url);
   check->client = client;
-  check->threat_type = SB_THREAT_TYPE_SAFE;
-  check->is_download = false;
   check->need_get_hash = full_hits.empty();
   check->prefix_hits.swap(prefix_hits);
   check->full_hits.swap(full_hits);
@@ -403,10 +432,9 @@ void SafeBrowsingDatabaseManager::DoStopOnIOThread() {
   while (!queued_checks_.empty()) {
     QueuedCheck queued = queued_checks_.front();
     if (queued.client) {
-      SafeBrowsingCheck sb_check;
+      SafeBrowsingCheck sb_check(queued.check_type);
       sb_check.urls.push_back(queued.url);
       sb_check.client = queued.client;
-      sb_check.threat_type = SB_THREAT_TYPE_SAFE;
       queued.client->OnSafeBrowsingResult(sb_check);
     }
     queued_checks_.pop_front();
@@ -438,10 +466,8 @@ void SafeBrowsingDatabaseManager::DoStopOnIOThread() {
   for (CurrentChecks::iterator it = checks_.begin();
        it != checks_.end(); ++it) {
     SafeBrowsingCheck* check = *it;
-    if (check->client) {
-      check->threat_type = SB_THREAT_TYPE_SAFE;
+    if (check->client)
       check->client->OnSafeBrowsingResult(*check);
-    }
   }
   STLDeleteElements(&checks_);
 
@@ -561,12 +587,14 @@ void SafeBrowsingDatabaseManager::OnCheckDone(SafeBrowsingCheck* check) {
     check->start = base::TimeTicks::Now();
     // Note: If |this| is deleted or stopped, the protocol_manager will
     // be destroyed as well - hence it's OK to do unretained in this case.
+    bool is_download = check->check_type == safe_browsing_util::BINURL ||
+                       check->check_type == safe_browsing_util::BINHASH;
     sb_service_->protocol_manager()->GetFullHash(
         check->prefix_hits,
         base::Bind(&SafeBrowsingDatabaseManager::HandleGetHashResults,
                    base::Unretained(this),
                    check),
-        check->is_download);
+        is_download);
   } else {
     // We may have cached results for previous GetHash queries.  Since
     // this data comes from cache, don't histogram hits.
@@ -630,10 +658,9 @@ void SafeBrowsingDatabaseManager::DatabaseLoadComplete() {
     // client's handler function (because normally it's being directly called by
     // the client).  Since we're not the client, we have to convey this result.
     if (check.client && CheckBrowseUrl(check.url, check.client)) {
-      SafeBrowsingCheck sb_check;
+      SafeBrowsingCheck sb_check(check.check_type);
       sb_check.urls.push_back(check.url);
       sb_check.client = check.client;
-      sb_check.threat_type = SB_THREAT_TYPE_SAFE;
       check.client->OnSafeBrowsingResult(sb_check);
     }
     queued_checks_.pop_front();
@@ -738,12 +765,12 @@ void SafeBrowsingDatabaseManager::OnHandleGetHashResults(
     SafeBrowsingCheck* check,
     const std::vector<SBFullHashResult>& full_hashes) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
-  bool is_download = check->is_download;
+  safe_browsing_util::ListType check_type = check->check_type;
   SBPrefix prefix = check->prefix_hits[0];
   GetHashRequests::iterator it = gethash_requests_.find(prefix);
   if (check->prefix_hits.size() > 1 || it == gethash_requests_.end()) {
     const bool hit = HandleOneCheck(check, full_hashes);
-    RecordGetHashCheckStatus(hit, is_download, full_hashes);
+    RecordGetHashCheckStatus(hit, check_type, full_hashes);
     return;
   }
 
@@ -755,7 +782,7 @@ void SafeBrowsingDatabaseManager::OnHandleGetHashResults(
     if (HandleOneCheck(*r, full_hashes))
       hit = true;
   }
-  RecordGetHashCheckStatus(hit, is_download, full_hashes);
+  RecordGetHashCheckStatus(hit, check_type, full_hashes);
 
   gethash_requests_.erase(it);
 }
@@ -766,28 +793,34 @@ bool SafeBrowsingDatabaseManager::HandleOneCheck(
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
   DCHECK(check);
 
-  // Always calculate the index, for recording hits.
-  int index = -1;
-  if (!check->urls.empty()) {
-    for (size_t i = 0; i < check->urls.size(); ++i) {
-      index = safe_browsing_util::GetUrlHashIndex(check->urls[i], full_hashes);
-      if (index != -1)
-        break;
-    }
-  } else {
-    index = safe_browsing_util::GetHashIndex(*(check->full_hash), full_hashes);
+  for (std::vector<GURL>::iterator it = check->urls.begin();
+       it != check->urls.end(); ++it) {
+    int index = safe_browsing_util::GetUrlHashIndex(*it, full_hashes);
+    if (index == -1)
+      continue;
+    SBThreatType threat =
+        GetThreatTypeFromListname(full_hashes[index].list_name);
+    if (threat != SB_THREAT_TYPE_SAFE)
+      check->url_threats[*it] = threat;
+    else
+      NOTREACHED();
+  }
+
+  for (std::vector<SBFullHash>::iterator it = check->full_hashes.begin();
+       it != check->full_hashes.end(); ++it) {
+    int index = safe_browsing_util::GetHashIndex(*it, full_hashes);
+    if (index == -1)
+      continue;
+    SBThreatType threat =
+        GetThreatTypeFromListname(full_hashes[index].list_name);
+    if (threat != SB_THREAT_TYPE_SAFE)
+      check->full_hash_threats[*it] = threat;
+    else
+      NOTREACHED();
   }
 
-  // |client| is NULL if the request was cancelled.
-  if (check->client) {
-    check->threat_type = SB_THREAT_TYPE_SAFE;
-    if (index != -1) {
-      check->threat_type = GetThreatTypeFromListname(
-          full_hashes[index].list_name);
-    }
-  }
   SafeBrowsingCheckDone(check);
-  return (index != -1);
+  return check->url_threats.empty() && check->full_hash_threats.empty();
 }
 
 void SafeBrowsingDatabaseManager::CheckDownloadHashOnSBThread(
@@ -795,9 +828,11 @@ void SafeBrowsingDatabaseManager::CheckDownloadHashOnSBThread(
   DCHECK_EQ(MessageLoop::current(), safe_browsing_thread_->message_loop());
   DCHECK(enable_download_protection_);
 
-  if (!database_->ContainsDownloadHashPrefix(check->full_hash->prefix)) {
+  DCHECK_EQ(1u, check->full_hashes.size());
+  SBFullHash full_hash = check->full_hashes[0];
+
+  if (!database_->ContainsDownloadHashPrefix(full_hash.prefix)) {
     // Good, we don't have hash for this url prefix.
-    check->threat_type = SB_THREAT_TYPE_SAFE;
     BrowserThread::PostTask(
         BrowserThread::IO, FROM_HERE,
         base::Bind(&SafeBrowsingDatabaseManager::CheckDownloadHashDone, this,
@@ -806,7 +841,7 @@ void SafeBrowsingDatabaseManager::CheckDownloadHashOnSBThread(
   }
 
   check->need_get_hash = true;
-  check->prefix_hits.push_back(check->full_hash->prefix);
+  check->prefix_hits.push_back(full_hash.prefix);
   BrowserThread::PostTask(
       BrowserThread::IO, FROM_HERE,
       base::Bind(&SafeBrowsingDatabaseManager::OnCheckDone, this, check));
@@ -821,7 +856,6 @@ void SafeBrowsingDatabaseManager::CheckDownloadUrlOnSBThread(
 
   if (!database_->ContainsDownloadUrl(check->urls, &prefix_hits)) {
     // Good, we don't have hash for this url prefix.
-    check->threat_type = SB_THREAT_TYPE_SAFE;
     BrowserThread::PostTask(
         BrowserThread::IO, FROM_HERE,
         base::Bind(&SafeBrowsingDatabaseManager::CheckDownloadUrlDone, this,
@@ -845,7 +879,6 @@ void SafeBrowsingDatabaseManager::TimeoutCallback(SafeBrowsingCheck* check) {
     return;
 
   DCHECK(checks_.find(check) != checks_.end());
-  DCHECK_EQ(check->threat_type, SB_THREAT_TYPE_SAFE);
   if (check->client) {
     check->client->OnSafeBrowsingResult(*check);
     check->client = NULL;
@@ -872,7 +905,9 @@ void SafeBrowsingDatabaseManager::SafeBrowsingCheckDone(
   if (!enabled_)
     return;
 
-  VLOG(1) << "SafeBrowsingCheckDone: " << check->threat_type;
+  VLOG(1) << "SafeBrowsingCheckDone: " <<
+              check->url_threats.size() << " URL hits, " <<
+              check->full_hash_threats.size() << " hash hits";
   DCHECK(checks_.find(check) != checks_.end());
   if (check->client)
     check->client->OnSafeBrowsingResult(*check);
@@ -880,15 +915,13 @@ void SafeBrowsingDatabaseManager::SafeBrowsingCheckDone(
   delete check;
 }
 
-void SafeBrowsingDatabaseManager::StartDownloadCheck(SafeBrowsingCheck* check,
-                                                     Client* client,
-                                                     const base::Closure& task,
-                                                     int64 timeout_ms) {
-
+void SafeBrowsingDatabaseManager::StartSafeBrowsingCheck(
+    SafeBrowsingCheck* check,
+    Client* client,
+    const base::Closure& task,
+    const base::TimeDelta& timeout) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
   check->client = client;
-  check->threat_type = SB_THREAT_TYPE_SAFE;
-  check->is_download = true;
   check->timeout_factory_.reset(
       new base::WeakPtrFactory<SafeBrowsingDatabaseManager>(this));
   checks_.insert(check);
@@ -898,5 +931,5 @@ void SafeBrowsingDatabaseManager::StartDownloadCheck(SafeBrowsingCheck* check,
   MessageLoop::current()->PostDelayedTask(FROM_HERE,
       base::Bind(&SafeBrowsingDatabaseManager::TimeoutCallback,
                  check->timeout_factory_->GetWeakPtr(), check),
-      base::TimeDelta::FromMilliseconds(timeout_ms));
+      timeout);
 }