Add QuicMemoryCacheBackend::AddResponseWithEarlyHints()

This method will be used to write unittests/browser_tests in
the Chromium repository.

PiperOrigin-RevId: 363061099
Change-Id: Iff231ffc8f5192e57982f69b629c6cb8e7d1cd8a
diff --git a/quic/tools/quic_backend_response.h b/quic/tools/quic_backend_response.h
index a7edd78..1c5cad8 100644
--- a/quic/tools/quic_backend_response.h
+++ b/quic/tools/quic_backend_response.h
@@ -50,11 +50,20 @@
 
   ~QuicBackendResponse();
 
+  const std::vector<spdy::Http2HeaderBlock>& early_hints() const {
+    return early_hints_;
+  }
   SpecialResponseType response_type() const { return response_type_; }
   const spdy::Http2HeaderBlock& headers() const { return headers_; }
   const spdy::Http2HeaderBlock& trailers() const { return trailers_; }
   const absl::string_view body() const { return absl::string_view(body_); }
 
+  void AddEarlyHints(const spdy::Http2HeaderBlock& headers) {
+    spdy::Http2HeaderBlock hints = headers.Clone();
+    hints[":status"] = "103";
+    early_hints_.push_back(std::move(hints));
+  }
+
   void set_response_type(SpecialResponseType response_type) {
     response_type_ = response_type;
   }
@@ -70,6 +79,7 @@
   }
 
  private:
+  std::vector<spdy::Http2HeaderBlock> early_hints_;
   SpecialResponseType response_type_;
   spdy::Http2HeaderBlock headers_;
   spdy::Http2HeaderBlock trailers_;
diff --git a/quic/tools/quic_memory_cache_backend.cc b/quic/tools/quic_memory_cache_backend.cc
index 1e116ca..f1acc4a 100644
--- a/quic/tools/quic_memory_cache_backend.cc
+++ b/quic/tools/quic_memory_cache_backend.cc
@@ -199,7 +199,7 @@
                                          absl::string_view response_body) {
   AddResponseImpl(host, path, QuicBackendResponse::REGULAR_RESPONSE,
                   std::move(response_headers), response_body,
-                  Http2HeaderBlock());
+                  Http2HeaderBlock(), std::vector<spdy::Http2HeaderBlock>());
 }
 
 void QuicMemoryCacheBackend::AddResponse(absl::string_view host,
@@ -209,7 +209,19 @@
                                          Http2HeaderBlock response_trailers) {
   AddResponseImpl(host, path, QuicBackendResponse::REGULAR_RESPONSE,
                   std::move(response_headers), response_body,
-                  std::move(response_trailers));
+                  std::move(response_trailers),
+                  std::vector<spdy::Http2HeaderBlock>());
+}
+
+void QuicMemoryCacheBackend::AddResponseWithEarlyHints(
+    absl::string_view host,
+    absl::string_view path,
+    spdy::Http2HeaderBlock response_headers,
+    absl::string_view response_body,
+    const std::vector<spdy::Http2HeaderBlock>& early_hints) {
+  AddResponseImpl(host, path, QuicBackendResponse::REGULAR_RESPONSE,
+                  std::move(response_headers), response_body,
+                  Http2HeaderBlock(), early_hints);
 }
 
 void QuicMemoryCacheBackend::AddSpecialResponse(
@@ -217,7 +229,7 @@
     absl::string_view path,
     SpecialResponseType response_type) {
   AddResponseImpl(host, path, response_type, Http2HeaderBlock(), "",
-                  Http2HeaderBlock());
+                  Http2HeaderBlock(), std::vector<spdy::Http2HeaderBlock>());
 }
 
 void QuicMemoryCacheBackend::AddSpecialResponse(
@@ -227,7 +239,8 @@
     absl::string_view response_body,
     SpecialResponseType response_type) {
   AddResponseImpl(host, path, response_type, std::move(response_headers),
-                  response_body, Http2HeaderBlock());
+                  response_body, Http2HeaderBlock(),
+                  std::vector<spdy::Http2HeaderBlock>());
 }
 
 QuicMemoryCacheBackend::QuicMemoryCacheBackend() : cache_initialized_(false) {}
@@ -362,7 +375,8 @@
     SpecialResponseType response_type,
     Http2HeaderBlock response_headers,
     absl::string_view response_body,
-    Http2HeaderBlock response_trailers) {
+    Http2HeaderBlock response_trailers,
+    const std::vector<spdy::Http2HeaderBlock>& early_hints) {
   QuicWriterMutexLock lock(&response_mutex_);
 
   QUICHE_DCHECK(!host.empty())
@@ -378,6 +392,9 @@
   new_response->set_headers(std::move(response_headers));
   new_response->set_body(response_body);
   new_response->set_trailers(std::move(response_trailers));
+  for (auto& headers : early_hints) {
+    new_response->AddEarlyHints(headers);
+  }
   QUIC_DVLOG(1) << "Add response with key " << key;
   responses_[key] = std::move(new_response);
 }
diff --git a/quic/tools/quic_memory_cache_backend.h b/quic/tools/quic_memory_cache_backend.h
index 5e4b267..8ba6d60 100644
--- a/quic/tools/quic_memory_cache_backend.h
+++ b/quic/tools/quic_memory_cache_backend.h
@@ -110,6 +110,14 @@
                    absl::string_view response_body,
                    spdy::Http2HeaderBlock response_trailers);
 
