Support HttpOnly cookie on Web Socket

net/websockets/websocket_job.cc

Index: net/websockets/websocket_job.cc
diff --git a/net/websockets/websocket_job.cc b/net/websockets/websocket_job.cc
new file mode 100644
index 0000000000000000000000000000000000000000..3ba0c36d6eb8b8fb9caa87f53a20f3f26861ca8a
--- /dev/null
+++ b/net/websockets/websocket_job.cc
@@ -0,0 +1,378 @@
+// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "net/websockets/websocket_job.h"
+
+#include "googleurl/src/gurl.h"
+#include "net/base/net_errors.h"
+#include "net/base/cookie_policy.h"
+#include "net/base/cookie_store.h"
+#include "net/http/http_util.h"
+#include "net/url_request/url_request_context.h"
+
+namespace net {
+
+// lower-case header names.
+static const char* const kCookieHeaders[] = {
+  "cookie", "cookie2"
+};
+static const char* const kSetCookieHeaders[] = {
+  "set-cookie", "set-cookie2"
+};
+
+static SocketStreamJob* WebSocketJobFactory(
+    const GURL& url, SocketStream::Delegate* delegate) {
+  WebSocketJob* job = new WebSocketJob(delegate);
+  job->InitSocketStream(new SocketStream(url, job));
+  return job;
+}
+
+class WebSocketJobInitSingleton {
+ private:
+  friend struct DefaultSingletonTraits<WebSocketJobInitSingleton>;
+  WebSocketJobInitSingleton() {
+    SocketStreamJob::RegisterProtocolFactory("ws", WebSocketJobFactory);
+    SocketStreamJob::RegisterProtocolFactory("wss", WebSocketJobFactory);
+  }
+};
+
+static void ParseHandshakeMessage(
+    const char* handshake_message, int len,
+    std::string* status_line,
+    std::string* header) {
+  size_t i = base::StringPiece(handshake_message, len).find_first_of("\r\n");
+  if (i == base::StringPiece::npos) {
+    *status_line = std::string(handshake_message, len);
+    *header = "";
+    return;
+  }
+  *status_line = std::string(handshake_message, i + 2);
+  *header = std::string(handshake_message + i + 2, len - i - 2);
+}
+
+static void FetchResponseCookies(
+    const char* handshake_message, int len,
+    std::vector<std::string>* response_cookies) {
+  std::string handshake_response(handshake_message, len);
+  HttpUtil::HeadersIterator iter(handshake_response.begin(),
+                                 handshake_response.end(), "\r\n");
+  while (iter.GetNext()) {
+    for (size_t i = 0; i < arraysize(kSetCookieHeaders); i++) {
+      if (LowerCaseEqualsASCII(iter.name_begin(), iter.name_end(),
+                               kSetCookieHeaders[i])) {
+        response_cookies->push_back(iter.values());
+      }
+    }
+  }
+}
+
+// static
+void WebSocketJob::EnsureInit() {
+  Singleton<WebSocketJobInitSingleton>::get();
+}
+
+WebSocketJob::WebSocketJob(SocketStream::Delegate* delegate)
+    : delegate_(delegate),
+      state_(INITIALIZED),
+      handshake_request_sent_(0),
+      handshake_response_header_length_(0),
+      response_cookies_save_index_(0),
+      ALLOW_THIS_IN_INITIALIZER_LIST(can_get_cookies_callback_(
+          this, &WebSocketJob::OnCanGetCookiesCompleted)),
+      ALLOW_THIS_IN_INITIALIZER_LIST(can_set_cookie_callback_(
+          this, &WebSocketJob::OnCanSetCookieCompleted)) {
+}
+
+WebSocketJob::~WebSocketJob() {
+  DCHECK_EQ(CLOSED, state_);
+  DCHECK(!delegate_);
+  DCHECK(!socket_.get());
+}
+
+void WebSocketJob::Connect() {
+  DCHECK(socket_.get());
+  DCHECK_EQ(state_, INITIALIZED);
+  state_ = CONNECTING;
+  socket_->Connect();
+}
+
+bool WebSocketJob::SendData(const char* data, int len) {
+  switch (state_) {
+    case INITIALIZED:
+      return false;
+
+    case CONNECTING:
+      return SendHandshakeRequest(data, len);
+
+    case OPEN:
+      return socket_->SendData(data, len);
+
+    case CLOSED:
+      return false;
+  }
+  return false;
+}
+
+void WebSocketJob::Close() {
+  state_ = CLOSED;
+  socket_->Close();
+}
+
+void WebSocketJob::RestartWithAuth(
+    const std::wstring& username,
+    const std::wstring& password) {
+  state_ = CONNECTING;
+  socket_->RestartWithAuth(username, password);
+}
+
+void WebSocketJob::DetachDelegate() {
+  state_ = CLOSED;
+  delegate_ = NULL;
+  socket_->DetachDelegate();
+  socket_ = NULL;
+}
+
+void WebSocketJob::OnConnected(
+    SocketStream* socket, int max_pending_send_allowed) {
+  if (delegate_)
+    delegate_->OnConnected(socket, max_pending_send_allowed);
+}
+
+void WebSocketJob::OnSentData(SocketStream* socket, int amount_sent) {
+  if (state_ == CONNECTING) {
+    OnSentHandshakeRequest(socket, amount_sent);
+    return;
+  }
+  if (delegate_)
+    delegate_->OnSentData(socket, amount_sent);
+}
+
+void WebSocketJob::OnReceivedData(
+    SocketStream* socket, const char* data, int len) {
+  if (state_ == CONNECTING) {
+    OnReceivedHandshakeResponse(socket, data, len);
+    return;
+  }
+  if (delegate_)
+    delegate_->OnReceivedData(socket, data, len);
+}
+
+void WebSocketJob::OnClose(SocketStream* socket) {
+  state_ = CLOSED;
+  SocketStream::Delegate* delegate = delegate_;
+  delegate_ = NULL;
+  socket_ = NULL;
+  if (delegate)
+    delegate->OnClose(socket);
+}
+
+void WebSocketJob::OnAuthRequired(
+    SocketStream* socket, AuthChallengeInfo* auth_info) {
+  if (delegate_)
+    delegate_->OnAuthRequired(socket, auth_info);
+}
+
+void WebSocketJob::OnError(const SocketStream* socket, int error) {
+  if (delegate_)
+    delegate_->OnError(socket, error);
+}
+
+bool WebSocketJob::SendHandshakeRequest(const char* data, int len) {
+  DCHECK_EQ(state_, CONNECTING);
+  if (!handshake_request_.empty()) {
+    // if we're already sending handshake message, don't send any more data
+    // until handshake is completed.
+    return false;
+  }
+  original_handshake_request_.append(data, len);
+  original_handshake_request_header_length_ =
+      HttpUtil::LocateEndOfHeaders(original_handshake_request_.data(),
+                                   original_handshake_request_.size(), 0);
+  if (original_handshake_request_header_length_ > 0) {
+    // handshake message is completed.
+    AddCookieHeaderAndSend();
+  }
+  // Just buffered in original_handshake_request_.
+  return true;
+}
+
+void WebSocketJob::AddCookieHeaderAndSend() {
+  AddRef();  // Balanced in OnCanGetCookiesCompleted
+
+  int policy = OK;
+  if (socket_->context()->cookie_policy()) {
+    GURL url_for_cookies = GetURLForCookies();
+    policy = socket_->context()->cookie_policy()->CanGetCookies(
+        url_for_cookies,
+        url_for_cookies,
+        &can_get_cookies_callback_);
+    if (policy == ERR_IO_PENDING)
+      return;  // Wait for completion callback
+  }
+  OnCanGetCookiesCompleted(policy);
+}
+
+void WebSocketJob::OnCanGetCookiesCompleted(int policy) {
+  if (socket_ && delegate_ && state_ == CONNECTING) {
+    std::string handshake_request_status_line;
+    std::string handshake_request_header;
+    ParseHandshakeMessage(original_handshake_request_.data(),
+                          original_handshake_request_header_length_,
+                          &handshake_request_status_line,
+                          &handshake_request_header);
+
+    // Remove cookie headers.
+    handshake_request_header = HttpUtil::StripHeaders(
+        handshake_request_header,
+        kCookieHeaders, arraysize(kCookieHeaders));
+
+    if (policy == OK) {
+      // Add cookies, including HttpOnly cookies.
+      if (socket_->context()->cookie_store()) {
+        CookieOptions cookie_options;
+        cookie_options.set_include_httponly();
+        std::string cookie =
+            socket_->context()->cookie_store()->GetCookiesWithOptions(
+                GetURLForCookies(), cookie_options);
+        if (!cookie.empty()) {
+          HttpUtil::AppendHeaderIfMissing("Cookie", cookie,
+                                          &handshake_request_header);
+        }
+      }
+    }
+
+    // Simply ignore rest data in original request header after
+    // original_handshake_request_header_length_, because websocket protocol
+    // doesn't allow sending message before handshake is completed.
+    // TODO(ukai): report as error?
+    handshake_request_ =
+          handshake_request_status_line + handshake_request_header + "\r\n";
+
+    handshake_request_sent_ = 0;
+    socket_->SendData(handshake_request_.data(),
+                      handshake_request_.size());
+  }
+  Release();  // Balance AddRef taken in AddCookieHeaderAndSend
+}
+
+void WebSocketJob::OnSentHandshakeRequest(
+    SocketStream* socket, int amount_sent) {
+  DCHECK_EQ(state_, CONNECTING);
+  handshake_request_sent_ += amount_sent;
+  if (handshake_request_sent_ >= handshake_request_.size()) {
+    // handshake request has been sent.
+    // notify original size of handshake request to delegate.
+    if (delegate_)
+      delegate_->OnSentData(socket, original_handshake_request_.size());
+  }
+}
+
+void WebSocketJob::OnReceivedHandshakeResponse(
+    SocketStream* socket, const char* data, int len) {
+  DCHECK_EQ(state_, CONNECTING);
+  handshake_response_.append(data, len);
+  handshake_response_header_length_ = HttpUtil::LocateEndOfHeaders(
+      handshake_response_.data(),
+      handshake_response_.size(), 0);
+  if (handshake_response_header_length_ > 0) {
+    // handshake message is completed.
+    SaveCookiesAndNotifyHeaderComplete();
+  }
+}
+
+void WebSocketJob::SaveCookiesAndNotifyHeaderComplete() {
+  // handshake message is completed.
+  DCHECK(handshake_response_.data());
+  DCHECK_GT(handshake_response_header_length_, 0);
+
+  response_cookies_.clear();
+  response_cookies_save_index_ = 0;
+
+  FetchResponseCookies(handshake_response_.data(),
+                       handshake_response_header_length_,
+                       &response_cookies_);
+
+  // Now, loop over the response cookies, and attempt to persist each.
+  SaveNextCookie();
+}
+
+void WebSocketJob::SaveNextCookie() {
+  if (response_cookies_save_index_ == response_cookies_.size()) {
+    response_cookies_.clear();
+    response_cookies_save_index_ = 0;
+
+    std::string handshake_response_status_line;
+    std::string handshake_response_header;
+    ParseHandshakeMessage(handshake_response_.data(),
+                          handshake_response_header_length_,
+                          &handshake_response_status_line,
+                          &handshake_response_header);
+    // Remove cookie headers.
+    std::string filtered_handshake_response_header =
+        HttpUtil::StripHeaders(
+            handshake_response_header,
+            kSetCookieHeaders, arraysize(kSetCookieHeaders));
+    std::string remaining_data =
+        std::string(handshake_response_.data() +
+                    handshake_response_header_length_,
+                    handshake_response_.size() -
+                    handshake_response_header_length_);
+    std::string received_data =
+        handshake_response_status_line +
+        filtered_handshake_response_header +
+        "\r\n" +
+        remaining_data;
+    state_ = OPEN;
+    if (delegate_)
+      delegate_->OnReceivedData(socket_,
+                                received_data.data(), received_data.size());
+    return;
+  }
+
+  AddRef();  // Balanced in OnCanSetCookieCompleted
+
+  int policy = OK;
+  if (socket_->context()->cookie_policy()) {
+    GURL url_for_cookies = GetURLForCookies();
+    policy = socket_->context()->cookie_policy()->CanSetCookie(
+        url_for_cookies,
+        url_for_cookies,
+        response_cookies_[response_cookies_save_index_],
+        &can_set_cookie_callback_);
+    if (policy == ERR_IO_PENDING)
+      return;  // Wait for completion callback
+  }
+
+  OnCanSetCookieCompleted(policy);
+}
+
+void WebSocketJob::OnCanSetCookieCompleted(int policy) {
+  if (socket_ && delegate_ && state_ == CONNECTING) {
+    if ((policy == OK || policy == OK_FOR_SESSION_ONLY) &&
+        socket_->context()->cookie_store()) {
+      CookieOptions options;
+      options.set_include_httponly();
+      if (policy == OK_FOR_SESSION_ONLY)
+        options.set_force_session();
+      GURL url_for_cookies = GetURLForCookies();
+      socket_->context()->cookie_store()->SetCookieWithOptions(
+          url_for_cookies, response_cookies_[response_cookies_save_index_],
+          options);
+    }
+    response_cookies_save_index_++;
+    SaveNextCookie();
+  }
+  Release();  // Balance AddRef taken in SaveNextCookie
+}
+
+GURL WebSocketJob::GetURLForCookies() const {
+  GURL url = socket_->url();
+  std::string scheme = socket_->is_secure() ? "https" : "http";
+  url_canon::Replacements<char> replacements;
+  replacements.SetScheme(scheme.c_str(),
+                         url_parse::Component(0, scheme.length()));
+  return url.ReplaceComponents(replacements);
+}
+
+}  // namespace net