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

chrome/browser/safe_browsing/database_manager.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/database_manager.h"
 
 #include "base/bind.h"
 #include "base/bind_helpers.h"
 #include "base/callback.h"
 #include "base/command_line.h"
@@ skipping 17 matching lines @@
 #include "chrome/common/chrome_paths.h"
 #include "chrome/common/chrome_switches.h"
 #include "chrome/common/url_constants.h"
 #include "content/public/browser/browser_thread.h"
 #include "content/public/browser/notification_service.h"
 
 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()) {
     result = SafeBrowsingProtocolManager::GET_HASH_FULL_HASH_EMPTY;
   } else if (hit) {
     result = SafeBrowsingProtocolManager::GET_HASH_FULL_HASH_HIT;
   } 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) {
 }
 
 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_hash.get()) {
-    OnCheckDownloadHashResult(
-        safe_browsing_util::SBFullHashToString(*check.full_hash),
-        check.threat_type);
+  } 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 {
     NOTREACHED();
   }
 }
 
 SafeBrowsingDatabaseManager::SafeBrowsingDatabaseManager(
     const scoped_refptr<SafeBrowsingService>& service)
     : sb_service_(service),
       database_(NULL),
       enabled_(false),
       enable_download_protection_(false),
       enable_csd_whitelist_(false),
       enable_download_whitelist_(false),
       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();
   enable_download_protection_ =
       !cmdline->HasSwitch(switches::kSbDisableDownloadProtection);
 
   // We only download the csd-whitelist if client-side phishing detection is
   // enabled.
   enable_csd_whitelist_ =
       !cmdline->HasSwitch(switches::kDisableClientSidePhishingDetection);
 
   // TODO(noelutz): remove this boolean variable since it should always be true
   // if SafeBrowsing is enabled.  Unfortunately, we have no test data for this
   // 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() {
   // We should have already been shut down. If we're still enabled, then the
   // database isn't going to be closed properly, which could lead to corruption.
   DCHECK(!enabled_);
 }
 
 bool SafeBrowsingDatabaseManager::CanCheckUrl(const GURL& url) const {
   return url.SchemeIs(chrome::kFtpScheme) ||
          url.SchemeIs(chrome::kHttpScheme) ||
          url.SchemeIs(chrome::kHttpsScheme);
 }
 
 bool SafeBrowsingDatabaseManager::CheckDownloadUrl(
     const std::vector<GURL>& url_chain,
     Client* client) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
   if (!enabled_ || !enable_download_protection_)
     return true;
 
   // 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;
 }
 
 bool SafeBrowsingDatabaseManager::CheckDownloadHash(
     const std::string& full_hash,
     Client* client) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
   DCHECK(!full_hash.empty());
   if (!enabled_ || !enable_download_protection_ || full_hash.empty())
     return true;
 
   // 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;
 }
 
 bool SafeBrowsingDatabaseManager::MatchCsdWhitelistUrl(const GURL& url) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
   if (!enabled_ || !enable_csd_whitelist_ || !MakeDatabaseAvailable()) {
     // There is something funky going on here -- for example, perhaps the user
     // has not restarted since enabling metrics reporting, so we haven't
     // enabled the csd whitelist yet.  Just to be safe we return true in this
     // case.
@@ skipping 24 matching lines @@
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
   if (!enabled_)
     return true;
 
   if (!CanCheckUrl(url))
     return true;
 
   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;
     queued_checks_.push_back(check);
     return false;
   }
 
   std::string list;
   std::vector<SBPrefix> prefix_hits;
   std::vector<SBFullHashResult> full_hits;
 
   bool prefix_match =
       database_->ContainsBrowseUrl(url, &list, &prefix_hits, &full_hits,
           sb_service_->protocol_manager()->last_update());
 
   UMA_HISTOGRAM_TIMES("SB2.FilterCheck", base::TimeTicks::Now() - start);
 
   if (!prefix_match)
     return true;  // URL is okay.
 
   // 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);
   checks_.insert(check);
 
   BrowserThread::PostTask(
       BrowserThread::IO, FROM_HERE,
       base::Bind(&SafeBrowsingDatabaseManager::OnCheckDone, this, check));
 
   return false;
@@ skipping 135 matching lines @@
   if (!enabled_)
     return;
 
   enabled_ = false;
 
   // Delete queued checks, calling back any clients with 'SB_THREAT_TYPE_SAFE'.
   // If we don't do this here we may fail to close the database below.
   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();
   }
 
   // Close the database.  We don't simply DeleteSoon() because if a close is
   // already pending, we'll double-free, and we don't set |database_| to NULL
   // because if there is still anything running on the db thread, it could
   // create a new database object (via GetDatabase()) that would then leak.
   CloseDatabase();
