Fix various webRequest events not getting Set-Cookie even with extraHeaders.

This fixes onResponseStarted, onBeforeRedirect, onCompleted and onAuthRequired.

Bug: 935104
Change-Id: Id545fcd4ffb82f1e73d22ad67441b379c3fe31be
Reviewed-on: https://chromium-review.googlesource.com/c/1488444
Commit-Queue: John Abd-El-Malek <jam@chromium.org>
Reviewed-by: Karan Bhatia <karandeepb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#636240}

chrome/browser/extensions/api/web_request/web_request_apitest.cc

@@ -558,6 +558,8 @@ IN_PROC_BROWSER_TEST_F(ExtensionWebRequestApiTest,
 }
 
 IN_PROC_BROWSER_TEST_F(ExtensionWebRequestApiTest, WebRequestExtraHeaders) {
+  CancelLoginDialog login_dialog_helper;
+
   ASSERT_TRUE(StartEmbeddedTestServer());
   ASSERT_TRUE(RunExtensionSubtest("webrequest", "test_extra_headers.html"))
       << message_;

chrome/test/data/extensions/api_test/webrequest/test_extra_headers.js

@@ -59,21 +59,59 @@ runTests([
 
   function testSpecialResponseHeadersVisible() {
     var url = getSetCookieUrl('foo', 'bar');
-    var extraHeadersListener = callbackPass(function(details) {
+    var extraHeadersListenerCalledCount = 0;
+    function extraHeadersListener(details) {
+      extraHeadersListenerCalledCount++;
       checkHeaders(details.responseHeaders, ['set-cookie'], []);
-    });
+    }
     chrome.webRequest.onHeadersReceived.addListener(extraHeadersListener,
         {urls: [url]}, ['responseHeaders', 'extraHeaders']);
+    chrome.webRequest.onResponseStarted.addListener(extraHeadersListener,
+        {urls: [url]}, ['responseHeaders', 'extraHeaders']);
+    chrome.webRequest.onCompleted.addListener(extraHeadersListener,
+        {urls: [url]}, ['responseHeaders', 'extraHeaders']);
 
-    var standardListener = callbackPass(function(details) {
+    var standardListenerCalledCount = 0;
+    function standardListener(details) {
+      standardListenerCalledCount++;
       checkHeaders(details.responseHeaders, [], ['set-cookie']);
-    });
+    }
     chrome.webRequest.onHeadersReceived.addListener(standardListener,
         {urls: [url]}, ['responseHeaders']);
+    chrome.webRequest.onResponseStarted.addListener(standardListener,
+        {urls: [url]}, ['responseHeaders']);
+    chrome.webRequest.onCompleted.addListener(standardListener,
+        {urls: [url]}, ['responseHeaders']);
 
     navigateAndWait(url, function() {
+      chrome.test.assertEq(3, standardListenerCalledCount);
+      chrome.test.assertEq(3, extraHeadersListenerCalledCount);
       chrome.webRequest.onHeadersReceived.removeListener(extraHeadersListener);
+      chrome.webRequest.onResponseStarted.removeListener(extraHeadersListener);
+      chrome.webRequest.onCompleted.removeListener(extraHeadersListener);
       chrome.webRequest.onHeadersReceived.removeListener(standardListener);
+      chrome.webRequest.onResponseStarted.removeListener(standardListener);
+      chrome.webRequest.onCompleted.removeListener(standardListener);
+    });
+  },
+
+  function testSpecialResponseHeadersVisibleForAuth() {
+    var url = getServerURL('auth-basic?set-cookie-if-challenged');
+    var extraHeadersListener = callbackPass(function(details) {
+      checkHeaders(details.responseHeaders, ['set-cookie'], []);
+    });
+    chrome.webRequest.onAuthRequired.addListener(extraHeadersListener,
+        {urls: [url]}, ['responseHeaders', 'extraHeaders']);
+
+    var standardListener = callbackPass(function(details) {
+      checkHeaders(details.responseHeaders, [], ['set-cookie']);
+    });
+    chrome.webRequest.onAuthRequired.addListener(standardListener,
+        {urls: [url]}, ['responseHeaders']);
+
+    navigateAndWait(url, function() {
+      chrome.webRequest.onAuthRequired.removeListener(extraHeadersListener);
+      chrome.webRequest.onAuthRequired.removeListener(standardListener);
     });
   },
 
@@ -100,6 +138,36 @@ runTests([
     });
   },
 
+  function testModifySpecialResponseHeaders() {
+    var url = getSetCookieUrl('foo', 'bar');
+    var headersListener = callbackPass(function(details) {
+      checkHeaders(details.responseHeaders, ['set-cookie'], []);
+      details.responseHeaders.push({name: 'X-New-Header',
+                                    value: 'Foo'});
+      return {responseHeaders: details.responseHeaders};
+    });
+    chrome.webRequest.onHeadersReceived.addListener(headersListener,
+        {urls: [url]}, ['responseHeaders', 'blocking', 'extraHeaders']);
+
+    var responseListener = callbackPass(function(details) {
+      checkHeaders(details.responseHeaders, ['set-cookie', 'x-new-header'], []);
+    });
+    chrome.webRequest.onResponseStarted.addListener(responseListener,
+        {urls: [url]}, ['responseHeaders', 'extraHeaders']);
+
+    var completedListener = callbackPass(function(details) {
+      checkHeaders(details.responseHeaders, ['set-cookie', 'x-new-header'], []);
+    });
+    chrome.webRequest.onCompleted.addListener(completedListener,
+        {urls: [url]}, ['responseHeaders', 'extraHeaders']);
+
+    navigateAndWait(url, function(tab) {
+      chrome.webRequest.onHeadersReceived.removeListener(headersListener);
+      chrome.webRequest.onResponseStarted.removeListener(responseListener);
+      chrome.webRequest.onCompleted.removeListener(completedListener);
+    });
+  },
+
   function testCannotModifySpecialRequestHeadersWithoutExtraHeaders() {
     // Set a cookie so the cookie request header is set.
     navigateAndWait(getSetCookieUrl('foo', 'bar'), function() {

extensions/browser/api/web_request/web_request_proxying_url_loader_factory.cc

@@ -184,11 +184,16 @@ void WebRequestProxyingURLLoaderFactory::InProgressRequest::
 
 void WebRequestProxyingURLLoaderFactory::InProgressRequest::OnReceiveResponse(
     const network::ResourceResponseHead& head) {
-  current_response_ = head;
   on_receive_response_received_ = true;
   if (current_request_uses_header_client_) {
+    // Use the headers we got from OnHeadersReceived as that'll contain
+    // Set-Cookie if it existed.
+    auto saved_headers = current_response_.headers;
+    current_response_ = head;
+    current_response_.headers = saved_headers;
     ContinueToResponseStarted(net::OK);
   } else {
+    current_response_ = head;
     HandleResponseOrRedirectHeaders(
         base::BindRepeating(&InProgressRequest::ContinueToResponseStarted,
                             weak_factory_.GetWeakPtr()));
@@ -205,10 +210,15 @@ void WebRequestProxyingURLLoaderFactory::InProgressRequest::OnReceiveRedirect(
     return;
   }
 
-  current_response_ = head;
   if (current_request_uses_header_client_) {
+    // Use the headers we got from OnHeadersReceived as that'll contain
+    // Set-Cookie if it existed.
+    auto saved_headers = current_response_.headers;
+    current_response_ = head;
+    current_response_.headers = saved_headers;
     ContinueToBeforeRedirect(redirect_info, net::OK);
   } else {
+    current_response_ = head;
     HandleResponseOrRedirectHeaders(
         base::BindRepeating(&InProgressRequest::ContinueToBeforeRedirect,
                             weak_factory_.GetWeakPtr(), redirect_info));
@@ -267,11 +277,17 @@ void WebRequestProxyingURLLoaderFactory::InProgressRequest::HandleAuthRequest(
     WebRequestAPI::AuthRequestCallback callback) {
   DCHECK(!auth_credentials_);
 
+  // If |current_request_uses_header_client_| is true, |current_response_|
+  // should already hold the correct set of response headers (including
+  // Set-Cookie). So we don't use |response_headers| since it won't have the
+  // Set-Cookie headers.
+  if (!current_request_uses_header_client_) {
+    network::ResourceResponseHead head;
+    head.headers = response_headers;
+    current_response_ = head;
+  }
   // We first need to simulate |onHeadersReceived| for the response headers
   // which indicated a need to authenticate.
-  network::ResourceResponseHead head;
-  head.headers = response_headers;
-  current_response_ = head;
   HandleResponseOrRedirectHeaders(base::BindRepeating(base::BindRepeating(
       &InProgressRequest::ContinueAuthRequest, weak_factory_.GetWeakPtr(),
       base::RetainedRef(auth_info), base::Passed(&callback))));
@@ -556,8 +572,15 @@ void WebRequestProxyingURLLoaderFactory::InProgressRequest::
 
   DCHECK(on_headers_received_callback_);
   base::Optional<std::string> headers;
-  if (override_headers_)
+  if (override_headers_) {
     headers = override_headers_->raw_headers();
+    if (current_request_uses_header_client_) {
+      // Make sure to update current_response_,  since when OnReceiveResponse
+      // is called we will not use its headers as it might be missing the
+      // Set-Cookie line (as that gets stripped over IPC).
+      current_response_.headers = override_headers_;
+    }
+  }
   std::move(on_headers_received_callback_).Run(net::OK, headers, redirect_url_);
   override_headers_ = nullptr;