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);
@@ 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()),
+      : 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;
   }
 
   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 =
         IsDangerous(threat_type) ?
         DownloadProtectionService::DANGEROUS :
         DownloadProtectionService::SAFE;
     BrowserThread::PostTask(BrowserThread::UI,
                             FROM_HERE,
                             base::Bind(callback_, result));

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

This is racy. |callback_| will be scheduled on the UI thread prior to
IdentifyReferrerChain() below. Hence it's technically possible for the
DownloadItem to proceed to CheckClientDownloadRequest(). If it does, then the
referrer chain will not be present on the DownloadItem by then.

This won't happen in practice due to the history and FILE thread hops that
happen in-between, but those are implementation details and not things that are
guaranteed by the APIs nor are they enforced via tests.

It would be more robust to schedule |callback_| after IdentifyReferrerChain.

Note that the ReportMalware task is OK because it doesn't modify our
DownloadItem.

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

You're right. Moved callback to the end of this function.

     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() &&

[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";
     }
 
     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_) {

[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_;
 
  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(
+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