Abort response body copy for in-memory cache when response is too large

Instead of checking the size after copying the whole response, it would
be efficient to abort the copy at the time when we know the content size
exceeds the limit.

Bug: 1339708
Change-Id: Id8893c2df8aa874cdee102f517c2f877009a9876
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3758458
Reviewed-by: Yutaka Hirano <yhirano@chromium.org>
Commit-Queue: Kenichi Ishibashi <bashi@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1024579}

services/network/network_service_memory_cache.cc

@@ -5,6 +5,7 @@
 #include "services/network/network_service_memory_cache.h"
 
 #include <algorithm>
+#include <limits>
 
 #include "base/bit_cast.h"
 #include "base/memory/ref_counted_memory.h"
@@ -162,6 +163,9 @@ NetworkServiceMemoryCache::NetworkServiceMemoryCache(
       max_per_entry_bytes_(kNetworkServiceMemoryCacheMaxPerEntrySize.Get()) {
   DCHECK(network_context_);
   DCHECK_GE(max_total_bytes_, max_per_entry_bytes_);
+  DCHECK_GE(static_cast<size_t>(std::numeric_limits<int>::max()),
+            max_per_entry_bytes_);
+
   memory_pressure_listener_ = std::make_unique<base::MemoryPressureListener>(
       FROM_HERE,
       base::BindRepeating(&NetworkServiceMemoryCache::OnMemoryPressure,
@@ -222,6 +226,9 @@ NetworkServiceMemoryCache::MaybeCreateWriter(
     return nullptr;
   }
 
+  if (response->content_length > static_cast<int>(max_per_entry_bytes_))
+    return nullptr;
+
   net::ValidationType validation_type = response->headers->RequiresValidation(
       response->request_time, response->response_time, GetCurrentTime());
   if (validation_type != net::VALIDATION_NONE)
@@ -237,8 +244,8 @@ NetworkServiceMemoryCache::MaybeCreateWriter(
       /*single_key_checksum=*/"");
 
   return std::make_unique<NetworkServiceMemoryCacheWriter>(
-      weak_ptr_factory_.GetWeakPtr(), GetNextTraceId(), std::move(cache_key),
-      url_request, request_destination, response);
+      weak_ptr_factory_.GetWeakPtr(), GetNextTraceId(), max_per_entry_bytes_,
+      std::move(cache_key), url_request, request_destination, response);
 }
 
 void NetworkServiceMemoryCache::StoreResponse(
@@ -248,6 +255,8 @@ void NetworkServiceMemoryCache::StoreResponse(
     const net::HttpVaryData& vary_data,
     mojom::URLResponseHeadPtr response_head,
     std::vector<unsigned char> data) {
+  DCHECK_GE(max_per_entry_bytes_, data.size());
+
   // TODO(https://crbug.com/1339708): Consider caching a response that doesn't
   // have contents.
   if (status.error_code != net::OK || data.size() == 0)
@@ -261,9 +270,6 @@ void NetworkServiceMemoryCache::StoreResponse(
       /*min=*/1, /*exclusive_max=*/50000000,
       /*buckets=*/50);
 
-  if (max_per_entry_bytes_ < data.size())
-    return;
-
   // Record fresness of the response in seconds.
   net::HttpResponseHeaders::FreshnessLifetimes lifetimes =
       response_head->headers->GetFreshnessLifetimes(

services/network/network_service_memory_cache_unittest.cc

@@ -97,6 +97,28 @@ std::unique_ptr<net::test_server::HttpResponse> CacheableResponseHandler(
   return response;
 }
 
+// Similar to above, but doesn't send Content-Length header.
+std::unique_ptr<net::test_server::HttpResponse>
+CacheableWithoutContentLengthHandler(
+    const net::test_server::HttpRequest& request) {
+  if (request.GetURL().path_piece() != "/cacheable_without_content_length")
+    return nullptr;
+
+  uint64_t body_size = 64;
+  std::string query_body_size;
+  if (net::GetValueForKeyInQuery(request.GetURL(), "body-size",
+                                 &query_body_size)) {
+    EXPECT_TRUE(base::StringToUint64(query_body_size, &body_size));
+  }
+
+  constexpr const char kHeader[] =
+      "HTTP/1.1 200 OK\n"
+      "Content-Type: text/plain\n";
+  auto response = std::make_unique<net::test_server::RawHttpResponse>(
+      kHeader, std::string(body_size, 'a'));
+  return response;
+}
+
 // Used for cross origin read blocking check.
 std::unique_ptr<net::test_server::HttpResponse> CorbCheckHandler(
     const net::test_server::HttpRequest& request) {
@@ -143,6 +165,8 @@ class NetworkServiceMemoryCacheTest : public testing::Test {
     test_server_.AddDefaultHandlers();
     test_server_.RegisterRequestHandler(
         base::BindRepeating(&CacheableResponseHandler));
+    test_server_.RegisterRequestHandler(
+        base::BindRepeating(&CacheableWithoutContentLengthHandler));
     test_server_.RegisterRequestHandler(base::BindRepeating(&CorbCheckHandler));
     ASSERT_TRUE(test_server_.Start());
 
@@ -444,6 +468,15 @@ TEST_F(NetworkServiceMemoryCacheTest, CanServe_ResponseTooLarge) {
   ASSERT_FALSE(CanServeFromMemoryCache(request));
 }
 
+TEST_F(NetworkServiceMemoryCacheTest,
+       CanServe_ResponseTooLargeWithoutContentLength) {
+  ResourceRequest request = CreateRequest(base::StringPrintf(
+      "/cacheable_without_content_length?body-size=%d", kMaxPerEntrySize + 1));
+  StoreResponseToMemoryCache(request);
+
+  ASSERT_FALSE(CanServeFromMemoryCache(request));
+}
+
 TEST_F(NetworkServiceMemoryCacheTest,
        CanServe_SplitCacheByNetworkIsolationKeyEnabled) {
   base::test::ScopedFeatureList feature_list;

services/network/network_service_memory_cache_writer.cc

@@ -15,12 +15,14 @@ namespace network {
 NetworkServiceMemoryCacheWriter::NetworkServiceMemoryCacheWriter(
     base::WeakPtr<NetworkServiceMemoryCache> cache,
     uint64_t trace_id,
+    size_t max_bytes,
     std::string cache_key,
     net::URLRequest* url_request,
     mojom::RequestDestination request_destination,
     const mojom::URLResponseHeadPtr& response_head)
     : cache_(std::move(cache)),
       trace_id_(trace_id),
+      max_bytes_(max_bytes),
       cache_key_(std::move(cache_key)),
       url_request_(url_request),
       request_destination_(request_destination),
@@ -37,13 +39,18 @@ NetworkServiceMemoryCacheWriter::~NetworkServiceMemoryCacheWriter() {
                                   TRACE_ID_LOCAL(trace_id_));
 }
 
-void NetworkServiceMemoryCacheWriter::OnDataRead(const char* buf, int result) {
+bool NetworkServiceMemoryCacheWriter::OnDataRead(const char* buf, int result) {
+  DCHECK_GT(result, 0);
   TRACE_EVENT_NESTABLE_ASYNC_INSTANT1(
       "loading", "NetworkServiceMemoryCacheWriter::OnDataRead",
       TRACE_ID_LOCAL(trace_id_), "result", result);
 
-  if (result > 0)
-    received_data_.insert(received_data_.end(), buf, buf + result);
+  DCHECK_GE(max_bytes_, received_data_.size());
+  if (static_cast<size_t>(result) > max_bytes_ - received_data_.size())
+    return false;
+
+  received_data_.insert(received_data_.end(), buf, buf + result);
+  return true;
 }
 
 void NetworkServiceMemoryCacheWriter::OnCompleted(

services/network/network_service_memory_cache_writer.h

@@ -32,15 +32,18 @@ class COMPONENT_EXPORT(NETWORK_SERVICE) NetworkServiceMemoryCacheWriter {
   NetworkServiceMemoryCacheWriter(
       base::WeakPtr<NetworkServiceMemoryCache> cache,
       uint64_t trace_id,
+      size_t max_bytes,
       std::string cache_key,
       net::URLRequest* request,
       mojom::RequestDestination request_destination,
       const mojom::URLResponseHeadPtr& response_head);
 
   ~NetworkServiceMemoryCacheWriter();
 
-  // Called when the owner received response content.
-  void OnDataRead(const char* buf, int result);
+  // Called when the owner received response content. Returns false when `this`
+  // doesn't want to duplicate the response any longer (e.g. the content is too
+  // large).
+  bool OnDataRead(const char* buf, int result);
 
   // Called when the owner completed.
   void OnCompleted(const URLLoaderCompletionStatus& status);
@@ -51,6 +54,9 @@ class COMPONENT_EXPORT(NETWORK_SERVICE) NetworkServiceMemoryCacheWriter {
   // Used for tracing.
   const uint64_t trace_id_;
 
+  // The maximum size of `received_data_`.
+  const size_t max_bytes_;
+
   std::string cache_key_;
 
   // `url_request_` must outlive `this`. The owner owns `url_request_`.

services/network/url_loader.cc

@@ -1751,9 +1751,12 @@ void URLLoader::DidRead(int num_bytes, bool completed_synchronously) {
   DCHECK(read_in_progress_);
   read_in_progress_ = false;
 
-  if (memory_cache_writer_ && pending_write_) {
-    memory_cache_writer_->OnDataRead(
-        pending_write_->buffer() + pending_write_buffer_offset_, num_bytes);
+  if (memory_cache_writer_ && pending_write_ && num_bytes > 0) {
+    if (!memory_cache_writer_->OnDataRead(
+            pending_write_->buffer() + pending_write_buffer_offset_,
+            num_bytes)) {
+      memory_cache_writer_.reset();
+    }
   }
 
   size_t new_data_offset = pending_write_buffer_offset_;