Move download attribution right after CheckDownloadUrl

chrome/browser/safe_browsing/download_protection_service.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/download_protection_service.h"
 
 #include <stddef.h>
 
 #include <memory>
 
 #include "base/bind.h"
 #include "base/callback_helpers.h"
 #include "base/command_line.h"
 #include "base/compiler_specific.h"
 #include "base/format_macros.h"
 #include "base/macros.h"
 #include "base/memory/weak_ptr.h"
 #include "base/metrics/field_trial.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/metrics/sparse_histogram.h"
 #include "base/rand_util.h"
+#include "base/scoped_observer.h"
 #include "base/sequenced_task_runner_helpers.h"
 #include "base/sha1.h"
 #include "base/stl_util.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/string_split.h"
 #include "base/strings/string_util.h"
 #include "base/strings/stringprintf.h"
 #include "base/task/cancelable_task_tracker.h"
 #include "base/threading/sequenced_worker_pool.h"
 #include "base/time/time.h"
@@ skipping 49 matching lines @@
 const int64_t kDownloadRequestTimeoutMs = 7000;
 // We sample 1% of whitelisted downloads to still send out download pings.
 const double kWhitelistDownloadSampleRate = 0.01;
 
 // The number of user gestures we trace back for download attribution.
 const int kDownloadAttributionUserGestureLimit = 2;
 
 const char kDownloadExtensionUmaName[] = "SBClientDownload.DownloadExtensions";
 const char kUnsupportedSchemeUmaPrefix[] = "SBClientDownload.UnsupportedScheme";
 
+const void* const kDownloadReferrerChainDataKey =
+    &kDownloadReferrerChainDataKey;
+
 enum WhitelistType {
   NO_WHITELIST_MATCH,
   URL_WHITELIST,
   SIGNATURE_WHITELIST,
   WHITELIST_TYPE_MAX
 };
 
 void RecordCountOfWhitelistedDownload(WhitelistType type) {
   UMA_HISTOGRAM_ENUMERATION("SBClientDownload.CheckWhitelistResult", type,
                             WHITELIST_TYPE_MAX);
@@ skipping 55 matching lines @@
   // ALWAYS ADD NEW VALUES BEFORE THIS ONE.
   DOWNLOAD_CHECKS_MAX
 };
 
 }  // namespace
 
 // Parent SafeBrowsing::Client class used to lookup the bad binary
 // URL and digest list.  There are two sub-classes (one for each list).
 class DownloadSBClient
     : public SafeBrowsingDatabaseManager::Client,
