WebSocket over SPDY: protocol switch

This change enable to switch protocol from WebSocket to SPDY.
SpdyWebsocketStream class is not included, then SPDY connection
always fail now.

BUG=42320
TEST=net_unittests --gtest_filter=WebSocket\*

Review URL: http://codereview.chromium.org/7062043

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@89318 0039d316-1c4b-4281-b951-d872f2087c98

8 files changed, 202 insertions, 49 deletions

diff --git a/net/base/net_error_list.h b/net/base/net_error_list.h
index 9edc9ce..6cf561b 100644
--- a/net/base/net_error_list.h
+++ b/net/base/net_error_list.h
@@ -248,6 +248,10 @@ NET_ERROR(DNS_SERVER_REQUIRES_TCP, -144)
 //     operation for policy reasons.
 NET_ERROR(DNS_SERVER_FAILED, -145)
 
+// Connection was aborted for switching to another ptotocol.
+// WebSocket abort SocketStream connection when alternate protocol is found.
+NET_ERROR(PROTOCOL_SWITCHED, -146)
+
 // Certificate error codes
 //
 // The values of certificate error codes must be consecutive.
diff --git a/net/socket_stream/socket_stream.cc b/net/socket_stream/socket_stream.cc
index a36603f..5915ba4 100644
--- a/net/socket_stream/socket_stream.cc
+++ b/net/socket_stream/socket_stream.cc
@@ -257,9 +257,10 @@ void SocketStream::CopyAddrInfo(struct addrinfo* head) {
 
 void SocketStream::DoClose() {
   closing_ = true;
-  // If next_state_ is STATE_TCP_CONNECT, it's waiting other socket establishing
-  // connection.  If next_state_ is STATE_AUTH_REQUIRED, it's waiting for
-  // restarting.  In these states, we'll close the SocketStream now.
+  // If next_state_ is STATE_TCP_CONNECT, it's waiting other socket
+  // establishing connection.  If next_state_ is STATE_AUTH_REQUIRED, it's
+  // waiting for restarting.  In these states, we'll close the SocketStream
+  // now.
   if (next_state_ == STATE_TCP_CONNECT || next_state_ == STATE_AUTH_REQUIRED) {
     DoLoop(ERR_ABORTED);
     return;
@@ -470,6 +471,14 @@ int SocketStream::DoResolveProxy() {
     return ERR_INVALID_ARGUMENT;
   }
 
+  // TODO(toyoshim): Check server advertisement of SPDY through the HTTP
+  // Alternate-Protocol header, then switch to SPDY if SPDY is available.
+  // Usually we already have a session to the SPDY server because JavaScript
+  // running WebSocket itself would be served by SPDY. But, in some situation
+  // (E.g. Used by Chrome Extensions or used for cross origin connection), this
+  // connection might be the first one. At that time, we should check
+  // Alternate-Protocol header here for ws:// or TLS NPN extension for wss:// .
+
   return proxy_service()->ResolveProxy(
       proxy_url_, &proxy_info_, &io_callback_, &pac_request_, net_log_);
 }
@@ -541,6 +550,9 @@ int SocketStream::DoResolveHostComplete(int result) {
     result = delegate_->OnStartOpenConnection(this, &io_callback_);
     if (result == ERR_IO_PENDING)
       metrics_->OnWaitConnection();
+    else if (result == ERR_PROTOCOL_SWITCHED)
+      // TODO(toyoshim): Add metrics events for protocol switch.
+      next_state_ = STATE_CLOSE;
   } else {
     next_state_ = STATE_CLOSE;
   }
@@ -548,6 +560,10 @@ int SocketStream::DoResolveHostComplete(int result) {
   return result;
 }
 
+const ProxyServer& SocketStream::proxy_server() const {
+  return proxy_info_.proxy_server();
+}
+
 int SocketStream::DoTcpConnect(int result) {
   if (result != OK) {
     next_state_ = STATE_CLOSE;
@@ -853,6 +869,10 @@ int SocketStream::DoSSLConnectComplete(int result) {
     }
   }
 
+  // TODO(toyoshim): Upgrade to SPDY through TLS NPN extension if possible.
+  // If we use HTTPS and this is the first connection to the SPDY server,
+  // we should take care of TLS NPN extension here.
+
   if (result == OK)
     result = DidEstablishConnection();
   else
diff --git a/net/socket_stream/socket_stream.h b/net/socket_stream/socket_stream.h
index 0b1bc10..4fb9ea3 100644
--- a/net/socket_stream/socket_stream.h
+++ b/net/socket_stream/socket_stream.h
@@ -159,6 +159,8 @@ class NET_API SocketStream : public base::RefCountedThreadSafe<SocketStream> {
   // back.
   virtual void DetachDelegate();
 
+  const ProxyServer& proxy_server() const;
+
   // Sets an alternative HostResolver. For testing purposes only.
   void SetHostResolver(HostResolver* host_resolver);
 
diff --git a/net/socket_stream/socket_stream_unittest.cc b/net/socket_stream/socket_stream_unittest.cc
index c1e2d7d..5cd0207 100644
--- a/net/socket_stream/socket_stream_unittest.cc
+++ b/net/socket_stream/socket_stream_unittest.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 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.
 
@@ -20,33 +20,39 @@
 
 struct SocketStreamEvent {
   enum EventType {
-    EVENT_CONNECTED, EVENT_SENT_DATA, EVENT_RECEIVED_DATA, EVENT_CLOSE,
-    EVENT_AUTH_REQUIRED,
+    EVENT_START_OPEN_CONNECTION, EVENT_CONNECTED, EVENT_SENT_DATA,
+    EVENT_RECEIVED_DATA, EVENT_CLOSE, EVENT_AUTH_REQUIRED,
   };
 
-  SocketStreamEvent(EventType type, net::SocketStream* socket_stream,
-                    int num, const std::string& str,
-                    net::AuthChallengeInfo* auth_challenge_info)
+  SocketStreamEvent(EventType type,
+                    net::SocketStream* socket_stream,
+                    int num,
+                    const std::string& str,
+                    net::AuthChallengeInfo* auth_challenge_info,
+                    net::CompletionCallback* callback)
       : event_type(type), socket(socket_stream), number(num), data(str),
-        auth_info(auth_challenge_info) {}
+        auth_info(auth_challenge_info), result(net::OK) {}
 
   EventType event_type;
   net::SocketStream* socket;
   int number;
   std::string data;
   scoped_refptr<net::AuthChallengeInfo> auth_info;
+  int result;
 };
 
 class SocketStreamEventRecorder : public net::SocketStream::Delegate {
  public:
   explicit SocketStreamEventRecorder(net::CompletionCallback* callback)
-      : on_connected_(NULL),
+      : on_start_open_connection_(NULL),
+        on_connected_(NULL),
         on_sent_data_(NULL),
         on_received_data_(NULL),
         on_close_(NULL),
         on_auth_required_(NULL),
         callback_(callback) {}
   virtual ~SocketStreamEventRecorder() {
+    delete on_start_open_connection_;
     delete on_connected_;
     delete on_sent_data_;
     delete on_received_data_;
@@ -54,6 +60,10 @@ class SocketStreamEventRecorder : public net::SocketStream::Delegate {
     delete on_auth_required_;
   }
 
+  void SetOnStartOpenConnection(
+      Callback1<SocketStreamEvent*>::Type* callback) {
+    on_start_open_connection_ = callback;
+  }
   void SetOnConnected(Callback1<SocketStreamEvent*>::Type* callback) {
     on_connected_ = callback;
   }
@@ -70,12 +80,21 @@ class SocketStreamEventRecorder : public net::SocketStream::Delegate {
     on_auth_required_ = callback;
   }
 
+  virtual int OnStartOpenConnection(net::SocketStream* socket,
+                                    net::CompletionCallback* callback) {
+    events_.push_back(
+        SocketStreamEvent(SocketStreamEvent::EVENT_START_OPEN_CONNECTION,
+                          socket, 0, std::string(), NULL, callback));
+    if (on_start_open_connection_)
+      on_start_open_connection_->Run(&events_.back());
+    return events_.back().result;
+  }
   virtual void OnConnected(net::SocketStream* socket,
                            int num_pending_send_allowed) {
     events_.push_back(
         SocketStreamEvent(SocketStreamEvent::EVENT_CONNECTED,
                           socket, num_pending_send_allowed, std::string(),
-                          NULL));
+                          NULL, NULL));
     if (on_connected_)
       on_connected_->Run(&events_.back());
   }
@@ -83,7 +102,7 @@ class SocketStreamEventRecorder : public net::SocketStream::Delegate {
                           int amount_sent) {
     events_.push_back(
         SocketStreamEvent(SocketStreamEvent::EVENT_SENT_DATA,
-                          socket, amount_sent, std::string(), NULL));
+                          socket, amount_sent, std::string(), NULL, NULL));
     if (on_sent_data_)
       on_sent_data_->Run(&events_.back());
   }
@@ -91,14 +110,14 @@ class SocketStreamEventRecorder : public net::SocketStream::Delegate {
                               const char* data, int len) {
     events_.push_back(
         SocketStreamEvent(SocketStreamEvent::EVENT_RECEIVED_DATA,
-                          socket, len, std::string(data, len), NULL));
+                          socket, len, std::string(data, len), NULL, NULL));
     if (on_received_data_)
       on_received_data_->Run(&events_.back());
   }
   virtual void OnClose(net::SocketStream* socket) {
     events_.push_back(
         SocketStreamEvent(SocketStreamEvent::EVENT_CLOSE,
-                          socket, 0, std::string(), NULL));
+                          socket, 0, std::string(), NULL, NULL));
     if (on_close_)
       on_close_->Run(&events_.back());
     if (callback_)
