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..3dac03c25e390e8e31ad082ea4fe09588bac67a9 100644
--- a/chrome/browser/safe_browsing/download_protection_service.cc
+++ b/chrome/browser/safe_browsing/download_protection_service.cc
@@ -108,6 +108,9 @@ const char DownloadProtectionService::kDownloadRequestUrl[] =
 const void* const DownloadProtectionService::kDownloadPingTokenKey
     = &kDownloadPingTokenKey;
 
+const void* const DownloadProtectionService::kDownloadReferrerChainDataKey
+    = &kDownloadReferrerChainDataKey;
+
 namespace {
 void RecordFileExtensionType(const std::string& metric_name,
                              const base::FilePath& file) {
@@ -166,20 +169,24 @@ class DownloadSBClient
       public base::RefCountedThreadSafe<DownloadSBClient> {
  public:
   DownloadSBClient(
-      const content::DownloadItem& item,
+      content::DownloadItem* item,
+      DownloadProtectionService* service,

[Nathan Parker, 2017/01/24 22:46:46]

Could service_ get destroyed before this DownloadSBClient is done?  If so, then
you'd need something like a weakptr on DownloadProtectionService.

Or, since SafeBrowsingService owns DownloadProtectionService, you could verify
that SBS destroys its SafeBrowsingService _before_ the DPS (thus closing out all
the Clients).  That might be fragile, unless we document it carefully.

[Jialiu Lin, 2017/01/25 00:11:19]

Done. Passed in a weakptr of DPS instead.

       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 +215,25 @@ class DownloadSBClient
           FROM_HERE,
           base::Bind(&DownloadSBClient::ReportMalware,
                      this, threat_type));
+    } else if (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,

[Nathan Parker, 2017/01/24 22:46:46]

Why is this triggered in CheckDone() and SAFE?  I thought the idea was to do it
at the beginning.

[Jialiu Lin, 2017/01/25 00:11:20]

This is not the final SAFE verdict. At this point, chrome just finished download
URL blacklist checking. If this download is already marked as dangerous
(!SB_THREAT_TYPE_SAFE) because of its url, warning will be shown WITHOUT
proceeding to content checking (thus, no ClientDownloadRequest will be sent
out).
So eventually, only these downloads pass blacklist checking and finishes actual
downloading will have a chance to send out its referrer chain info (as part of
ClientDownloadRequest). 

Potentially, we can include referrer info in the hit report too for these url
blacklisted downloads. But I need to talk to SB team first to make sure there is
an interested consumer of this data. 

+                       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 +255,32 @@ class DownloadSBClient
     ui_manager_->MaybeReportSafeBrowsingHit(hit_report);
   }
 
+  void IdentifyReferrerChain() {
+    DCHECK_CURRENTLY_ON(BrowserThread::UI);
+    std::unique_ptr<ReferrerChain> referrer_chain
+        = base::MakeUnique<ReferrerChain>();
+    service_->IdentifyReferrerChain(
+       item_->GetURL(), item_->GetWebContents(),
+       referrer_chain.get());
+    if (!referrer_chain->empty()) {
+      item_->SetUserData(
+        DownloadProtectionService::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_;
+  DownloadProtectionService* service_;
   DownloadProtectionService::CheckDownloadCallback callback_;
   scoped_refptr<SafeBrowsingUIManager> ui_manager_;
   base::TimeTicks start_time_;
@@ -261,11 +296,12 @@ class DownloadSBClient
 class DownloadUrlSBClient : public DownloadSBClient {
  public:
   DownloadUrlSBClient(
-      const content::DownloadItem& item,
+      content::DownloadItem* item,
+      DownloadProtectionService* service,

[Nathan Parker, 2017/01/24 22:46:46]

Same lifetime question here as above.

[Jialiu Lin, 2017/01/25 00:11:20]

Done.

       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 +309,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 +1042,15 @@ class DownloadProtectionService::CheckClientDownloadRequest
     }
     request.set_download_type(type_);
 
-    service_->AddReferrerChainToClientDownloadRequest(
-      item_->GetURL(),
-      item_->GetWebContents(),
-      &request);
+    ReferrerChainData* referrer_chain_data =
+      static_cast<ReferrerChainData*>(
+          item_->GetUserData(
+              DownloadProtectionService::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);
@@ -1613,11 +1654,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, this, callback, ui_manager_,
+                              database_manager_));
   // The client will release itself once it is done.
   BrowserThread::PostTask(
         BrowserThread::IO,
@@ -1820,10 +1862,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 +1876,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 +1904,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