+  // Add a response, with 103 Early Hints, to the cache.
+  void AddResponseWithEarlyHints(
+      absl::string_view host,
+      absl::string_view path,
+      spdy::Http2HeaderBlock response_headers,
+      absl::string_view response_body,
+      const std::vector<spdy::Http2HeaderBlock>& early_hints);
+
   // Simulate a special behavior at a particular path.
   void AddSpecialResponse(
       absl::string_view host,
@@ -152,7 +160,8 @@
                        QuicBackendResponse::SpecialResponseType response_type,
                        spdy::Http2HeaderBlock response_headers,
                        absl::string_view response_body,
-                       spdy::Http2HeaderBlock response_trailers);
+                       spdy::Http2HeaderBlock response_trailers,
+                       const std::vector<spdy::Http2HeaderBlock>& early_hints);
 
   std::string GetKey(absl::string_view host, absl::string_view path) const;
 
diff --git a/quic/tools/quic_simple_server_stream.cc b/quic/tools/quic_simple_server_stream.cc
index fef2d3e..6bc93cb 100644
--- a/quic/tools/quic_simple_server_stream.cc
+++ b/quic/tools/quic_simple_server_stream.cc
@@ -206,6 +206,13 @@
     return;
   }
 
+  // Send Early Hints first.
+  for (const auto& headers : response->early_hints()) {
+    QUIC_DVLOG(1) << "Stream " << id() << " sending an Early Hints response: "
+                  << headers.DebugString();
+    WriteHeaders(headers.Clone(), false, nullptr);
+  }
+
   if (response->response_type() == QuicBackendResponse::CLOSE_CONNECTION) {
     QUIC_DVLOG(1) << "Special response: closing connection.";
     OnUnrecoverableError(QUIC_NO_ERROR, "Toy server forcing close");
diff --git a/quic/tools/quic_simple_server_stream_test.cc b/quic/tools/quic_simple_server_stream_test.cc
index 13b0335..dd3b8cd 100644
--- a/quic/tools/quic_simple_server_stream_test.cc
+++ b/quic/tools/quic_simple_server_stream_test.cc
@@ -59,12 +59,17 @@
   ~TestStream() override = default;
 
   MOCK_METHOD(void, WriteHeadersMock, (bool fin), ());
+  MOCK_METHOD(void, WriteEarlyHintsHeadersMock, (bool fin), ());
 
-  size_t WriteHeaders(spdy::Http2HeaderBlock /*header_block*/,
+  size_t WriteHeaders(spdy::Http2HeaderBlock header_block,
                       bool fin,
                       QuicReferenceCountedPointer<QuicAckListenerInterface>
                       /*ack_listener*/) override {
-    WriteHeadersMock(fin);
+    if (header_block[":status"] == "103") {
+      WriteEarlyHintsHeadersMock(fin);
+    } else {
+      WriteHeadersMock(fin);
+    }
     return 0;
   }
 
@@ -581,6 +586,50 @@
   EXPECT_EQ(*request_headers, session_.original_request_headers_);
 }
 
+TEST_P(QuicSimpleServerStreamTest, SendResponseWithEarlyHints) {
+  std::string host = "www.google.com";
+  std::string request_path = "/foo";
+  std::string body = "Yummm";
+
+  // Add a request and response with early hints.
+  spdy::Http2HeaderBlock* request_headers = stream_->mutable_headers();
+  (*request_headers)[":path"] = request_path;
+  (*request_headers)[":authority"] = host;
+  (*request_headers)[":method"] = "GET";
+
+  std::unique_ptr<char[]> buffer;
+  QuicByteCount header_length =
+      HttpEncoder::SerializeDataFrameHeader(body.length(), &buffer);
+  std::vector<spdy::Http2HeaderBlock> early_hints;
+  // Add two Early Hints.
+  const size_t kNumEarlyHintsResponses = 2;
+  for (size_t i = 0; i < kNumEarlyHintsResponses; ++i) {
+    spdy::Http2HeaderBlock hints;
+    hints["link"] = "</image.png>; rel=preload; as=image";
+    early_hints.push_back(std::move(hints));
+  }
+
+  response_headers_[":status"] = "200";
+  response_headers_["content-length"] = "5";
+  memory_cache_backend_.AddResponseWithEarlyHints(
+      host, request_path, std::move(response_headers_), body, early_hints);
+  QuicStreamPeer::SetFinReceived(stream_);
+
+  InSequence s;
+  for (size_t i = 0; i < kNumEarlyHintsResponses; ++i) {
+    EXPECT_CALL(*stream_, WriteEarlyHintsHeadersMock(false));
+  }
+  EXPECT_CALL(*stream_, WriteHeadersMock(false));
+  if (UsesHttp3()) {
+    EXPECT_CALL(session_, WritevData(_, header_length, _, NO_FIN, _, _));
+  }
+  EXPECT_CALL(session_, WritevData(_, body.length(), _, FIN, _, _));
+
+  stream_->DoSendResponse();
+  EXPECT_FALSE(QuicStreamPeer::read_side_closed(stream_));
+  EXPECT_TRUE(stream_->write_side_closed());
+}
+
 TEST_P(QuicSimpleServerStreamTest, PushResponseOnClientInitiatedStream) {
   // EXPECT_QUIC_BUG tests are expensive so only run one instance of them.
   if (GetParam() != AllSupportedVersions()[0]) {