@@ -108,7 +127,7 @@ class SocketStreamEventRecorder : public net::SocketStream::Delegate {
                               net::AuthChallengeInfo* auth_info) {
     events_.push_back(
         SocketStreamEvent(SocketStreamEvent::EVENT_AUTH_REQUIRED,
-                          socket, 0, std::string(), auth_info));
+                          socket, 0, std::string(), auth_info, NULL));
     if (on_auth_required_)
       on_auth_required_->Run(&events_.back());
   }
@@ -133,6 +152,7 @@ class SocketStreamEventRecorder : public net::SocketStream::Delegate {
 
  private:
   std::vector<SocketStreamEvent> events_;
+  Callback1<SocketStreamEvent*>::Type* on_start_open_connection_;
   Callback1<SocketStreamEvent*>::Type* on_connected_;
   Callback1<SocketStreamEvent*>::Type* on_sent_data_;
   Callback1<SocketStreamEvent*>::Type* on_received_data_;
@@ -193,6 +213,10 @@ class SocketStreamTest : public PlatformTest {
     event->socket->Close();
   }
 
+  virtual void DoSwitchToSpdyTest(SocketStreamEvent* event) {
+    event->result = net::ERR_PROTOCOL_SWITCHED;
+  }
+
   static const char* kWebSocketHandshakeRequest;
   static const char* kWebSocketHandshakeResponse;
 
@@ -275,14 +299,16 @@ TEST_F(SocketStreamTest, CloseFlushPendingWrite) {
   callback.WaitForResult();
 
   const std::vector<SocketStreamEvent>& events = delegate->GetSeenEvents();
-  EXPECT_EQ(6U, events.size());
+  ASSERT_EQ(7U, events.size());
 
-  EXPECT_EQ(SocketStreamEvent::EVENT_CONNECTED, events[0].event_type);
-  EXPECT_EQ(SocketStreamEvent::EVENT_SENT_DATA, events[1].event_type);
-  EXPECT_EQ(SocketStreamEvent::EVENT_RECEIVED_DATA, events[2].event_type);
-  EXPECT_EQ(SocketStreamEvent::EVENT_SENT_DATA, events[3].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_START_OPEN_CONNECTION,
+            events[0].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_CONNECTED, events[1].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_SENT_DATA, events[2].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_RECEIVED_DATA, events[3].event_type);
   EXPECT_EQ(SocketStreamEvent::EVENT_SENT_DATA, events[4].event_type);
-  EXPECT_EQ(SocketStreamEvent::EVENT_CLOSE, events[5].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_SENT_DATA, events[5].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_CLOSE, events[6].event_type);
 }
 
 TEST_F(SocketStreamTest, BasicAuthProxy) {
@@ -344,13 +370,46 @@ TEST_F(SocketStreamTest, BasicAuthProxy) {
   callback.WaitForResult();
 
   const std::vector<SocketStreamEvent>& events = delegate->GetSeenEvents();
-  EXPECT_EQ(3U, events.size());
+  ASSERT_EQ(4U, events.size());
 
-  EXPECT_EQ(SocketStreamEvent::EVENT_AUTH_REQUIRED, events[0].event_type);
-  EXPECT_EQ(SocketStreamEvent::EVENT_CONNECTED, events[1].event_type);
-  EXPECT_EQ(SocketStreamEvent::EVENT_CLOSE, events[2].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_START_OPEN_CONNECTION,
+            events[0].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_AUTH_REQUIRED, events[1].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_CONNECTED, events[2].event_type);
+  EXPECT_EQ(SocketStreamEvent::EVENT_CLOSE, events[3].event_type);
 
   // TODO(eroman): Add back NetLogTest here...
 }
 
+TEST_F(SocketStreamTest, SwitchToSpdy) {
+  TestCompletionCallback callback;
+
+  scoped_ptr<SocketStreamEventRecorder> delegate(
+      new SocketStreamEventRecorder(&callback));
+  SocketStreamTest* test = this;
+  delegate->SetOnStartOpenConnection(NewCallback(
+      test, &SocketStreamTest::DoSwitchToSpdyTest));
+
+  MockHostResolver host_resolver;
+
+  scoped_refptr<SocketStream> socket_stream(
+      new SocketStream(GURL("ws://example.com/demo"), delegate.get()));
+
+  socket_stream->set_context(new TestURLRequestContext());
+  socket_stream->SetHostResolver(&host_resolver);
+
+  socket_stream->Connect();
+
+  int result = callback.WaitForResult();
+
+  const std::vector<SocketStreamEvent>& events = delegate->GetSeenEvents();
+  ASSERT_EQ(2U, events.size());
+
+  EXPECT_EQ(SocketStreamEvent::EVENT_START_OPEN_CONNECTION,
+            events[0].event_type);
+  EXPECT_EQ(net::ERR_PROTOCOL_SWITCHED, events[0].result);
+  EXPECT_EQ(SocketStreamEvent::EVENT_CLOSE, events[1].event_type);
+  EXPECT_EQ(net::OK, result);
+}
+
 }  // namespace net
diff --git a/net/websockets/websocket_job.cc b/net/websockets/websocket_job.cc
index fc67f9a..77fc058 100644
--- a/net/websockets/websocket_job.cc
+++ b/net/websockets/websocket_job.cc
@@ -13,7 +13,11 @@
 #include "net/base/net_log.h"
 #include "net/base/cookie_store.h"
 #include "net/base/io_buffer.h"
+#include "net/http/http_network_session.h"
+#include "net/http/http_transaction_factory.h"
 #include "net/http/http_util.h"
+#include "net/spdy/spdy_session.h"
+#include "net/spdy/spdy_session_pool.h"
 #include "net/url_request/url_request_context.h"
 #include "net/websockets/websocket_frame_handler.h"
 #include "net/websockets/websocket_handshake_handler.h"
@@ -110,7 +114,8 @@ bool WebSocketJob::SendData(const char* data, int len) {
         // pending bytes less than max_pending_send_allowed, so when sending
         // larger message than max_pending_send_allowed should not be buffered.
         // If we don't call OnSentData, WebCore::SocketStreamHandle would stop
-        // sending more data when pending data reaches max_pending_send_allowed.
+        // sending more data when pending data reaches
+        // max_pending_send_allowed.
         // TODO(ukai): Fix this to support compression for larger message.
         int err = 0;
         if (!send_frame_handler_->GetCurrentBuffer() &&
@@ -119,7 +124,7 @@ bool WebSocketJob::SendData(const char* data, int len) {
           current_buffer_ = new DrainableIOBuffer(
               send_frame_handler_->GetCurrentBuffer(),
               send_frame_handler_->GetCurrentBufferSize());
-          return socket_->SendData(
+          return SendDataInternal(
               current_buffer_->data(), current_buffer_->BytesRemaining());
         }
         return err >= 0;
@@ -139,7 +144,7 @@ void WebSocketJob::Close() {
     return;
   }
   state_ = CLOSED;
-  socket_->Close();
+  CloseInternal();
 }
 
 void WebSocketJob::RestartWithAuth(
@@ -173,11 +178,14 @@ int WebSocketJob::OnStartOpenConnection(
   state_ = CONNECTING;
   addresses_ = socket->address_list();
   WebSocketThrottle::GetInstance()->PutInQueue(this);
-  if (!waiting_)
-    return OK;
+  if (!waiting_) {
+    int result = TrySpdyStream();
+    if (result != ERR_IO_PENDING)
+      return result;
+  }
   callback_ = callback;
   AddRef();  // Balanced when callback_ becomes NULL.
-  return ERR_IO_PENDING;
+  return ERR_IO_PENDING;  // Wakeup will be called later.
 }
 
 void WebSocketJob::OnConnected(
@@ -440,6 +448,33 @@ const AddressList& WebSocketJob::address_list() const {
   return addresses_;
 }
 
+int WebSocketJob::TrySpdyStream() {
+  if (!socket_.get())
+    return ERR_FAILED;
+
+  if (websocket_over_spdy_enabled_) {
+    // Check if we have a SPDY session available.
+    // If so, use it to create the websocket stream.
+    HttpTransactionFactory* factory =
+        socket_->context()->http_transaction_factory();
+    if (factory) {
+      scoped_refptr<HttpNetworkSession> session = factory->GetSession();
+      if (session.get()) {
+        SpdySessionPool* spdy_pool = session->spdy_session_pool();
+        const HostPortProxyPair pair(HostPortPair::FromURL(socket_->url()),
+                                     socket_->proxy_server());
+        if (spdy_pool->HasSession(pair)) {
+          // TODO(toyoshim): Switch to SpdyWebSocketStream here by returning
+          // ERR_PROTOCOL_SWITCHED.
+        }
+      }
+    }
+  }
+  // No SPDY session was available.
+  // Fallback to connecting a new socket.
+  return OK;
+}
+
 void WebSocketJob::SetWaiting() {
   waiting_ = true;
 }
@@ -455,20 +490,30 @@ void WebSocketJob::Wakeup() {
   DCHECK(callback_);
   MessageLoopForIO::current()->PostTask(
       FROM_HERE,
-      NewRunnableMethod(this,
-                        &WebSocketJob::DoCallback));
+      NewRunnableMethod(this, &WebSocketJob::RetryPendingIO));
 }
 
-void WebSocketJob::DoCallback() {
+void WebSocketJob::RetryPendingIO() {
+  int result = TrySpdyStream();
   // |callback_| may be NULL if OnClose() or DetachDelegate() was called.
   if (callback_) {
     net::CompletionCallback* callback = callback_;
     callback_ = NULL;
-    callback->Run(net::OK);
+    callback->Run(result);
     Release();  // Balanced with OnStartOpenConnection().
   }
 }
 
+bool WebSocketJob::SendDataInternal(const char* data, int length) {
+  // TODO(toyoshim): Call protocol specific SendData().
+  return socket_->SendData(data, length);
+}
+
+void WebSocketJob::CloseInternal() {
+  // TODO(toyoshim): Call protocol specific Close().
+  socket_->Close();
+}
+
 void WebSocketJob::SendPending() {
   if (current_buffer_)
     return;
@@ -477,13 +522,13 @@ void WebSocketJob::SendPending() {
   if (send_frame_handler_->UpdateCurrentBuffer(false) <= 0) {
     // No more data to send.
     if (state_ == CLOSING)
-      socket_->Close();
+      CloseInternal();
     return;
   }
   current_buffer_ = new DrainableIOBuffer(
       send_frame_handler_->GetCurrentBuffer(),
       send_frame_handler_->GetCurrentBufferSize());
-  socket_->SendData(current_buffer_->data(), current_buffer_->BytesRemaining());
+  SendDataInternal(current_buffer_->data(), current_buffer_->BytesRemaining());
 }
 
 }  // namespace net
diff --git a/net/websockets/websocket_job.h b/net/websockets/websocket_job.h
index b554469..369ef56 100644
--- a/net/websockets/websocket_job.h
+++ b/net/websockets/websocket_job.h
@@ -62,17 +62,13 @@ class NET_API WebSocketJob
   // SocketStream::Delegate methods.
   virtual int OnStartOpenConnection(
       SocketStream* socket, CompletionCallback* callback);
-  virtual void OnConnected(
-      SocketStream* socket, int max_pending_send_allowed);
-  virtual void OnSentData(
-      SocketStream* socket, int amount_sent);
-  virtual void OnReceivedData(
-      SocketStream* socket, const char* data, int len);
+  virtual void OnConnected(SocketStream* socket, int max_pending_send_allowed);
+  virtual void OnSentData(SocketStream* socket, int amount_sent);
+  virtual void OnReceivedData(SocketStream* socket, const char* data, int len);
   virtual void OnClose(SocketStream* socket);
   virtual void OnAuthRequired(
       SocketStream* socket, AuthChallengeInfo* auth_info);
-  virtual void OnError(
-      const SocketStream* socket, int error);
+  virtual void OnError(const SocketStream* socket, int error);
 
  private:
   friend class WebSocketThrottle;
@@ -91,11 +87,14 @@ class NET_API WebSocketJob
   GURL GetURLForCookies() const;
 
   const AddressList& address_list() const;
+  int TrySpdyStream();
   void SetWaiting();
   bool IsWaiting() const;
   void Wakeup();
-  void DoCallback();
+  void RetryPendingIO();
 
+  bool SendDataInternal(const char* data, int length);
+  void CloseInternal();
   void SendPending();
 
   static bool websocket_over_spdy_enabled_;
diff --git a/net/websockets/websocket_job_unittest.cc b/net/websockets/websocket_job_unittest.cc
index 8951a4e..b19fb37 100644
--- a/net/websockets/websocket_job_unittest.cc
+++ b/net/websockets/websocket_job_unittest.cc
@@ -59,7 +59,8 @@ class MockSocketStreamDelegate : public SocketStream::Delegate {
   }
   virtual ~MockSocketStreamDelegate() {}
 
-  virtual void OnConnected(SocketStream* socket, int max_pending_send_allowed) {
+  virtual void OnConnected(SocketStream* socket,
+                           int max_pending_send_allowed) {
   }
   virtual void OnSentData(SocketStream* socket, int amount_sent) {
     amount_sent_ += amount_sent;
@@ -278,6 +279,9 @@ const char* WebSocketJobTest::kHandshakeResponseWithCookie =
     "8jKS'y:G*Co,Wxa-";
 
 TEST_F(WebSocketJobTest, SimpleHandshake) {
+  // TODO(toyoshim): We need to consider both spdy-enabled and spdy-disabled
+  // configuration.
+  WebSocketJob::set_websocket_over_spdy_enabled(true);
   GURL url("ws://example.com/demo");
   MockSocketStreamDelegate delegate;
   InitWebSocketJob(url, &delegate);
@@ -302,6 +306,7 @@ TEST_F(WebSocketJobTest, SimpleHandshake) {
 }
 
 TEST_F(WebSocketJobTest, SlowHandshake) {
+  WebSocketJob::set_websocket_over_spdy_enabled(true);
   GURL url("ws://example.com/demo");
   MockSocketStreamDelegate delegate;
   InitWebSocketJob(url, &delegate);
@@ -341,6 +346,7 @@ TEST_F(WebSocketJobTest, SlowHandshake) {
 }
 
 TEST_F(WebSocketJobTest, HandshakeWithCookie) {
+  WebSocketJob::set_websocket_over_spdy_enabled(true);
   GURL url("ws://example.com/demo");
   GURL cookieUrl("http://example.com/demo");
   CookieOptions cookie_options;
@@ -382,6 +388,7 @@ TEST_F(WebSocketJobTest, HandshakeWithCookie) {
 }
 
 TEST_F(WebSocketJobTest, HandshakeWithCookieButNotAllowed) {
+  WebSocketJob::set_websocket_over_spdy_enabled(true);
   GURL url("ws://example.com/demo");
   GURL cookieUrl("http://example.com/demo");
   CookieOptions cookie_options;
@@ -422,6 +429,7 @@ TEST_F(WebSocketJobTest, HandshakeWithCookieButNotAllowed) {
 }
 
 TEST_F(WebSocketJobTest, HSTSUpgrade) {
+  WebSocketJob::set_websocket_over_spdy_enabled(true);
   GURL url("ws://upgrademe.com/");
   MockSocketStreamDelegate delegate;
   scoped_refptr<SocketStreamJob> job = SocketStreamJob::CreateSocketStreamJob(
diff --git a/net/websockets/websocket_throttle_unittest.cc b/net/websockets/websocket_throttle_unittest.cc
index 7eb18fa5..e89b2c1 100644
--- a/net/websockets/websocket_throttle_unittest.cc
+++ b/net/websockets/websocket_throttle_unittest.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 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.
 
@@ -10,6 +10,7 @@
 #include "net/base/sys_addrinfo.h"
 #include "net/base/test_completion_callback.h"
 #include "net/socket_stream/socket_stream.h"
+#include "net/url_request/url_request_test_util.h"
 #include "net/websockets/websocket_job.h"
 #include "net/websockets/websocket_throttle.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -64,6 +65,8 @@ class WebSocketThrottleTest : public PlatformTest {
   static void MockSocketStreamConnect(
       SocketStream* socket, struct addrinfo* head) {
     socket->CopyAddrInfo(head);
+    // TODO(toyoshim): We should introduce additional tests on cases via proxy.
+    socket->proxy_info_.UseDirect();
     // In SocketStream::Connect(), it adds reference to socket, which is
     // balanced with SocketStream::Finish() that is finally called from
     // SocketStream::Close() or SocketStream::DetachDelegate(), when
@@ -77,7 +80,11 @@ class WebSocketThrottleTest : public PlatformTest {
 };
 
 TEST_F(WebSocketThrottleTest, Throttle) {
+  scoped_refptr<URLRequestContext> context(new TestURLRequestContext);
   DummySocketStreamDelegate delegate;
+  // TODO(toyoshim): We need to consider both spdy-enabled and spdy-disabled
+  // configuration.
+  WebSocketJob::set_websocket_over_spdy_enabled(true);
 
   // For host1: 1.2.3.4, 1.2.3.5, 1.2.3.6
   struct addrinfo* addr = AddAddr(1, 2, 3, 4, NULL);
@@ -86,6 +93,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<WebSocketJob> w1(new WebSocketJob(&delegate));
   scoped_refptr<SocketStream> s1(
       new SocketStream(GURL("ws://host1/"), w1.get()));
+  s1->set_context(context.get());
   w1->InitSocketStream(s1.get());
   WebSocketThrottleTest::MockSocketStreamConnect(s1, addr);
   DeleteAddrInfo(addr);
@@ -106,6 +114,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<WebSocketJob> w2(new WebSocketJob(&delegate));
   scoped_refptr<SocketStream> s2(
       new SocketStream(GURL("ws://host2/"), w2.get()));
+  s2->set_context(context.get());
   w2->InitSocketStream(s2.get());
   WebSocketThrottleTest::MockSocketStreamConnect(s2, addr);
   DeleteAddrInfo(addr);
@@ -125,6 +134,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<WebSocketJob> w3(new WebSocketJob(&delegate));
   scoped_refptr<SocketStream> s3(
       new SocketStream(GURL("ws://host3/"), w3.get()));
+  s3->set_context(context.get());
   w3->InitSocketStream(s3.get());
   WebSocketThrottleTest::MockSocketStreamConnect(s3, addr);
   DeleteAddrInfo(addr);
@@ -144,6 +154,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<WebSocketJob> w4(new WebSocketJob(&delegate));
   scoped_refptr<SocketStream> s4(
       new SocketStream(GURL("ws://host4/"), w4.get()));
+  s4->set_context(context.get());
   w4->InitSocketStream(s4.get());
   WebSocketThrottleTest::MockSocketStreamConnect(s4, addr);
   DeleteAddrInfo(addr);
@@ -162,6 +173,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<WebSocketJob> w5(new WebSocketJob(&delegate));
   scoped_refptr<SocketStream> s5(
       new SocketStream(GURL("ws://host5/"), w5.get()));
+  s5->set_context(context.get());
   w5->InitSocketStream(s5.get());
   WebSocketThrottleTest::MockSocketStreamConnect(s5, addr);
   DeleteAddrInfo(addr);
@@ -180,6 +192,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<WebSocketJob> w6(new WebSocketJob(&delegate));
   scoped_refptr<SocketStream> s6(
       new SocketStream(GURL("ws://host6/"), w6.get()));
+  s6->set_context(context.get());
   w6->InitSocketStream(s6.get());
   WebSocketThrottleTest::MockSocketStreamConnect(s6, addr);
   DeleteAddrInfo(addr);
@@ -281,7 +294,9 @@ TEST_F(WebSocketThrottleTest, Throttle) {
 }
 
 TEST_F(WebSocketThrottleTest, NoThrottleForDuplicateAddress) {
+  scoped_refptr<URLRequestContext> context(new TestURLRequestContext);
   DummySocketStreamDelegate delegate;
+  WebSocketJob::set_websocket_over_spdy_enabled(true);
 
   // For localhost: 127.0.0.1, 127.0.0.1
   struct addrinfo* addr = AddAddr(127, 0, 0, 1, NULL);
@@ -289,6 +304,7 @@ TEST_F(WebSocketThrottleTest, NoThrottleForDuplicateAddress) {
   scoped_refptr<WebSocketJob> w1(new WebSocketJob(&delegate));
   scoped_refptr<SocketStream> s1(
       new SocketStream(GURL("ws://localhost/"), w1.get()));
+  s1->set_context(context.get());
   w1->InitSocketStream(s1.get());
   WebSocketThrottleTest::MockSocketStreamConnect(s1, addr);
   DeleteAddrInfo(addr);