Add |SetImageDownloadLimit| to ImageFetcher to limit downloaded bytes

components/image_fetcher/core/image_data_fetcher.cc

 // Copyright 2016 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 "components/image_fetcher/core/image_data_fetcher.h"
 
+#include <utility>
+
 #include "net/base/load_flags.h"
 #include "net/http/http_response_headers.h"
 #include "net/http/http_status_code.h"
 #include "net/url_request/url_fetcher.h"
 #include "net/url_request/url_request_context_getter.h"
 #include "net/url_request/url_request_status.h"
 #include "url/gurl.h"
 
 using data_use_measurement::DataUseUserData;
 
@@ skipping 21 matching lines @@
       data_use_service_name_(DataUseUserData::IMAGE_FETCHER_UNTAGGED),
       next_url_fetcher_id_(0) {}
 
 ImageDataFetcher::~ImageDataFetcher() {}
 
 void ImageDataFetcher::SetDataUseServiceName(
     DataUseServiceName data_use_service_name) {
   data_use_service_name_ = data_use_service_name;
 }
 
+void ImageDataFetcher::SetImageDownloadLimit(
+    base::Optional<int64_t> max_download_bytes) {
+  max_download_bytes_ = max_download_bytes;
+}
+
 void ImageDataFetcher::FetchImageData(
     const GURL& image_url,
     const ImageDataFetcherCallback& callback) {
   FetchImageData(
       image_url, callback, /*referrer=*/std::string(),
       net::URLRequest::CLEAR_REFERRER_ON_TRANSITION_FROM_SECURE_TO_INSECURE);
 }
 
 void ImageDataFetcher::FetchImageData(
     const GURL& image_url,
@@ skipping 12 matching lines @@
   request->url_fetcher->SetReferrerPolicy(referrer_policy);
   request->url_fetcher->SetLoadFlags(net::LOAD_DO_NOT_SEND_COOKIES |
                                      net::LOAD_DO_NOT_SAVE_COOKIES |
                                      net::LOAD_DO_NOT_SEND_AUTH_DATA);
   request->url_fetcher->Start();
 
   pending_requests_[request->url_fetcher.get()] = std::move(request);
 }
 
 void ImageDataFetcher::OnURLFetchComplete(const net::URLFetcher* source) {
-  auto request_iter = pending_requests_.find(source);
-  DCHECK(request_iter != pending_requests_.end());
-
+  DCHECK(pending_requests_.find(source) != pending_requests_.end());
   bool success = source->GetStatus().status() == net::URLRequestStatus::SUCCESS;
 
   RequestMetadata metadata;
   if (success && source->GetResponseHeaders()) {
     source->GetResponseHeaders()->GetMimeType(&metadata.mime_type);
     metadata.http_response_code = source->GetResponseHeaders()->response_code();
     success &= (metadata.http_response_code == net::HTTP_OK);
   }
   metadata.from_http_cache = source->WasCached();
 
   std::string image_data;
   if (success) {
     source->GetResponseAsString(&image_data);
   }
+  FinishRequest(source, metadata, image_data);
+}
+
+void ImageDataFetcher::OnURLFetchDownloadProgress(
+    const net::URLFetcher* source,
+    int64_t current,
+    int64_t total,
+    int64_t current_network_bytes) {
+  if (!max_download_bytes_.has_value()) {
+    return;
+  }
+  if (total <= max_download_bytes_.value() &&
+      current <= max_download_bytes_.value()) {
+    return;
+  }
+  DCHECK(pending_requests_.find(source) != pending_requests_.end());
+  DLOG(WARNING) << "Image data exceeded download size limit.";
+  RequestMetadata metadata;
+  metadata.http_response_code = net::URLFetcher::RESPONSE_CODE_INVALID;
+
+  FinishRequest(source, metadata, /*image_data=*/std::string());
+}
+
+void ImageDataFetcher::FinishRequest(const net::URLFetcher* source,
+                                     const RequestMetadata& metadata,
+                                     const std::string& image_data) {
+  auto request_iter = pending_requests_.find(source);
+  DCHECK(request_iter != pending_requests_.end());
   request_iter->second->callback.Run(image_data, metadata);
-
-  // Remove the finished request.
   pending_requests_.erase(request_iter);
 }
 
 }  // namespace image_fetcher