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>
 
@@ skipping 70 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);
 }
 
 }  // namespace
 
 const char DownloadProtectionService::kDownloadRequestUrl[] =
     "https://sb-ssl.google.com/safebrowsing/clientreport/download";
 
 const void* const DownloadProtectionService::kDownloadPingTokenKey
     = &kDownloadPingTokenKey;
 
+

[Nathan Parker, 2017/01/26 19:29:10]

nit: 3 blank lines

[Jialiu Lin, 2017/01/26 22:51:26]

Done.

+
 namespace {
 void RecordFileExtensionType(const std::string& metric_name,
                              const base::FilePath& file) {
   UMA_HISTOGRAM_SPARSE_SLOWLY(
       metric_name, FileTypePolicies::GetInstance()->UmaValueForFile(file));
 }
 
 void RecordArchivedArchiveFileExtensionType(const base::FilePath& file) {
   UMA_HISTOGRAM_SPARSE_SLOWLY(
       "SBClientDownload.ArchivedArchiveExtensions",
@@ skipping 38 matching lines @@
 
 }  // 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 base::RefCountedThreadSafe<DownloadSBClient> {
  public:
   DownloadSBClient(
-      const content::DownloadItem& item,
+      content::DownloadItem* item,
+      const base::WeakPtr<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()),
+      : 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);
+    Profile* profile = Profile::FromBrowserContext(item->GetBrowserContext());
     extended_reporting_level_ =
         profile ? GetExtendedReportingLevel(*profile->GetPrefs())
                 : SBER_LEVEL_OFF;
   }
 
   virtual void StartCheck() = 0;
   virtual bool IsDangerous(SBThreatType threat_type) const = 0;
 
  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));
     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 (service_ && service_->navigation_observer_manager() &&
+        base::FeatureList::IsEnabled(
+            SafeBrowsingNavigationObserverManager::kDownloadAttribution)) {
+        // 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 (service_) {

[asanka, 2017/01/26 18:54:47]

The DownloadProtectionService and the SafeBrowsingService outlives the main loop
(i.e. the UI thread).

[Jialiu Lin, 2017/01/26 22:51:26]

Thanks for confirming this. Removed checking for service_ and made it a regular
pointer (instead of WeakPtr).

+      std::unique_ptr<ReferrerChain> referrer_chain
+          = base::MakeUnique<ReferrerChain>();
+      service_->IdentifyReferrerChain(
+         item_->GetURL(), item_->GetWebContents(),
+         referrer_chain.get());
+      if (!referrer_chain->empty()) {
+        item_->SetUserData(

[asanka, 2017/01/26 18:54:47]

DownloadItem can go away abruptly for a number of reasons. I'd suggest:

  * Adding DownloadSBClient as a DownloadItem::Observer
  * Listening for DownloadItem::Observer::OnDownloadDestroyed(), and if seen,
setting item_ to nullptr
  * Condition all UI thread interactions on item_ not being nullptr.

You don't need to test for service_ nor retain a weak ref to it since it
outlives the UI thread if it was ever alive.

[Nathan Parker, 2017/01/26 19:29:10]

The CheckDone() method above is called on the IO thread.  I feel like we need to
ensure that DPS hasn't been destroyed already. I think it depends on the
destruction order within SafeBrowsingService. 

[looking..]

The DB_manager (in SBS) will be destroyed before ServicesDelegate (which owns
DPS), so then we're safe all around.

[Jialiu Lin, 2017/01/26 22:51:26]

Done.

[Jialiu Lin, 2017/01/26 22:51:26]

Thanks for checking this out!

+          kDownloadReferrerChainDataKey,
+          new ReferrerChainData(std::move(referrer_chain)));
+      }
+    }
+  }
+
   void UpdateDownloadCheckStats(SBStatsType stat_type) {
     UMA_HISTOGRAM_ENUMERATION("SB2.DownloadChecks",
                               stat_type,
                               DOWNLOAD_CHECKS_MAX);
   }
-
+  // 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_;
+  base::WeakPtr<DownloadProtectionService> service_;
   DownloadProtectionService::CheckDownloadCallback callback_;
   scoped_refptr<SafeBrowsingUIManager> ui_manager_;
   base::TimeTicks start_time_;
 
  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,
+      const base::WeakPtr<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)) {
+        database_manager_->CheckDownloadUrl(item_->GetUrlChain(), this)) {
       CheckDone(SB_THREAT_TYPE_SAFE);
     } else {
       AddRef();  // SafeBrowsingService takes a pointer not a scoped_refptr.
     }
   }
 
   bool IsDangerous(SBThreatType threat_type) const override {
     return threat_type == SB_THREAT_TYPE_BINARY_MALWARE_URL;
   }
 
@@ skipping 712 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 520 matching lines @@
 DownloadProtectionService::DownloadProtectionService(
     SafeBrowsingService* sb_service)
     : request_context_getter_(sb_service ? sb_service->url_request_context()
                                          : nullptr),
       enabled_(false),
       binary_feature_extractor_(new BinaryFeatureExtractor()),
       download_request_timeout_ms_(kDownloadRequestTimeoutMs),
       feedback_service_(
           new DownloadFeedbackService(request_context_getter_.get(),
                                       BrowserThread::GetBlockingPool())),
-      whitelist_sample_rate_(kWhitelistDownloadSampleRate) {
+      whitelist_sample_rate_(kWhitelistDownloadSampleRate),
+      weak_ptr_factory_(this) {
   if (sb_service) {
     ui_manager_ = sb_service->ui_manager();
     database_manager_ = sb_service->database_manager();
     navigation_observer_manager_ = sb_service->navigation_observer_manager();
     ParseManualBlacklistFlag();
   }
 }
 
 DownloadProtectionService::~DownloadProtectionService() {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
@@ skipping 42 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, weak_ptr_factory_.GetWeakPtr(), 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(
+void DownloadProtectionService::IdentifyReferrerChain(
     const GURL& download_url,
     content::WebContents* web_contents,
-    ClientDownloadRequest* out_request) {
+    ReferrerChain* out_referrer_chain) {
   if (!base::FeatureList::IsEnabled(
       SafeBrowsingNavigationObserverManager::kDownloadAttribution) ||
       !navigation_observer_manager_) {
     return;
   }
 
   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);
+          out_referrer_chain);
   UMA_HISTOGRAM_COUNTS_100(
       "SafeBrowsing.ReferrerURLChainSize.DownloadAttribution",
-      attribution_chain.size());
+      out_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());
 }
 
 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