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..228e0ea539f8494aabaf2c5c82b7c4ae61aab1c1 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,
@@ -163,23 +166,30 @@ enum SBStatsType {
 // 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()),
+      : 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);
+    if (item_)

[asanka, 2017/01/27 00:10:48]

|item_| had better be true here :-). For cases like this, it's best to just
DCHECK since there's no point constructing a DownloadSBClient if we don't have a
DownloadItem.

Might as well DCHECK service_ as well.

[Jialiu Lin, 2017/01/27 01:29:46]

You're right. DCHECK makes more sense.

+      item_->AddObserver(this);
+    Profile* profile = Profile::FromBrowserContext(item_->GetBrowserContext());
     extended_reporting_level_ =
         profile ? GetExtendedReportingLevel(*profile->GetPrefs())
                 : SBER_LEVEL_OFF;
@@ -188,9 +198,18 @@ class DownloadSBClient
   virtual void StartCheck() = 0;
   virtual bool IsDangerous(SBThreatType threat_type) const = 0;
 
+  // Implements DownloadItem::Observer.
+  void OnDownloadDestroyed(content::DownloadItem* download) override {
+    download->RemoveObserver(this);
+    item_ = nullptr;
+  }
+
  protected:
   friend class base::RefCountedThreadSafe<DownloadSBClient>;
-  ~DownloadSBClient() override {}
+  ~DownloadSBClient() override {
+    if (item_)

[asanka, 2017/01/27 00:10:48]

FYI: This is totally fine. But another way to do this is to use a
base::ScopedObserver<DownloadItem, DownloadItem::Observer>. This way we don't
need to worry about manually removing the observer at the destruction time.

[Jialiu Lin, 2017/01/27 01:29:46]

ScopedObserver is much cleaner. Thanks!

+      item_->RemoveObserver(this);
+  }
 
   void CheckDone(SBThreatType threat_type) {
     DownloadProtectionService::DownloadCheckResult result =
@@ -208,14 +227,25 @@ class DownloadSBClient
           FROM_HERE,
           base::Bind(&DownloadSBClient::ReportMalware,
                      this, threat_type));
+    } else if (service_ && service_->navigation_observer_manager() &&

[asanka, 2017/01/27 00:10:49]

There's a race here because CheckDone() could execute after the UI thread has
been shutdown and service_ has been freed already. In that case this will be a
UAF.

We only verified that service_ is safe to access on the UI thread :-). One way
to avoid this problem is to do check service_->navigation_observer_manager() in
the DownloadSBCLient constructor and store a bool indicating whether we need to
identify the referrer chain here. That way you don't have a racy dependency at
CheckDone().

[Jialiu Lin, 2017/01/27 01:29:45]

Moved navigation_observer_manager() and feature checking to DownloadSBClient
constructor.

+        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 +267,34 @@ class DownloadSBClient
     ui_manager_->MaybeReportSafeBrowsingHit(hit_report);
   }
 
+  void IdentifyReferrerChain() {
+    DCHECK_CURRENTLY_ON(BrowserThread::UI);
+    if (item_) {

[asanka, 2017/01/27 00:10:49]

Nit: return early if !item_. That way you'll have one fewer set of {}s to worry
about.

[Jialiu Lin, 2017/01/27 01:29:45]

Done.

+      std::unique_ptr<ReferrerChain> referrer_chain
+          = base::MakeUnique<ReferrerChain>();
+      service_->IdentifyReferrerChain(

[asanka, 2017/01/27 00:10:49]

Minor nit: Something nice to have would be for DPS::IdentifyReferrerChain to
just return a unique_ptr<ReferrerChain>.

[Jialiu Lin, 2017/01/27 01:29:46]

Done.

+         item_->GetURL(), item_->GetWebContents(),
+         referrer_chain.get());
+      if (!referrer_chain->empty()) {
+        item_->SetUserData(
+          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 +310,12 @@ class 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) { }
@@ -1006,10 +1056,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);
@@ -1613,11 +1667,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 +1875,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 +1889,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 +1917,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