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 f2da285818686c0ed101da419ded4d3093dc202d..246dca7afa5695322a6be57b8c6fa46e12c95d07 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,17 +52,16 @@ void RecordGetHashCheckStatus(
   } else {
     result = SafeBrowsingProtocolManager::GET_HASH_FULL_HASH_MISS;
   }
-  SafeBrowsingProtocolManager::RecordGetHashResult(is_download, result);
+  SafeBrowsingProtocolManager::RecordGetHashResult(check_type, 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) {
 }
 
@@ -75,18 +70,38 @@ SafeBrowsingDatabaseManager::SafeBrowsingCheck::~SafeBrowsingCheck() {}
 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],
+            check.url_threats.empty() ? SB_THREAT_TYPE_SAFE
+                                      : check.url_threats.begin()->second);
+        break;
+      case safe_browsing_util::BINURL:
+        OnCheckDownloadUrlResult(
+            check.urls,
+            check.url_threats.empty() ? SB_THREAT_TYPE_SAFE
+                                      : SB_THREAT_TYPE_BINARY_MALWARE_URL);
+        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]),
+            check.full_hash_threats.empty()
+                ? SB_THREAT_TYPE_SAFE
+                : SB_THREAT_TYPE_BINARY_MALWARE_HASH);
+        break;
+      default:
+        NOTREACHED();
     }
-  } else if (check.full_hash.get()) {
-    OnCheckDownloadHashResult(
-        safe_browsing_util::SBFullHashToString(*check.full_hash),
-        check.threat_type);
   } else {
     NOTREACHED();
   }
@@ -103,8 +118,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)) {
   DCHECK(sb_service_ != NULL);
 
   CommandLine* cmdline = CommandLine::ForCurrentProcess();
@@ -121,7 +138,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 +161,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 +182,16 @@ 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);
+  SBFullHash sb_hash;
+  safe_browsing_util::StringToSBFullHash(full_hash, &sb_hash);
+  check->full_hashes.push_back(sb_hash);
+  StartSafeBrowsingCheck(
       check,
       client,
       base::Bind(&SafeBrowsingDatabaseManager::CheckDownloadHashOnSBThread,this,
                  check),
-      download_hashcheck_timeout_ms_);
+      download_hash_check_timeout_);
   return false;
 }
 
@@ -219,6 +236,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 +259,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);
@@ -401,10 +418,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();
@@ -436,10 +452,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_);
 
@@ -559,12 +573,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.
@@ -627,10 +643,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();
@@ -733,12 +748,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;
   }
 
@@ -750,7 +765,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);
 }
@@ -761,28 +776,26 @@ 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;
+  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) {
+      check->url_threats[*it] =
+          GetThreatTypeFromListname(full_hashes[index].list_name);
     }
-  } else {
-    index = safe_browsing_util::GetHashIndex(*(check->full_hash), full_hashes);
   }
 
-  // |client| is NULL if the request was cancelled.
-  if (check->client) {
-    check->threat_type = SB_THREAT_TYPE_SAFE;
+  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) {
-      check->threat_type = GetThreatTypeFromListname(
-          full_hashes[index].list_name);
+      check->full_hash_threats[*it] =
+          GetThreatTypeFromListname(full_hashes[index].list_name);
     }
   }
+
   SafeBrowsingCheckDone(check);
-  return (index != -1);
+  return check->url_threats.empty() && check->full_hash_threats.empty();
 }
 
 void SafeBrowsingDatabaseManager::CheckDownloadHashOnSBThread(
@@ -790,9 +803,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,
@@ -801,7 +816,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));
@@ -816,7 +831,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,
@@ -840,7 +854,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;
@@ -867,7 +880,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);
@@ -875,15 +890,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);
@@ -893,5 +906,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);
 }