4 files changed, 111 insertions, 28 deletions

diff --git a/net/url_request/url_request.cc b/net/url_request/url_request.cc
index 335f974..d6c34f1 100644
--- a/net/url_request/url_request.cc
+++ b/net/url_request/url_request.cc
@@ -352,26 +352,8 @@ std::string URLRequest::StripPostSpecificHeaders(const std::string& headers) {
       "content-length",
       "origin"
   };
-
-  std::string stripped_headers;
-  net::HttpUtil::HeadersIterator it(headers.begin(), headers.end(), "\r\n");
-
-  while (it.GetNext()) {
-    bool is_post_specific = false;
-    for (size_t i = 0; i < ARRAYSIZE_UNSAFE(kPostHeaders); ++i) {
-      if (LowerCaseEqualsASCII(it.name_begin(), it.name_end(),
-                               kPostHeaders[i])) {
-        is_post_specific = true;
-        break;
-      }
-    }
-    if (!is_post_specific) {
-      // Assume that name and values are on the same line.
-      stripped_headers.append(it.name_begin(), it.values_end());
-      stripped_headers.append("\r\n");
-    }
-  }
-  return stripped_headers;
+  return net::HttpUtil::StripHeaders(
+      headers, kPostHeaders, arraysize(kPostHeaders));
 }
 
 int URLRequest::Redirect(const GURL& location, int http_status_code) {
diff --git a/net/url_request/url_request_http_job.cc b/net/url_request/url_request_http_job.cc
index 4180bef..42b8182 100644
--- a/net/url_request/url_request_http_job.cc
+++ b/net/url_request/url_request_http_job.cc
@@ -21,6 +21,7 @@
 #include "net/http/http_response_info.h"
 #include "net/http/http_transaction.h"
 #include "net/http/http_transaction_factory.h"
+#include "net/http/http_util.h"
 #include "net/url_request/url_request.h"
 #include "net/url_request/url_request_context.h"
 #include "net/url_request/url_request_error_job.h"
@@ -299,10 +300,27 @@ void URLRequestHttpJob::SetAuth(const std::wstring& username,
     server_auth_state_ = net::AUTH_STATE_HAVE_AUTH;
   }
 
+  RestartTransactionWithAuth(username, password);
+}
+
+void URLRequestHttpJob::RestartTransactionWithAuth(
+    const std::wstring& username,
+    const std::wstring& password) {
+
   // These will be reset in OnStartCompleted.
   response_info_ = NULL;
   response_cookies_.clear();
 
+  // Update the cookies, since the cookie store may have been updated from the
+  // headers in the 401/407. Since cookies were already appended to
+  // extra_headers by AddExtraHeaders(), we need to strip them out.
+  static const char* const cookie_name[] = { "cookie" };
+  request_info_.extra_headers = net::HttpUtil::StripHeaders(
+      request_info_.extra_headers, cookie_name, arraysize(cookie_name));
+  // TODO(eroman): this ordering is inconsistent with non-restarted request,
+  // where cookies header appears second from the bottom.
+  request_info_.extra_headers += AssembleRequestCookies();
+
   // No matter what, we want to report our status as IO pending since we will
   // be notifying our consumer asynchronously via OnStartCompleted.
   SetStatus(URLRequestStatus(URLRequestStatus::IO_PENDING, 0));
@@ -474,6 +492,15 @@ void URLRequestHttpJob::NotifyHeadersComplete() {
     }
   }
 
+  // The HTTP transaction may be restarted several times for the purposes
+  // of sending authorization information. Each time it restarts, we get
+  // notified of the headers completion so that we can update the cookie store.
+  if (transaction_->IsReadyToRestartForAuth()) {
+    DCHECK(!response_info_->auth_challenge.get());
+    RestartTransactionWithAuth(std::wstring(), std::wstring());
+    return;
+  }
+
   URLRequestJob::NotifyHeadersComplete();
 }
 
@@ -544,24 +571,32 @@ void URLRequestHttpJob::AddExtraHeaders() {
 
   URLRequestContext* context = request_->context();
   if (context) {
+    request_info_.extra_headers += AssembleRequestCookies();
+    if (!context->accept_language().empty())
+      request_info_.extra_headers += "Accept-Language: " +
+          context->accept_language() + "\r\n";
+    if (!context->accept_charset().empty())
+      request_info_.extra_headers += "Accept-Charset: " +
+          context->accept_charset() + "\r\n";
+  }
+}
+
+std::string URLRequestHttpJob::AssembleRequestCookies() {
+  URLRequestContext* context = request_->context();
+  if (context) {
     // Add in the cookie header.  TODO might we need more than one header?
     if (context->cookie_store() &&
         context->cookie_policy()->CanGetCookies(request_->url(),
-                                               request_->policy_url())) {
+                                                request_->policy_url())) {
       net::CookieMonster::CookieOptions options;
       options.set_include_httponly();
       std::string cookies = request_->context()->cookie_store()->
           GetCookiesWithOptions(request_->url(), options);
       if (!cookies.empty())
-        request_info_.extra_headers += "Cookie: " + cookies + "\r\n";
+        return "Cookie: " + cookies + "\r\n";
     }
-    if (!context->accept_language().empty())
-      request_info_.extra_headers += "Accept-Language: " +
-          context->accept_language() + "\r\n";
-    if (!context->accept_charset().empty())
-      request_info_.extra_headers += "Accept-Charset: " +
-          context->accept_charset() + "\r\n";
   }