@@ skipping 11 matching lines @@
     safe_browsing_thread_.reset();
   }
 
   // Delete pending checks, calling back any clients with 'SB_THREAT_TYPE_SAFE'.
   // We have to do this after the db thread returns because methods on it can
   // have copies of these pointers, so deleting them might lead to accessing
   // garbage.
   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_);
 
   gethash_requests_.clear();
 }
 
 bool SafeBrowsingDatabaseManager::DatabaseAvailable() const {
   base::AutoLock lock(database_lock_);
   return !closing_database_ && (database_ != NULL);
 }
@@ skipping 99 matching lines @@
       GetHashRequestors requestors;
       requestors.push_back(check);
       gethash_requests_[prefix] = requestors;
     }
 
     // Reset the start time so that we can measure the network time without the
     // database time.
     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.
     HandleOneCheck(check, check->full_hits);
   }
 }
 
 void SafeBrowsingDatabaseManager::GetAllChunksFromDatabase(
     GetChunksCallback callback) {
   DCHECK_EQ(MessageLoop::current(), safe_browsing_thread_->message_loop());
@@ skipping 42 matching lines @@
   // below will add the check back to the queue, and we'll infinite-loop.
   DCHECK(DatabaseAvailable());
   while (!queued_checks_.empty()) {
     QueuedCheck check = queued_checks_.front();
     DCHECK(!check.start.is_null());
     HISTOGRAM_TIMES("SB.QueueDelay", base::TimeTicks::Now() - check.start);
     // If CheckUrl() determines the URL is safe immediately, it doesn't call the
     // 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();
   }
 }
 
 void SafeBrowsingDatabaseManager::HandleChunkForDatabase(
     const std::string& list_name, SBChunkList* chunks) {
   DCHECK_EQ(MessageLoop::current(), safe_browsing_thread_->message_loop());
   if (chunks) {
@@ skipping 82 matching lines @@
   const std::vector<SBPrefix>& prefixes,
   const std::vector<SBFullHashResult>& full_hashes) {
   DCHECK_EQ(MessageLoop::current(), safe_browsing_thread_->message_loop());
   GetDatabase()->CacheHashResults(prefixes, full_hashes);
 }
 
 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;
   }
 
   // Call back all interested parties, noting if any has a hit.
   GetHashRequestors& requestors = it->second;
   bool hit = false;
   for (GetHashRequestors::iterator r = requestors.begin();
        r != requestors.end(); ++r) {
     if (HandleOneCheck(*r, full_hashes))
       hit = true;
   }
-  RecordGetHashCheckStatus(hit, is_download, full_hashes);
+  RecordGetHashCheckStatus(hit, check_type, full_hashes);
 
   gethash_requests_.erase(it);
 }
 
 bool SafeBrowsingDatabaseManager::HandleOneCheck(
     SafeBrowsingCheck* check,
     const std::vector<SBFullHashResult>& full_hashes) {
   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(
     SafeBrowsingCheck* check) {
   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,
                    check));
     return;
   }
 
   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));
 }
 
 void SafeBrowsingDatabaseManager::CheckDownloadUrlOnSBThread(
     SafeBrowsingCheck* check) {
   DCHECK_EQ(MessageLoop::current(), safe_browsing_thread_->message_loop());
   DCHECK(enable_download_protection_);
 
   std::vector<SBPrefix> prefix_hits;
 
   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,
                    check));
     return;
   }
 
   check->need_get_hash = true;
   check->prefix_hits.clear();
   check->prefix_hits = prefix_hits;
   BrowserThread::PostTask(
       BrowserThread::IO, FROM_HERE,
       base::Bind(&SafeBrowsingDatabaseManager::OnCheckDone, this, check));
 }
 
 void SafeBrowsingDatabaseManager::TimeoutCallback(SafeBrowsingCheck* check) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
   DCHECK(check);
 
   if (!enabled_)
     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;
   }
 }
 
 void SafeBrowsingDatabaseManager::CheckDownloadUrlDone(
     SafeBrowsingCheck* check) {
   DCHECK(enable_download_protection_);
   SafeBrowsingCheckDone(check);
 }
 
 void SafeBrowsingDatabaseManager::CheckDownloadHashDone(
     SafeBrowsingCheck* check) {
   DCHECK(enable_download_protection_);
   SafeBrowsingCheckDone(check);
 }
 
 void SafeBrowsingDatabaseManager::SafeBrowsingCheckDone(
     SafeBrowsingCheck* check) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
   DCHECK(check);
 
   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);
   checks_.erase(check);
   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);
 
   safe_browsing_thread_->message_loop()->PostTask(FROM_HERE, task);
 
   MessageLoop::current()->PostDelayedTask(FROM_HERE,
       base::Bind(&SafeBrowsingDatabaseManager::TimeoutCallback,
                  check->timeout_factory_->GetWeakPtr(), check),
-      base::TimeDelta::FromMilliseconds(timeout_ms));
+      timeout);
 }