Move URL whitelist checking b4 content analysis

Refactor checking of URL whitelist to happen before analyzing the
downloaded file, in order to improve SB check latency for whitelisted
URLs. Certificate whitelist checks still happen towards the end of
processing since they rely on file analysis.

Bug: 526841
Change-Id: I59e228dea251ff1a03e319efccf1c2a8aec89bc3
Reviewed-on: https://chromium-review.googlesource.com/587329
Commit-Queue: Varun Khaneja <vakh@chromium.org>
Reviewed-by: Varun Khaneja <vakh@chromium.org>
Reviewed-by: Jialiu Lin <jialiul@chromium.org>
Cr-Commit-Position: refs/heads/master@{#491890}

chrome/browser/safe_browsing/download_protection_service.cc

@@ -406,6 +406,24 @@ class DownloadProtectionService::CheckClientDownloadRequest
       is_incognito_ = item_->GetBrowserContext()->IsOffTheRecord();
     }
 
+    // If whitelist check passes, PostFinishTask() will be called to avoid
+    // analyzing file. Otherwise, AnalyzeFile() will be called to continue with
+    // analysis.
+    auto io_task_runner =
+        BrowserThread::GetTaskRunnerForThread(BrowserThread::IO);
+    cancelable_task_tracker_.PostTask(
+        io_task_runner.get(), FROM_HERE,
+        base::BindOnce(&CheckClientDownloadRequest::CheckUrlAgainstWhitelist,
+                       this));
+  }
+
+  void AnalyzeFile() {
+    // Returns if DownloadItem is destroyed during whitelist check.
+    if (item_ == nullptr) {
+      PostFinishTask(UNKNOWN, REASON_REQUEST_CANCELED);
+      return;
+    }
+
     DownloadCheckResultReason reason = REASON_MAX;
     if (!IsSupportedDownload(
         *item_, item_->GetTargetFilePath(), &reason, &type_)) {
@@ -523,6 +541,7 @@ class DownloadProtectionService::CheckClientDownloadRequest
   // unless a pending request is about to call FinishRequest.
   void Cancel() {
     DCHECK_CURRENTLY_ON(BrowserThread::UI);
+    cancelable_task_tracker_.TryCancelAll();
     if (fetcher_.get()) {
       // The DownloadProtectionService is going to release its reference, so we
       // might be destroyed before the URLFetcher completes.  Cancel the
@@ -676,13 +695,19 @@ class DownloadProtectionService::CheckClientDownloadRequest
 
   void OnFileFeatureExtractionDone() {
     // This can run in any thread, since it just posts more messages.
+    if (item_ == nullptr) {
+      PostFinishTask(UNKNOWN, REASON_REQUEST_CANCELED);
+      return;
+    }
 
     // TODO(noelutz): DownloadInfo should also contain the IP address of
     // every URL in the redirect chain.  We also should check whether the
     // download URL is hosted on the internal network.
     BrowserThread::PostTask(
         BrowserThread::IO, FROM_HERE,
-        base::BindOnce(&CheckClientDownloadRequest::CheckWhitelists, this));
+        base::BindOnce(
+            &CheckClientDownloadRequest::CheckCertificateChainAgainstWhitelist,
+            this));
 
     // We wait until after the file checks finish to start the timeout, as
     // windows can cause permissions errors if the timeout fired while we were
@@ -750,6 +775,10 @@ class DownloadProtectionService::CheckClientDownloadRequest
   void OnZipAnalysisFinished(const ArchiveAnalyzerResults& results) {
     DCHECK_CURRENTLY_ON(BrowserThread::UI);
     DCHECK_EQ(ClientDownloadRequest::ZIPPED_EXECUTABLE, type_);
+    if (item_ == nullptr) {
+      PostFinishTask(UNKNOWN, REASON_REQUEST_CANCELED);
+      return;
+    }
     if (!service_)
       return;
 
@@ -838,6 +867,10 @@ class DownloadProtectionService::CheckClientDownloadRequest
   void OnDmgAnalysisFinished(const ArchiveAnalyzerResults& results) {
     DCHECK_CURRENTLY_ON(BrowserThread::UI);
     DCHECK_EQ(ClientDownloadRequest::MAC_EXECUTABLE, type_);
+    if (item_ == nullptr) {
+      PostFinishTask(UNKNOWN, REASON_REQUEST_CANCELED);
+      return;
+    }
     if (!service_)
       return;
 
@@ -897,7 +930,7 @@ class DownloadProtectionService::CheckClientDownloadRequest
            base::RandDouble() < service_->whitelist_sample_rate();
   }
 
-  void CheckWhitelists() {
+  void CheckCertificateChainAgainstWhitelist() {
     DCHECK_CURRENTLY_ON(BrowserThread::IO);
 
     if (!database_manager_.get()) {
@@ -905,22 +938,6 @@ class DownloadProtectionService::CheckClientDownloadRequest
       return;
     }
 
-    const GURL& url = url_chain_.back();
-    // TODO(asanka): This may acquire a lock on the SB DB on the IO thread.
-    if (url.is_valid() && database_manager_->MatchDownloadWhitelistUrl(url)) {
-      DVLOG(2) << url << " is on the download whitelist.";
-      RecordCountOfWhitelistedDownload(URL_WHITELIST);
-      if (ShouldSampleWhitelistedDownload()) {
-        skipped_url_whitelist_ = true;
-      } else {
-        // TODO(grt): Continue processing without uploading so that
-        // ClientDownloadRequest callbacks can be run even for this type of safe
-        // download.
-        PostFinishTask(SAFE, REASON_WHITELISTED_URL);
-        return;
-      }
-    }
-
     if (!skipped_url_whitelist_ && signature_info_.trusted()) {
       for (int i = 0; i < signature_info_.certificate_chain_size(); ++i) {
         if (CertificateChainIsWhitelisted(
@@ -954,6 +971,40 @@ class DownloadProtectionService::CheckClientDownloadRequest
         base::BindOnce(&CheckClientDownloadRequest::GetTabRedirects, this));
   }
 
+  // TODO(jialiul): Make this function static and pass it the necessary objects
+  // instead of accessing private members of CheckClientDownloadRequest class.
+  // This way, if CheckClientDownloadRequest object is deleted in middle of
+  // whitelist check, browser will not crash.
+  void CheckUrlAgainstWhitelist() {
+    DCHECK_CURRENTLY_ON(BrowserThread::IO);
+
+    if (!database_manager_.get()) {
+      PostFinishTask(UNKNOWN, REASON_SB_DISABLED);
+      return;
+    }
+
+    const GURL& url = url_chain_.back();
+    // TODO(vakh): This may acquire a lock on the SB DB on the IO thread.
+    if (url.is_valid() && database_manager_->MatchDownloadWhitelistUrl(url)) {
+      DVLOG(2) << url << " is on the download whitelist.";
+      RecordCountOfWhitelistedDownload(URL_WHITELIST);
+      if (ShouldSampleWhitelistedDownload()) {
+        skipped_url_whitelist_ = true;
+      } else {
+        // TODO(grt): Continue processing without uploading so that
+        // ClientDownloadRequest callbacks can be run even for this type of safe
+        // download.
+        PostFinishTask(SAFE, REASON_WHITELISTED_URL);
+        return;
+      }
+    }
+
+    // Posts task to continue with analysis.
+    BrowserThread::PostTask(
+        BrowserThread::UI, FROM_HERE,
+        base::BindOnce(&CheckClientDownloadRequest::AnalyzeFile, this));
+  }
+
   void GetTabRedirects() {
     DCHECK_CURRENTLY_ON(BrowserThread::UI);
     if (!service_)
@@ -1330,6 +1381,7 @@ class DownloadProtectionService::CheckClientDownloadRequest
   bool skipped_certificate_whitelist_;
   bool is_extended_reporting_;
   bool is_incognito_;
+  base::CancelableTaskTracker cancelable_task_tracker_;
   base::WeakPtrFactory<CheckClientDownloadRequest> weakptr_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(CheckClientDownloadRequest);

chrome/browser/safe_browsing/download_protection_service_unittest.cc

@@ -676,11 +676,11 @@ TEST_F(DownloadProtectionServiceTest,
     FILE_PATH_LITERAL("a.exe"));  // final_path
 
   EXPECT_CALL(*binary_feature_extractor_.get(), CheckSignature(tmp_path_, _))
-      .Times(4);
+      .Times(3);
   EXPECT_CALL(*binary_feature_extractor_.get(),
               ExtractImageFeatures(
                   tmp_path_, BinaryFeatureExtractor::kDefaultOptions, _, _))
-      .Times(4);
+      .Times(3);
 
   // We should not get whilelist checks for other URLs than specified below.
   EXPECT_CALL(*sb_service_->mock_database_manager(),
@@ -771,11 +771,11 @@ TEST_F(DownloadProtectionServiceTest,
     FILE_PATH_LITERAL("a.tmp"),   // tmp_path
     FILE_PATH_LITERAL("a.exe"));  // final_path
   EXPECT_CALL(*binary_feature_extractor_.get(), CheckSignature(tmp_path_, _))
-      .Times(4);
+      .Times(1);
   EXPECT_CALL(*binary_feature_extractor_.get(),
               ExtractImageFeatures(
                   tmp_path_, BinaryFeatureExtractor::kDefaultOptions, _, _))
-      .Times(6);
+      .Times(3);
   // Assume http://www.whitelist.com/a.exe is on the whitelist.
   EXPECT_CALL(*sb_service_->mock_database_manager(),
               MatchDownloadWhitelistUrl(_)).Times(0);
@@ -2082,24 +2082,16 @@ TEST_F(DownloadProtectionServiceTest, TestDownloadItemDestroyed) {
                 MatchDownloadWhitelistUrl(_))
         .WillRepeatedly(Return(false));
 
-#if defined(OS_MACOSX)
     // Expects that MockDownloadItem will go out of scope while asynchronous
-    // processing is parsing file metadata, and thus ExtractFileOrDmgFeatures()
-    // will return rather than continuing to process the download.
+    // processing is checking whitelist, and thus will return after whitelist
+    // check rather than continuing to process the download, since
+    // OnDownloadDestroyed will be called to terminate the processing.
     EXPECT_CALL(*binary_feature_extractor_.get(), CheckSignature(tmp_path_, _))
         .Times(0);
     EXPECT_CALL(*binary_feature_extractor_.get(),
                 ExtractImageFeatures(
                     tmp_path_, BinaryFeatureExtractor::kDefaultOptions, _, _))
         .Times(0);
-#else
-    // Expects synchronous processing that continues to extract features from
-    // download even after MockDownloadItem goes out of scope.
-    EXPECT_CALL(*binary_feature_extractor_.get(), CheckSignature(tmp_path_, _));
-    EXPECT_CALL(*binary_feature_extractor_.get(),
-                ExtractImageFeatures(
-                    tmp_path_, BinaryFeatureExtractor::kDefaultOptions, _, _));
-#endif
 
     download_service_->CheckClientDownload(
         &item,
@@ -2132,11 +2124,12 @@ TEST_F(DownloadProtectionServiceTest,
         item.reset();
         return false;
       }));
-  EXPECT_CALL(*binary_feature_extractor_.get(), CheckSignature(tmp_path_, _));
+  EXPECT_CALL(*binary_feature_extractor_.get(), CheckSignature(tmp_path_, _))
+      .Times(0);
   EXPECT_CALL(*binary_feature_extractor_.get(),
-              ExtractImageFeatures(tmp_path_,
-                                   BinaryFeatureExtractor::kDefaultOptions,
-                                   _, _));
+              ExtractImageFeatures(
+                  tmp_path_, BinaryFeatureExtractor::kDefaultOptions, _, _))
+      .Times(0);
 
   RunLoop run_loop;
   download_service_->CheckClientDownload(