Return chrome-proxy-ect header code (post M59 branch)

components/data_reduction_proxy/core/browser/data_reduction_proxy_network_delegate_unittest.cc

Index: components/data_reduction_proxy/core/browser/data_reduction_proxy_network_delegate_unittest.cc
diff --git a/components/data_reduction_proxy/core/browser/data_reduction_proxy_network_delegate_unittest.cc b/components/data_reduction_proxy/core/browser/data_reduction_proxy_network_delegate_unittest.cc
index 756102e5bbc498535070b98b71211139955014c6..c52934818a7600fc66aec0a430bded8bb400812a 100644
--- a/components/data_reduction_proxy/core/browser/data_reduction_proxy_network_delegate_unittest.cc
+++ b/components/data_reduction_proxy/core/browser/data_reduction_proxy_network_delegate_unittest.cc
@@ -273,6 +273,79 @@ class DataReductionProxyNetworkDelegateTest : public testing::Test {
     test_context_->EnableDataReductionProxyWithSecureProxyCheckSuccess();
   }
 
+  // Build the sockets by adding appropriate mock data for
+  // |effective_connection_types.size()| number of requests. Data for
+  // chrome-Proxy-ect header is added to the mock data if |expect_ect_header|
+  // is true. |reads_list|, |mock_writes| and |writes_list| should be empty, and
+  // are owned by the caller.
+  void BuildSocket(const std::string& response_headers,
+                   const std::string& response_body,
+                   bool expect_ect_header,
+                   const std::vector<net::EffectiveConnectionType>&
+                       effective_connection_types,
+                   std::vector<net::MockRead>* reads_list,
+                   std::vector<std::string>* mock_writes,
+                   std::vector<net::MockWrite>* writes_list) {
+    EXPECT_LT(0u, effective_connection_types.size());
+    EXPECT_TRUE(reads_list->empty());
+    EXPECT_TRUE(mock_writes->empty());
+    EXPECT_TRUE(writes_list->empty());
+
+    for (size_t i = 0; i < effective_connection_types.size(); ++i) {
+      reads_list->push_back(net::MockRead(response_headers.c_str()));
+      reads_list->push_back(net::MockRead(response_body.c_str()));
+    }
+    reads_list->push_back(net::MockRead(net::SYNCHRONOUS, net::OK));
+
+    std::string prefix = std::string("GET ")
+                             .append(kTestURL)
+                             .append(" HTTP/1.1\r\n")
+                             .append("Host: ")
+                             .append(GURL(kTestURL).host())
+                             .append(
+                                 "\r\n"
+                                 "Proxy-Connection: keep-alive\r\n"
+                                 "User-Agent:\r\n"
+                                 "Accept-Encoding: gzip, deflate\r\n"
+                                 "Accept-Language: en-us,fr\r\n");
+
+    if (io_data()->test_request_options()->GetHeaderValueForTesting().empty()) {
+      // Force regeneration of Chrome-Proxy header.
+      io_data()->test_request_options()->SetSecureSession("123");
+    }
+
+    EXPECT_FALSE(
+        io_data()->test_request_options()->GetHeaderValueForTesting().empty());
+
+    std::string suffix =
+        std::string("Chrome-Proxy: ") +
+        io_data()->test_request_options()->GetHeaderValueForTesting() +
+        std::string("\r\n\r\n");
+
+    mock_socket_factory_.AddSSLSocketDataProvider(&ssl_socket_data_provider_);
+
+    for (net::EffectiveConnectionType effective_connection_type :
+         effective_connection_types) {
+      std::string ect_header;
+      if (expect_ect_header) {
+        ect_header = "chrome-proxy-ect: " +
+                     std::string(net::GetNameForEffectiveConnectionType(
+                         effective_connection_type)) +
+                     "\r\n";
+      }
+
+      std::string mock_write = prefix + ect_header + suffix;
+      mock_writes->push_back(mock_write);
+      writes_list->push_back(net::MockWrite(mock_writes->back().c_str()));
+    }
+
+    EXPECT_FALSE(socket_);
+    socket_ = base::MakeUnique<net::StaticSocketDataProvider>(
+        reads_list->data(), reads_list->size(), writes_list->data(),
+        writes_list->size());
+    mock_socket_factory_.AddSocketDataProvider(socket_.get());
+  }
+
   static void VerifyHeaders(bool expected_data_reduction_proxy_used,
                             bool expected_lofi_used,
                             const net::HttpRequestHeaders& headers) {
@@ -399,6 +472,10 @@ class DataReductionProxyNetworkDelegateTest : public testing::Test {
                                          "User-Agent:\r\n");
 
     std::string accept_language_header("Accept-Language: en-us,fr\r\n");
+    std::string ect_header = "chrome-proxy-ect: " +
+                             std::string(net::GetNameForEffectiveConnectionType(
+                                 net::EFFECTIVE_CONNECTION_TYPE_UNKNOWN)) +
+                             "\r\n";
 
     // Brotli is included in accept-encoding header only if the request went
     // to the network (i.e., it was not a cached response), and if data
@@ -414,13 +491,14 @@ class DataReductionProxyNetworkDelegateTest : public testing::Test {
         std::string("\r\n\r\n");
 
     std::string mock_write = prefix_headers + accept_language_header +
-                             accept_encoding_header + suffix_headers;
+                             ect_header + accept_encoding_header +
+                             suffix_headers;
 
     if (expect_cached || !expect_brotli) {
       // Order of headers is different if the headers were modified by data
       // reduction proxy network delegate.
       mock_write = prefix_headers + accept_encoding_header +
-                   accept_language_header + suffix_headers;
+                   accept_language_header + ect_header + suffix_headers;
     }
 
     net::MockWrite writes[] = {net::MockWrite(mock_write.c_str())};
@@ -516,6 +594,7 @@ class DataReductionProxyNetworkDelegateTest : public testing::Test {
         "www.google.com\r\nProxy-Connection: "
         "keep-alive\r\nUser-Agent:\r\nAccept-Encoding: gzip, "
         "deflate\r\nAccept-Language: en-us,fr\r\n"
+        "chrome-proxy-ect: 4G\r\n"
         "Chrome-Proxy: " +
         io_data()->test_request_options()->GetHeaderValueForTesting() +
         (page_id_value.empty() ? "" : (", " + page_id_value)) + "\r\n\r\n";
@@ -549,6 +628,51 @@ class DataReductionProxyNetworkDelegateTest : public testing::Test {
     base::RunLoop().RunUntilIdle();
   }
 
+  // Fetches a request while the effective connection type is set to
+  // |effective_connection_type|. Verifies that the request headers include the
+  // chrome-proxy-ect header only if |expect_ect_header| is true. The response
+  // must be served from the cache if |expect_cached| is true.
+  void FetchURLRequestAndVerifyECTHeader(
+      net::EffectiveConnectionType effective_connection_type,
+      bool expect_ect_header,
+      bool expect_cached) {
+    test_network_quality_estimator()->set_effective_connection_type(
+        effective_connection_type);
+
+    net::TestDelegate delegate;
+    std::unique_ptr<net::URLRequest> request =
+        context_.CreateRequest(GURL(kTestURL), net::IDLE, &delegate);
+
+    request->Start();
+    base::RunLoop().RunUntilIdle();
+
+    EXPECT_EQ(140, request->received_response_content_length());
+    EXPECT_EQ(expect_cached, request->was_cached());
+    EXPECT_EQ(expect_cached, request->GetTotalSentBytes() == 0);
+    EXPECT_EQ(expect_cached, request->GetTotalReceivedBytes() == 0);
+
+    net::HttpRequestHeaders sent_request_headers;
+    EXPECT_NE(expect_cached,
+              request->GetFullRequestHeaders(&sent_request_headers));
+
+    if (expect_cached) {
+      // Request headers are missing. Return since there is nothing left to
+      // check.
+      return;
+    }
+
+    // Verify that chrome-proxy-ect header is present in the request headers
+    // only if |expect_ect_header| is true.
+    std::string ect_value;
+    EXPECT_EQ(expect_ect_header, sent_request_headers.GetHeader(
+                                     chrome_proxy_ect_header(), &ect_value));
+
+    if (!expect_ect_header)
+      return;
+    EXPECT_EQ(net::GetNameForEffectiveConnectionType(effective_connection_type),
+              ect_value);
+  }
+
   void DelegateStageDone(int result) {}
 
   void NotifyNetworkDelegate(net::URLRequest* request,
@@ -1376,6 +1500,11 @@ TEST_F(DataReductionProxyNetworkDelegateTest,
   net::ProxyRetryInfoMap proxy_retry_info;
 
   // Send a request and verify the page ID is 1.
+  network_delegate()->NotifyBeforeStartTransaction(
+      request.get(),
+      base::Bind(&DataReductionProxyNetworkDelegateTest::DelegateStageDone,
+                 base::Unretained(this)),
+      &headers);
   network_delegate()->NotifyBeforeSendHeaders(
       request.get(), data_reduction_proxy_info, proxy_retry_info, &headers);
   DataReductionProxyData* data =
@@ -1388,6 +1517,11 @@ TEST_F(DataReductionProxyNetworkDelegateTest,
                                      nullptr);
   request->SetLoadFlags(net::LOAD_MAIN_FRAME_DEPRECATED);
 
+  network_delegate()->NotifyBeforeStartTransaction(
+      request.get(),
+      base::Bind(&DataReductionProxyNetworkDelegateTest::DelegateStageDone,
+                 base::Unretained(this)),
+      &headers);
   network_delegate()->NotifyBeforeSendHeaders(
       request.get(), data_reduction_proxy_info, proxy_retry_info, &headers);
   data = DataReductionProxyData::GetData(*request.get());
@@ -1409,6 +1543,89 @@ TEST_F(DataReductionProxyNetworkDelegateTest,
   EXPECT_EQ(1u, data->page_id().value());
 }
 
+// Test that effective connection type is correctly added to the request
+// headers when it is enabled using field trial. The server is varying on the
+// effective connection type (ECT).
+TEST_F(DataReductionProxyNetworkDelegateTest, ECTHeaderEnabledWithVary) {
+  Init(USE_SECURE_PROXY, false /* enable_brotli_globally */);
+
+  std::string response_headers =
+      "HTTP/1.1 200 OK\r\n"
+      "Content-Length: 140\r\n"
+      "Via: 1.1 Chrome-Compression-Proxy\r\n"
+      "Cache-Control: max-age=1200\r\n"
+      "Vary: chrome-proxy-ect\r\n"
+      "x-original-content-length: 200\r\n\r\n";
+
+  int response_body_size = 140;
+  std::string response_body(base::checked_cast<size_t>(response_body_size),
+                            ' ');
+
+  std::vector<net::MockRead> reads_list;
+  std::vector<std::string> mock_writes;
+  std::vector<net::MockWrite> writes_list;
+
+  std::vector<net::EffectiveConnectionType> effective_connection_types;
+  effective_connection_types.push_back(net::EFFECTIVE_CONNECTION_TYPE_SLOW_2G);
+  effective_connection_types.push_back(net::EFFECTIVE_CONNECTION_TYPE_2G);
+
+  BuildSocket(response_headers, response_body, true, effective_connection_types,
+              &reads_list, &mock_writes, &writes_list);
+
+  // Add 2 socket providers since 2 requests in this test are fetched from the
+  // network.
+  FetchURLRequestAndVerifyECTHeader(effective_connection_types[0], true, false);
+
+  // When the ECT is set to the same value, fetching the same resource should
+  // result in a cache hit.
+  FetchURLRequestAndVerifyECTHeader(effective_connection_types[0], true, true);
+
+  // When the ECT is set to a different value, the response should not be
+  // served from the cache.
+  FetchURLRequestAndVerifyECTHeader(effective_connection_types[1], true, false);
+}
+
+// Test that effective connection type is correctly added to the request
+// headers when it is enabled using field trial. The server is not varying on
+// the effective connection type (ECT).
+TEST_F(DataReductionProxyNetworkDelegateTest, ECTHeaderEnabledWithoutVary) {
+  Init(USE_SECURE_PROXY, false /* enable_brotli_globally */);
+
+  std::string response_headers =
+      "HTTP/1.1 200 OK\r\n"
+      "Content-Length: 140\r\n"
+      "Via: 1.1 Chrome-Compression-Proxy\r\n"
+      "Cache-Control: max-age=1200\r\n"
+      "x-original-content-length: 200\r\n\r\n";
+
+  int response_body_size = 140;
+  std::string response_body(base::checked_cast<size_t>(response_body_size),
+                            ' ');
+
+  std::vector<net::MockRead> reads_list;
+  std::vector<std::string> mock_writes;
+  std::vector<net::MockWrite> writes_list;
+
+  std::vector<net::EffectiveConnectionType> effective_connection_types;
+  effective_connection_types.push_back(net::EFFECTIVE_CONNECTION_TYPE_SLOW_2G);
+  effective_connection_types.push_back(net::EFFECTIVE_CONNECTION_TYPE_2G);
+
+  BuildSocket(response_headers, response_body, true, effective_connection_types,
+              &reads_list, &mock_writes, &writes_list);
+
+  // Add 1 socket provider since 1 request in this test is fetched from the
+  // network.
+  FetchURLRequestAndVerifyECTHeader(effective_connection_types[0], true, false);
+
+  // When the ECT is set to the same value, fetching the same resource should
+  // result in a cache hit.
+  FetchURLRequestAndVerifyECTHeader(effective_connection_types[0], true, true);
+
+  // When the ECT is set to a different value, the response should still be
+  // served from the cache.
+  FetchURLRequestAndVerifyECTHeader(effective_connection_types[1], true, true);
+}
+
 }  // namespace
 
 }  // namespace data_reduction_proxy