Refactor BinaryUploadService::Request to make getting data size async.

Bug: 999143, 999141
Change-Id: I395da7e76649e46bea9cc69bdb0009168833f90e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1793822Reviewed-by: Daniel Rubery <drubery@chromium.org>
Commit-Queue: Roger Tawa <rogerta@chromium.org>
Cr-Commit-Position: refs/heads/master@{#695020}

chrome/browser/safe_browsing/download_protection/binary_upload_service.cc

@@ -25,7 +25,6 @@
 namespace safe_browsing {
 namespace {
 
-const size_t kMaxUploadSizeBytes = 50 * 1024 * 1024;  // 50 MB
 const int kScanningTimeoutSeconds = 5 * 60;           // 5 minutes
 const char kSbBinaryUploadUrl[] =
     "https://safebrowsing.google.com/safebrowsing/uploads/webprotect";
@@ -55,12 +54,6 @@ void BinaryUploadService::UploadForDeepScanning(
   Request* raw_request = request.get();
   active_requests_[raw_request] = std::move(request);
 
-  if (raw_request->GetFileSize() > kMaxUploadSizeBytes) {
-    FinishRequest(raw_request, Result::FILE_TOO_LARGE,
-                  DeepScanningClientResponse());
-    return;
-  }
-
   if (!binary_fcm_service_) {
     FinishRequest(raw_request, Result::FAILED_TO_GET_TOKEN,
                   DeepScanningClientResponse());
@@ -96,16 +89,22 @@ void BinaryUploadService::OnGetInstanceID(Request* request,
   }
 
   request->set_fcm_token(instance_id);
-  request->GetFileContents(
-      base::BindOnce(&BinaryUploadService::OnGetFileContents,
-                     weakptr_factory_.GetWeakPtr(), request));
+  request->GetRequestData(base::BindOnce(&BinaryUploadService::OnGetRequestData,
+                                         weakptr_factory_.GetWeakPtr(),
+                                         request));
 }
 
-void BinaryUploadService::OnGetFileContents(Request* request,
-                                            const std::string& file_contents) {
+void BinaryUploadService::OnGetRequestData(Request* request,
+                                           Result result,
+                                           const Request::Data& data) {
   if (!IsActive(request))
     return;
 
+  if (result != Result::SUCCESS) {
+    FinishRequest(request, result, DeepScanningClientResponse());
+    return;
+  }
+
   net::NetworkTrafficAnnotationTag traffic_annotation =
       net::DefineNetworkTrafficAnnotation("safe_browsing_binary_upload", R"(
         semantics {
@@ -149,7 +148,7 @@ void BinaryUploadService::OnGetFileContents(Request* request,
   base::Base64Encode(metadata, &metadata);
 
   auto upload_request = MultipartUploadRequest::Create(
-      url_loader_factory_, GURL(kSbBinaryUploadUrl), metadata, file_contents,
+      url_loader_factory_, GURL(kSbBinaryUploadUrl), metadata, data.contents,
       traffic_annotation,
       base::BindOnce(&BinaryUploadService::OnUploadComplete,
                      weakptr_factory_.GetWeakPtr(), request));
@@ -255,6 +254,8 @@ void BinaryUploadService::FinishRequest(Request* request,
   }
 }
 
+BinaryUploadService::Request::Data::Data() = default;
+
 BinaryUploadService::Request::Request(Callback callback)
     : callback_(std::move(callback)) {}
 

chrome/browser/safe_browsing/download_protection/binary_upload_service.h

@@ -23,6 +23,9 @@ namespace safe_browsing {
 // and asynchronously retrieving a verdict.
 class BinaryUploadService {
  public:
+  // The maximum size of data that can be uploaded via this service.
+  constexpr static size_t kMaxUploadSizeBytes = 50 * 1024 * 1024;  // 50 MB
+
   BinaryUploadService(
       scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
       Profile* profile);
@@ -72,16 +75,22 @@ class BinaryUploadService {
     Request(Request&&) = delete;
     Request& operator=(Request&&) = delete;
 
-    // Asynchronously returns the file contents to upload.
-    // TODO(drubery): This could allocate up to 50MB of memory for a large file
-    // upload. We should see how often that causes errors, and possibly
-    // implement some sort of streaming interface so we don't use so much
-    // memory.
-    virtual void GetFileContents(
-        base::OnceCallback<void(const std::string&)> callback) = 0;
+    // Structure of data returned in the callback to GetRequestData().
+    struct Data {
+      Data();
+      std::string contents;
+    };
 
-    // Returns the content size.
-    virtual size_t GetFileSize() = 0;
+    // Asynchronously returns the file contents to upload.
+    // TODO(drubery): This could allocate up to kMaxUploadSizeBytes of memory
+    // for a large file upload. We should see how often that causes errors,
+    // and possibly implement some sort of streaming interface so we don't use
+    // so much memory.
+    //
+    // |result| is set to SUCCESS if getting the request data succeeded or
+    // some value describing the error.
+    using DataCallback = base::OnceCallback<void(Result, const Data&)>;
+    virtual void GetRequestData(DataCallback callback) = 0;
 
     // Returns the metadata to upload, as a DeepScanningClientRequest.
     const DeepScanningClientRequest& deep_scanning_request() const {
@@ -119,7 +128,9 @@ class BinaryUploadService {
 
   void OnGetInstanceID(Request* request, const std::string& token);
 
-  void OnGetFileContents(Request* request, const std::string& file_contents);
+  void OnGetRequestData(Request* request,
+                        Result result,
+                        const Request::Data& data);
 
   void OnUploadComplete(Request* request,
                         bool success,

chrome/browser/safe_browsing/download_protection/binary_upload_service_unittest.cc

@@ -28,9 +28,7 @@ class MockRequest : public BinaryUploadService::Request {
  public:
   explicit MockRequest(BinaryUploadService::Callback callback)
       : BinaryUploadService::Request(std::move(callback)) {}
-  MOCK_METHOD1(GetFileContents,
-               void(base::OnceCallback<void(const std::string&)>));
-  MOCK_METHOD0(GetFileSize, size_t());
+  MOCK_METHOD1(GetRequestData, void(DataCallback));
 };
 
 class FakeMultipartUploadRequest : public MultipartUploadRequest {
@@ -143,11 +141,13 @@ class BinaryUploadServiceTest : public testing::Test {
           *target_response = response;
         },
         scanning_result, scanning_response));
-    ON_CALL(*request, GetFileSize()).WillByDefault(Return(strlen("contents")));
-    ON_CALL(*request, GetFileContents(_))
+    ON_CALL(*request, GetRequestData(_))
         .WillByDefault(
-            Invoke([](base::OnceCallback<void(const std::string&)> callback) {
-              std::move(callback).Run("contents");
+            Invoke([](BinaryUploadService::Request::DataCallback callback) {
+              BinaryUploadService::Request::Data data;
+              data.contents = "contents";
+              std::move(callback).Run(BinaryUploadService::Result::SUCCESS,
+                                      data);
             }));
     return request;
   }
@@ -163,9 +163,16 @@ TEST_F(BinaryUploadServiceTest, FailsForLargeFile) {
   BinaryUploadService::Result scanning_result;
   DeepScanningClientResponse scanning_response;
 
+  ExpectInstanceID("valid id");
   std::unique_ptr<MockRequest> request =
       MakeRequest(&scanning_result, &scanning_response);
-  ON_CALL(*request, GetFileSize()).WillByDefault(Return(100 * 1024 * 1024));
+  ON_CALL(*request, GetRequestData(_))
+      .WillByDefault(
+          Invoke([](BinaryUploadService::Request::DataCallback callback) {
+            BinaryUploadService::Request::Data data;
+            std::move(callback).Run(BinaryUploadService::Result::FILE_TOO_LARGE,
+                                    data);
+          }));
   service_->UploadForDeepScanning(std::move(request));
 
   content::RunAllTasksUntilIdle();

chrome/browser/safe_browsing/download_protection/download_item_request.cc

@@ -43,7 +43,6 @@ DownloadItemRequest::DownloadItemRequest(download::DownloadItem* item,
                                          BinaryUploadService::Callback callback)
     : Request(std::move(callback)),
       item_(item),
-      download_item_renamed_(false),
       weakptr_factory_(this) {
   item_->AddObserver(this);
 }
@@ -53,43 +52,39 @@ DownloadItemRequest::~DownloadItemRequest() {
     item_->RemoveObserver(this);
 }
 
-void DownloadItemRequest::GetFileContents(
-    base::OnceCallback<void(const std::string&)> callback) {
+void DownloadItemRequest::GetRequestData(DataCallback callback) {
   if (item_ == nullptr) {
-    std::move(callback).Run("");
+    std::move(callback).Run(BinaryUploadService::Result::UNKNOWN, Data());
     return;
   }
 
-  pending_callbacks_.push_back(std::move(callback));
+  if (static_cast<size_t>(item_->GetTotalBytes()) >
+      BinaryUploadService::kMaxUploadSizeBytes) {
+    std::move(callback).Run(BinaryUploadService::Result::FILE_TOO_LARGE,
+                            Data());
+    return;
+  }
+
+  if (is_data_valid_) {
+    std::move(callback).Run(BinaryUploadService::Result::SUCCESS, data_);
+    return;
+  }
 
-  if (download_item_renamed_)
-    RunPendingGetFileContentsCallbacks();
+  pending_callbacks_.push_back(std::move(callback));
 }
 
 void DownloadItemRequest::RunPendingGetFileContentsCallbacks() {
   for (auto it = pending_callbacks_.begin(); it != pending_callbacks_.end();
        it++) {
-    base::PostTaskAndReplyWithResult(
-        FROM_HERE,
-        {base::ThreadPool(), base::TaskPriority::USER_VISIBLE,
-         base::MayBlock()},
-        base::BindOnce(&GetFileContentsBlocking, item_->GetFullPath()),
-        base::BindOnce(&DownloadItemRequest::OnGotFileContents,
-                       weakptr_factory_.GetWeakPtr(), std::move(*it)));
+    std::move(*it).Run(BinaryUploadService::Result::SUCCESS, data_);
   }
 
   pending_callbacks_.clear();
 }
 
-size_t DownloadItemRequest::GetFileSize() {
-  return item_ == nullptr ? 0 : item_->GetTotalBytes();
-}
-
 void DownloadItemRequest::OnDownloadUpdated(download::DownloadItem* download) {
-  if (download == item_ && item_->GetFullPath() == item_->GetTargetFilePath()) {
-    download_item_renamed_ = true;
-    RunPendingGetFileContentsCallbacks();
-  }
+  if (download == item_ && item_->GetFullPath() == item_->GetTargetFilePath())
+    ReadFile();
 }
 
 void DownloadItemRequest::OnDownloadDestroyed(
@@ -98,10 +93,19 @@ void DownloadItemRequest::OnDownloadDestroyed(
     item_ = nullptr;
 }
 
-void DownloadItemRequest::OnGotFileContents(
-    base::OnceCallback<void(const std::string&)> callback,
-    const std::string& contents) {
-  std::move(callback).Run(contents);
+void DownloadItemRequest::ReadFile() {
+  base::PostTaskAndReplyWithResult(
+      FROM_HERE,
+      {base::ThreadPool(), base::TaskPriority::USER_VISIBLE, base::MayBlock()},
+      base::BindOnce(&GetFileContentsBlocking, item_->GetFullPath()),
+      base::BindOnce(&DownloadItemRequest::OnGotFileContents,
+                     weakptr_factory_.GetWeakPtr()));
+}
+
+void DownloadItemRequest::OnGotFileContents(std::string contents) {
+  data_.contents = std::move(contents);
+  is_data_valid_ = true;
+  RunPendingGetFileContentsCallbacks();
 }
 
 }  // namespace safe_browsing

chrome/browser/safe_browsing/download_protection/download_item_request.h

@@ -27,17 +27,16 @@ class DownloadItemRequest : public BinaryUploadService::Request,
   DownloadItemRequest& operator=(DownloadItemRequest&&) = delete;
 
   // BinaryUploadService::Request implementation.
-  void GetFileContents(
-      base::OnceCallback<void(const std::string&)> callback) override;
-  size_t GetFileSize() override;
+  void GetRequestData(DataCallback callback) override;
 
   // download::DownloadItem::Observer implementation.
   void OnDownloadDestroyed(download::DownloadItem* download) override;
   void OnDownloadUpdated(download::DownloadItem* download) override;
 
  private:
-  void OnGotFileContents(base::OnceCallback<void(const std::string&)> callback,
-                         const std::string& contents);
+  void ReadFile();
+
+  void OnGotFileContents(std::string contents);
 
   // Calls to GetFileContents can be deferred if the download item is not yet
   // renamed to its final location. When ready, this method runs those
@@ -48,11 +47,15 @@ class DownloadItemRequest : public BinaryUploadService::Request,
   // thread. Unowned.
   download::DownloadItem* item_;
 
-  // Whether the download item has been renamed to its final destination yet.
-  bool download_item_renamed_;
+  // The file's data.
+  Data data_;
+
+  // Is the |data_| member valid?  It becomes valid once the file has been
+  // read successfully.
+  bool is_data_valid_ = false;
 
   // All pending callbacks to GetFileContents before the download item is ready.
-  std::vector<base::OnceCallback<void(const std::string&)>> pending_callbacks_;
+  std::vector<DataCallback> pending_callbacks_;
 
   base::WeakPtrFactory<DownloadItemRequest> weakptr_factory_;
 };

chrome/browser/safe_browsing/download_protection/download_item_request_unittest.cc

@@ -51,18 +51,15 @@ class DownloadItemRequestTest : public ::testing::Test {
   std::string download_contents_;
 };
 
-TEST_F(DownloadItemRequestTest, GetsSize) {
-  EXPECT_EQ(request_.GetFileSize(), download_contents_.size());
-}
-
 TEST_F(DownloadItemRequestTest, GetsContentsWaitsUntilRename) {
   ON_CALL(item_, GetFullPath())
       .WillByDefault(ReturnRef(download_temporary_path_));
 
   std::string download_contents = "";
-  request_.GetFileContents(base::BindOnce(
-      [](std::string* target_contents, const std::string& contents) {
-        *target_contents = contents;
+  request_.GetRequestData(base::BindOnce(
+      [](std::string* target_contents, BinaryUploadService::Result result,
+         const BinaryUploadService::Request::Data& data) {
+        *target_contents = data.contents;
       },
       &download_contents));
   content::RunAllTasksUntilIdle();