+      public content::DownloadItem::Observer,
       public base::RefCountedThreadSafe<DownloadSBClient> {
  public:
   DownloadSBClient(
-      const content::DownloadItem& item,
+      content::DownloadItem* item,
+      DownloadProtectionService* service,
       const DownloadProtectionService::CheckDownloadCallback& callback,
       const scoped_refptr<SafeBrowsingUIManager>& ui_manager,
       SBStatsType total_type,
       SBStatsType dangerous_type)
-      : sha256_hash_(item.GetHash()),
-        url_chain_(item.GetUrlChain()),
-        referrer_url_(item.GetReferrerUrl()),
+      : observer_(this),
+        item_(item),
+        sha256_hash_(item->GetHash()),
+        url_chain_(item->GetUrlChain()),
+        referrer_url_(item->GetReferrerUrl()),
+        service_(service),
         callback_(callback),
         ui_manager_(ui_manager),
         start_time_(base::TimeTicks::Now()),
         total_type_(total_type),
         dangerous_type_(dangerous_type) {
-    Profile* profile = Profile::FromBrowserContext(item.GetBrowserContext());
+    DCHECK_CURRENTLY_ON(BrowserThread::UI);
+    DCHECK(item_);
+    DCHECK(service_);
+    observer_.Add(item_);
+    Profile* profile = Profile::FromBrowserContext(item_->GetBrowserContext());
     extended_reporting_level_ =
         profile ? GetExtendedReportingLevel(*profile->GetPrefs())
                 : SBER_LEVEL_OFF;
+    download_attribution_enabled_ = service_->navigation_observer_manager() &&
+        base::FeatureList::IsEnabled(
+            SafeBrowsingNavigationObserverManager::kDownloadAttribution);
   }
 
   virtual void StartCheck() = 0;
   virtual bool IsDangerous(SBThreatType threat_type) const = 0;
 
+  // Implements DownloadItem::Observer.
+  void OnDownloadDestroyed(content::DownloadItem* download) override {
+    observer_.Remove(download);
+    item_ = nullptr;
+  }
+
  protected:
   friend class base::RefCountedThreadSafe<DownloadSBClient>;
   ~DownloadSBClient() override {}
 
   void CheckDone(SBThreatType threat_type) {
     DownloadProtectionService::DownloadCheckResult result =
         IsDangerous(threat_type) ?
         DownloadProtectionService::DANGEROUS :
         DownloadProtectionService::SAFE;
     BrowserThread::PostTask(BrowserThread::UI,
                             FROM_HERE,
                             base::Bind(callback_, result));

[asanka, 2017/01/27 01:40:40]

Forgot to move this? :-P

[Jialiu Lin, 2017/01/27 01:49:51]

Oops! Thanks for catching this. 

     UpdateDownloadCheckStats(total_type_);
     if (threat_type != SB_THREAT_TYPE_SAFE) {
       UpdateDownloadCheckStats(dangerous_type_);
       BrowserThread::PostTask(
           BrowserThread::UI,
           FROM_HERE,
           base::Bind(&DownloadSBClient::ReportMalware,
                      this, threat_type));
+    } else if (download_attribution_enabled_) {
+        // Identify download referrer chain, which will be used in
+        // ClientDownloadRequest.
+        BrowserThread::PostTask(
+            BrowserThread::UI,
+            FROM_HERE,
+            base::Bind(&DownloadSBClient::IdentifyReferrerChain,
+                       this));
     }
   }
 
   void ReportMalware(SBThreatType threat_type) {
     std::string post_data;
-    if (!sha256_hash_.empty())
+    if (!sha256_hash_.empty()) {
       post_data += base::HexEncode(sha256_hash_.data(),
                                    sha256_hash_.size()) + "\n";
+    }
     for (size_t i = 0; i < url_chain_.size(); ++i) {
       post_data += url_chain_[i].spec() + "\n";
     }
 
     safe_browsing::HitReport hit_report;
     hit_report.malicious_url = url_chain_.back();
     hit_report.page_url = url_chain_.front();
     hit_report.referrer_url = referrer_url_;
     hit_report.is_subresource = true;
     hit_report.threat_type = threat_type;
     // TODO(nparker) Replace this with database_manager_->GetThreatSource();
     hit_report.threat_source = safe_browsing::ThreatSource::LOCAL_PVER3;
     // TODO(nparker) Populate hit_report.population_id once Pver4 is used here.
     hit_report.post_data = post_data;
     hit_report.extended_reporting_level = extended_reporting_level_;
     hit_report.is_metrics_reporting_active =
         ChromeMetricsServiceAccessor::IsMetricsAndCrashReportingEnabled();
 
     ui_manager_->MaybeReportSafeBrowsingHit(hit_report);
   }
 
+  void IdentifyReferrerChain() {
+    DCHECK_CURRENTLY_ON(BrowserThread::UI);
+    if (!item_)
+      return;
+
+    item_->SetUserData(kDownloadReferrerChainDataKey,
+                       new ReferrerChainData(
+                           service_->IdentifyReferrerChain(
+                               item_->GetURL(), item_->GetWebContents())));
+  }
+
   void UpdateDownloadCheckStats(SBStatsType stat_type) {
     UMA_HISTOGRAM_ENUMERATION("SB2.DownloadChecks",
                               stat_type,
                               DOWNLOAD_CHECKS_MAX);
   }
 
+  ScopedObserver<content::DownloadItem,
+                 content::DownloadItem::Observer> observer_;
+  // The DownloadItem we are checking. Must be accessed only on UI thread.
+  content::DownloadItem* item_;
+  // Copies of data from |item_| for access on other threads.
   std::string sha256_hash_;
   std::vector<GURL> url_chain_;
   GURL referrer_url_;
+  DownloadProtectionService* service_;
   DownloadProtectionService::CheckDownloadCallback callback_;
   scoped_refptr<SafeBrowsingUIManager> ui_manager_;
   base::TimeTicks start_time_;
+  bool download_attribution_enabled_;
 
  private:
   const SBStatsType total_type_;
   const SBStatsType dangerous_type_;
   ExtendedReportingLevel extended_reporting_level_;
 
   DISALLOW_COPY_AND_ASSIGN(DownloadSBClient);
 };
 
 class DownloadUrlSBClient : public DownloadSBClient {
  public:
   DownloadUrlSBClient(
-      const content::DownloadItem& item,
+      content::DownloadItem* item,
+      DownloadProtectionService* service,
       const DownloadProtectionService::CheckDownloadCallback& callback,
       const scoped_refptr<SafeBrowsingUIManager>& ui_manager,
       const scoped_refptr<SafeBrowsingDatabaseManager>& database_manager)
-      : DownloadSBClient(item, callback, ui_manager,
+      : DownloadSBClient(item, service, callback, ui_manager,
                          DOWNLOAD_URL_CHECKS_TOTAL,
                          DOWNLOAD_URL_CHECKS_MALWARE),
         database_manager_(database_manager) { }
 
   void StartCheck() override {
     DCHECK_CURRENTLY_ON(BrowserThread::IO);
     if (!database_manager_.get() ||
         database_manager_->CheckDownloadUrl(url_chain_, this)) {
       CheckDone(SB_THREAT_TYPE_SAFE);
     } else {
@@ skipping 720 matching lines @@
     if (type_ == ClientDownloadRequest::SAMPLED_UNSUPPORTED_FILE) {
       request.set_file_basename(
           base::FilePath(item_->GetTargetFilePath().Extension())
               .AsUTF8Unsafe());
     } else {
       request.set_file_basename(
         item_->GetTargetFilePath().BaseName().AsUTF8Unsafe());
     }
     request.set_download_type(type_);
 
-    service_->AddReferrerChainToClientDownloadRequest(
-      item_->GetURL(),
-      item_->GetWebContents(),
-      &request);
+    ReferrerChainData* referrer_chain_data =
+      static_cast<ReferrerChainData*>(
+          item_->GetUserData(kDownloadReferrerChainDataKey));
+    if (referrer_chain_data &&
+        !referrer_chain_data->GetReferrerChain()->empty()) {
+      request.mutable_referrer_chain()->Swap(
+          referrer_chain_data->GetReferrerChain());
+    }
 
     if (archive_is_valid_ != ArchiveValid::UNSET)
       request.set_archive_valid(archive_is_valid_ == ArchiveValid::VALID);
     request.mutable_signature()->CopyFrom(signature_info_);
     if (image_headers_)
       request.set_allocated_image_headers(image_headers_.release());
     if (archived_executable_)
       request.mutable_archived_binary()->Swap(&archived_binary_);
     if (!request.SerializeToString(&client_download_request_data_)) {
       FinishRequest(UNKNOWN, REASON_INVALID_REQUEST_PROTO);
@@ skipping 583 matching lines @@
     const CheckDownloadCallback& callback) {
   scoped_refptr<CheckClientDownloadRequest> request(
       new CheckClientDownloadRequest(item, callback, this,
                                      database_manager_,
                                      binary_feature_extractor_.get()));
   download_requests_.insert(request);
   request->Start();
 }
 
 void DownloadProtectionService::CheckDownloadUrl(
-    const content::DownloadItem& item,
+    content::DownloadItem* item,
     const CheckDownloadCallback& callback) {
-  DCHECK(!item.GetUrlChain().empty());
+  DCHECK(!item->GetUrlChain().empty());
   scoped_refptr<DownloadUrlSBClient> client(
-      new DownloadUrlSBClient(item, callback, ui_manager_, database_manager_));
+      new DownloadUrlSBClient(item, this, callback, ui_manager_,
+                              database_manager_));
   // The client will release itself once it is done.
   BrowserThread::PostTask(
         BrowserThread::IO,
         FROM_HERE,
         base::Bind(&DownloadUrlSBClient::StartCheck, client));
 }
 
 bool DownloadProtectionService::IsSupportedDownload(
     const content::DownloadItem& item,
     const base::FilePath& target_path) const {
@@ skipping 182 matching lines @@
 // static
 GURL DownloadProtectionService::GetDownloadRequestUrl() {
   GURL url(kDownloadRequestUrl);
   std::string api_key = google_apis::GetAPIKey();
   if (!api_key.empty())
     url = url.Resolve("?key=" + net::EscapeQueryParamValue(api_key, true));
 
   return url;
 }
 
-void DownloadProtectionService::AddReferrerChainToClientDownloadRequest(
+std::unique_ptr<ReferrerChain> DownloadProtectionService::IdentifyReferrerChain(
     const GURL& download_url,
-    content::WebContents* web_contents,
-    ClientDownloadRequest* out_request) {
-  if (!base::FeatureList::IsEnabled(
-      SafeBrowsingNavigationObserverManager::kDownloadAttribution) ||
-      !navigation_observer_manager_) {
-    return;
-  }
-
+    content::WebContents* web_contents) {
+  std::unique_ptr<ReferrerChain> referrer_chain =
+      base::MakeUnique<ReferrerChain>();
   int download_tab_id = SessionTabHelper::IdForTab(web_contents);
   UMA_HISTOGRAM_BOOLEAN(
       "SafeBrowsing.ReferrerHasInvalidTabID.DownloadAttribution",
       download_tab_id == -1);
-  SafeBrowsingNavigationObserverManager::ReferrerChain attribution_chain;
   SafeBrowsingNavigationObserverManager::AttributionResult result =
       navigation_observer_manager_->IdentifyReferrerChainForDownload(
           download_url,
           download_tab_id,
           kDownloadAttributionUserGestureLimit,
-          &attribution_chain);
+          referrer_chain.get());
   UMA_HISTOGRAM_COUNTS_100(
       "SafeBrowsing.ReferrerURLChainSize.DownloadAttribution",
-      attribution_chain.size());
+      referrer_chain->size());
   UMA_HISTOGRAM_ENUMERATION(
       "SafeBrowsing.ReferrerAttributionResult.DownloadAttribution", result,
       SafeBrowsingNavigationObserverManager::ATTRIBUTION_FAILURE_TYPE_MAX);
-  for (auto& entry : attribution_chain)
-    out_request->add_referrer_chain()->Swap(entry.get());
+  return referrer_chain;
 }
 
 void DownloadProtectionService::AddReferrerChainToPPAPIClientDownloadRequest(
   const GURL& initiating_frame_url,
   int tab_id,
   bool has_user_gesture,
   ClientDownloadRequest* out_request) {
   if (!base::FeatureList::IsEnabled(
       SafeBrowsingNavigationObserverManager::kDownloadAttribution) ||
       !navigation_observer_manager_) {
     return;
   }
 
   UMA_HISTOGRAM_BOOLEAN(
       "SafeBrowsing.ReferrerHasInvalidTabID.DownloadAttribution",
       tab_id == -1);
-  SafeBrowsingNavigationObserverManager::ReferrerChain attribution_chain;
   SafeBrowsingNavigationObserverManager::AttributionResult result =
       navigation_observer_manager_->IdentifyReferrerChainForPPAPIDownload(
           initiating_frame_url,
           tab_id,
           has_user_gesture,
           kDownloadAttributionUserGestureLimit,
-          &attribution_chain);
+          out_request->mutable_referrer_chain());
   UMA_HISTOGRAM_COUNTS_100(
       "SafeBrowsing.ReferrerURLChainSize.PPAPIDownloadAttribution",
-      attribution_chain.size());
+      out_request->referrer_chain_size());
   UMA_HISTOGRAM_ENUMERATION(
       "SafeBrowsing.ReferrerAttributionResult.PPAPIDownloadAttribution", result,
       SafeBrowsingNavigationObserverManager::ATTRIBUTION_FAILURE_TYPE_MAX);
-  for (auto& entry : attribution_chain)
-    out_request->add_referrer_chain()->Swap(entry.get());
 }
 
 }  // namespace safe_browsing