Move download attribution right after CheckDownloadUrl

chrome/browser/safe_browsing/download_protection_service.cc

Index: chrome/browser/safe_browsing/download_protection_service.cc
diff --git a/chrome/browser/safe_browsing/download_protection_service.cc b/chrome/browser/safe_browsing/download_protection_service.cc
index e500413c8a437638eed1907f77fde09f1f351d6f..c4422891712289522e4d51622946df28a1ed945c 100644
--- a/chrome/browser/safe_browsing/download_protection_service.cc
+++ b/chrome/browser/safe_browsing/download_protection_service.cc
@@ -88,6 +88,9 @@ 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,
@@ -108,6 +111,8 @@ const char DownloadProtectionService::kDownloadRequestUrl[] =
 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) {
@@ -166,20 +171,24 @@ class DownloadSBClient
       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;
@@ -208,14 +217,25 @@ class DownloadSBClient
           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";
     }
@@ -237,15 +257,34 @@ class DownloadSBClient
     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_;
@@ -261,11 +300,12 @@ class 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) { }
@@ -273,7 +313,7 @@ class DownloadUrlSBClient : public DownloadSBClient {
   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.
@@ -1006,10 +1046,14 @@ class DownloadProtectionService::CheckClientDownloadRequest
     }
     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);
@@ -1550,7 +1594,8 @@ DownloadProtectionService::DownloadProtectionService(
       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();
@@ -1613,11 +1658,12 @@ void DownloadProtectionService::CheckClientDownload(
 }
 
 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,
@@ -1820,10 +1866,10 @@ GURL DownloadProtectionService::GetDownloadRequestUrl() {
   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_) {
@@ -1834,21 +1880,18 @@ void DownloadProtectionService::AddReferrerChainToClientDownloadRequest(
   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(
@@ -1865,22 +1908,19 @@ void DownloadProtectionService::AddReferrerChainToPPAPIClientDownloadRequest(
   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