+  return std::string();
 }
 
 void URLRequestHttpJob::FetchResponseCookies() {
diff --git a/net/url_request/url_request_http_job.h b/net/url_request/url_request_http_job.h
index b165d2e..dbbdfab 100644
--- a/net/url_request/url_request_http_job.h
+++ b/net/url_request/url_request_http_job.h
@@ -63,11 +63,15 @@ class URLRequestHttpJob : public URLRequestJob {
   void DestroyTransaction();
   void StartTransaction();
   void AddExtraHeaders();
+  std::string AssembleRequestCookies();
   void FetchResponseCookies();
 
   void OnStartCompleted(int result);
   void OnReadCompleted(int result);
 
+  void RestartTransactionWithAuth(const std::wstring& username,
+                                  const std::wstring& password);
+
   net::HttpRequestInfo request_info_;
   scoped_ptr<net::HttpTransaction> transaction_;
   const net::HttpResponseInfo* response_info_;
diff --git a/net/url_request/url_request_unittest.cc b/net/url_request/url_request_unittest.cc
index e134208..fb4a650 100644
--- a/net/url_request/url_request_unittest.cc
+++ b/net/url_request/url_request_unittest.cc
@@ -21,6 +21,7 @@
 #include "base/process_util.h"
 #include "base/string_piece.h"
 #include "base/string_util.h"
+#include "net/base/cookie_monster.h"
 #include "net/base/load_flags.h"
 #include "net/base/net_errors.h"
 #include "net/base/net_module.h"
@@ -46,9 +47,12 @@ class URLRequestHttpCacheContext : public URLRequestContext {
     http_transaction_factory_ =
         new net::HttpCache(net::HttpNetworkLayer::CreateFactory(proxy_service_),
                            disk_cache::CreateInMemoryCacheBackend(0));
+    // In-memory cookie store.
+    cookie_store_ = new net::CookieMonster();
   }
 
   virtual ~URLRequestHttpCacheContext() {
+    delete cookie_store_;
     delete http_transaction_factory_;
     delete proxy_service_;
   }
@@ -931,6 +935,64 @@ TEST_F(URLRequestTest, BasicAuth) {
   }
 }
 
+// Check that Set-Cookie headers in 401 responses are respected.
+// http://crbug.com/6450
+TEST_F(URLRequestTest, BasicAuthWithCookies) {
+  scoped_refptr<HTTPTestServer> server =
+      HTTPTestServer::CreateServer(L"", NULL);
+  ASSERT_TRUE(NULL != server.get());
+
+  GURL url_requiring_auth =
+      server->TestServerPage("auth-basic?set-cookie-if-challenged");
+
+  // Request a page that will give a 401 containing a Set-Cookie header.
+  // Verify that when the transaction is restarted, it includes the new cookie.
+  {
+    scoped_refptr<URLRequestContext> context = new URLRequestHttpCacheContext();
+    TestDelegate d;
+    d.set_username(L"user");
+    d.set_password(L"secret");
+
+    URLRequest r(url_requiring_auth, &d);
+    r.set_context(context);
+    r.Start();
+
+    MessageLoop::current()->Run();
+
+    EXPECT_TRUE(d.data_received().find("user/secret") != std::string::npos);
+
+    // Make sure we sent the cookie in the restarted transaction.
+    EXPECT_TRUE(d.data_received().find("Cookie: got_challenged=true")
+        != std::string::npos);
+  }
+
+  // Same test as above, except this time the restart is initiated earlier
+  // (without user intervention since identity is embedded in the URL).
+  {
+    scoped_refptr<URLRequestContext> context = new URLRequestHttpCacheContext();
+    TestDelegate d;
+
+    GURL::Replacements replacements;
+    std::string username("user2");
+    std::string password("secret");
+    replacements.SetUsernameStr(username);
+    replacements.SetPasswordStr(password);
+    GURL url_with_identity = url_requiring_auth.ReplaceComponents(replacements);
+
+    URLRequest r(url_with_identity, &d);
+    r.set_context(context);
+    r.Start();
+
+    MessageLoop::current()->Run();
+
+    EXPECT_TRUE(d.data_received().find("user2/secret") != std::string::npos);
+
+    // Make sure we sent the cookie in the restarted transaction.
+    EXPECT_TRUE(d.data_received().find("Cookie: got_challenged=true")
+        != std::string::npos);
+  }
+}
+
 // In this test, we do a POST which the server will 302 redirect.
 // The subsequent transaction should use GET, and should not send the
 // Content-Type header.