Implement handshake timeout on WebSocket

net/websockets/websocket_stream_test.cc

 // Copyright 2013 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_stream.h"
 
 #include <algorithm>
 #include <string>
 #include <utility>
 #include <vector>
 
 #include "base/compiler_specific.h"
 #include "base/memory/scoped_vector.h"
 #include "base/metrics/histogram.h"
 #include "base/metrics/histogram_samples.h"
 #include "base/metrics/statistics_recorder.h"
 #include "base/run_loop.h"
 #include "base/strings/stringprintf.h"
+#include "base/timer/mock_timer.h"
+#include "base/timer/timer.h"
 #include "net/base/net_errors.h"
 #include "net/base/test_data_directory.h"
 #include "net/http/http_request_headers.h"
 #include "net/http/http_response_headers.h"
 #include "net/socket/client_socket_handle.h"
 #include "net/socket/socket_test_util.h"
 #include "net/test/cert_test_util.h"
 #include "net/url_request/url_request_test_util.h"
 #include "net/websockets/websocket_basic_handshake_stream.h"
 #include "net/websockets/websocket_frame.h"
@@ skipping 43 matching lines @@
   socket_data->SetStop(reads_count + writes_count);
   return socket_data.Pass();
 }
 
 // Builder for a DeterministicSocketData that expects nothing. This does not
 // set the connect data, so the calling code must do that explicitly.
 scoped_ptr<DeterministicSocketData> BuildNullSocketData() {
   return make_scoped_ptr(new DeterministicSocketData(NULL, 0, NULL, 0));
 }
 
+class MockWeakTimer : public base::MockTimer,
+                      public base::SupportsWeakPtr<MockWeakTimer> {
+ public:
+  MockWeakTimer(bool retain_user_task, bool is_repeating)
+      : MockTimer(retain_user_task, is_repeating) {}
+};
+
 // A sub-class of WebSocketHandshakeStreamCreateHelper which always sets a
 // deterministic key to use in the WebSocket handshake.
 class DeterministicKeyWebSocketHandshakeStreamCreateHelper
     : public WebSocketHandshakeStreamCreateHelper {
  public:
   DeterministicKeyWebSocketHandshakeStreamCreateHelper(
       WebSocketStream::ConnectDelegate* connect_delegate,
       const std::vector<std::string>& requested_subprotocols)
       : WebSocketHandshakeStreamCreateHelper(connect_delegate,
                                              requested_subprotocols) {}
 
   virtual void OnStreamCreated(WebSocketBasicHandshakeStream* stream) OVERRIDE {
     stream->SetWebSocketKeyForTesting("dGhlIHNhbXBsZSBub25jZQ==");
   }
 };
 
 class WebSocketStreamCreateTest : public ::testing::Test {
  public:
   WebSocketStreamCreateTest() : has_failed_(false), ssl_fatal_(false) {}
 
   void CreateAndConnectCustomResponse(
       const std::string& socket_url,
       const std::string& socket_path,
       const std::vector<std::string>& sub_protocols,
       const std::string& origin,
       const std::string& extra_request_headers,
-      const std::string& response_body) {
+      const std::string& response_body,
+      scoped_ptr<base::Timer> timer = scoped_ptr<base::Timer>()) {
     url_request_context_host_.SetExpectations(
         WebSocketStandardRequest(socket_path, origin, extra_request_headers),
         response_body);
-    CreateAndConnectStream(socket_url, sub_protocols, origin);
+    CreateAndConnectStream(socket_url, sub_protocols, origin, timer.Pass());
   }
 
   // |extra_request_headers| and |extra_response_headers| must end in "\r\n" or
   // errors like "Unable to perform synchronous IO while stopped" will occur.
   void CreateAndConnectStandard(const std::string& socket_url,
                                 const std::string& socket_path,
                                 const std::vector<std::string>& sub_protocols,
                                 const std::string& origin,
                                 const std::string& extra_request_headers,
-                                const std::string& extra_response_headers) {
+                                const std::string& extra_response_headers,
+                                scoped_ptr<base::Timer> timer =
+                                scoped_ptr<base::Timer>()) {
     CreateAndConnectCustomResponse(
         socket_url,
         socket_path,
         sub_protocols,
         origin,
         extra_request_headers,
-        WebSocketStandardResponse(extra_response_headers));
+        WebSocketStandardResponse(extra_response_headers),
+        timer.Pass());
   }
 
   void CreateAndConnectRawExpectations(
       const std::string& socket_url,
       const std::vector<std::string>& sub_protocols,
       const std::string& origin,
-      scoped_ptr<DeterministicSocketData> socket_data) {
+      scoped_ptr<DeterministicSocketData> socket_data,
+      scoped_ptr<base::Timer> timer = scoped_ptr<base::Timer>()) {
     AddRawExpectations(socket_data.Pass());
-    CreateAndConnectStream(socket_url, sub_protocols, origin);
+    CreateAndConnectStream(socket_url, sub_protocols, origin, timer.Pass());
   }
 
   // Add additional raw expectations for sockets created before the final one.
   void AddRawExpectations(scoped_ptr<DeterministicSocketData> socket_data) {
     url_request_context_host_.AddRawExpectations(socket_data.Pass());
   }
 
   // A wrapper for CreateAndConnectStreamForTesting that knows about our default
   // parameters.
   void CreateAndConnectStream(const std::string& socket_url,
                               const std::vector<std::string>& sub_protocols,
-                              const std::string& origin) {
+                              const std::string& origin,
+                              scoped_ptr<base::Timer> timer) {
     for (size_t i = 0; i < ssl_data_.size(); ++i) {
       scoped_ptr<SSLSocketDataProvider> ssl_data(ssl_data_[i]);
       ssl_data_[i] = NULL;
       url_request_context_host_.AddSSLSocketDataProvider(ssl_data.Pass());
     }
     ssl_data_.clear();
     scoped_ptr<WebSocketStream::ConnectDelegate> connect_delegate(
         new TestConnectDelegate(this));
     WebSocketStream::ConnectDelegate* delegate = connect_delegate.get();
     scoped_ptr<WebSocketHandshakeStreamCreateHelper> create_helper(
         new DeterministicKeyWebSocketHandshakeStreamCreateHelper(
             delegate, sub_protocols));
     stream_request_ = ::net::CreateAndConnectStreamForTesting(
         GURL(socket_url),
         create_helper.Pass(),
         url::Origin(origin),
         url_request_context_host_.GetURLRequestContext(),
         BoundNetLog(),
-        connect_delegate.Pass());
+        connect_delegate.Pass(),
+        timer ? timer.Pass() : scoped_ptr<base::Timer>(
+            new base::Timer(false, false)));
   }
 
   static void RunUntilIdle() { base::RunLoop().RunUntilIdle(); }
 
   // A simple function to make the tests more readable. Creates an empty vector.
   static std::vector<std::string> NoSubProtocols() {
     return std::vector<std::string>();
   }
 
   const std::string& failure_message() const { return failure_message_; }
@@ skipping 911 matching lines @@
   socket_data->set_connect_data(
       MockConnect(ASYNC, ERR_CONNECTION_TIMED_OUT));
   CreateAndConnectRawExpectations("ws://localhost/", NoSubProtocols(),
                                   "http://localhost", socket_data.Pass());
   RunUntilIdle();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error in connection establishment: net::ERR_CONNECTION_TIMED_OUT",
             failure_message());
 }
 
+// The server doesn't respond to the opening handshake.
+TEST_F(WebSocketStreamCreateTest, HandshakeTimeout) {
+  scoped_ptr<DeterministicSocketData> socket_data(BuildNullSocketData());
+  socket_data->set_connect_data(MockConnect(SYNCHRONOUS, ERR_IO_PENDING));
+  scoped_ptr<MockWeakTimer> timer(new MockWeakTimer(false, false));
+  base::WeakPtr<MockWeakTimer> weak_timer = timer->AsWeakPtr();
+  CreateAndConnectRawExpectations("ws://localhost/", NoSubProtocols(),
+                                  "http://localhost", socket_data.Pass(),
+                                  timer.PassAs<base::Timer>());
+  EXPECT_FALSE(has_failed());
+  ASSERT_TRUE(weak_timer.get());
+  EXPECT_TRUE(weak_timer->IsRunning());
+
+  weak_timer->Fire();
+  RunUntilIdle();
+
+  EXPECT_TRUE(has_failed());
+  EXPECT_EQ("WebSocket opening handshake timed out", failure_message());
+  ASSERT_TRUE(weak_timer.get());
+  EXPECT_FALSE(weak_timer->IsRunning());
+}
+
+// When the connection establishes the timer should be stopped.
+TEST_F(WebSocketStreamCreateTest, HandshakeTimerOnSuccess) {
+  scoped_ptr<MockWeakTimer> timer(new MockWeakTimer(false, false));
+  base::WeakPtr<MockWeakTimer> weak_timer = timer->AsWeakPtr();
+
+  CreateAndConnectStandard(
+      "ws://localhost/", "/", NoSubProtocols(), "http://localhost", "", "",
+      timer.PassAs<base::Timer>());
+  EXPECT_FALSE(request_info_);

[Adam Rice, 2014/09/11 09:47:39]

I think you should remove these two assertions as they are not relevant to the
test (and they are covered by other test cases).

[yhirano, 2014/09/11 10:46:49]

Done.

+  EXPECT_FALSE(response_info_);
+  ASSERT_TRUE(weak_timer);
+  EXPECT_TRUE(weak_timer->IsRunning());
+
+  RunUntilIdle();
+  EXPECT_FALSE(has_failed());
+  EXPECT_TRUE(stream_);

[Adam Rice, 2014/09/11 09:47:39]

Also please remove these three assertions as above.

[yhirano, 2014/09/11 10:46:49]

I want to make sure the operation succeeds.
So I'd leave one assertion.

+  EXPECT_TRUE(request_info_);
+  EXPECT_TRUE(response_info_);
+  ASSERT_TRUE(weak_timer);
+  EXPECT_FALSE(weak_timer->IsRunning());
+}
+
+// When the connection fails the timer should be stopped.
+TEST_F(WebSocketStreamCreateTest, HandshakeTimerOnFailure) {
+  scoped_ptr<DeterministicSocketData> socket_data(BuildNullSocketData());
+  socket_data->set_connect_data(
+      MockConnect(SYNCHRONOUS, ERR_CONNECTION_REFUSED));
+  scoped_ptr<MockWeakTimer> timer(new MockWeakTimer(false, false));
+  base::WeakPtr<MockWeakTimer> weak_timer = timer->AsWeakPtr();
+  CreateAndConnectRawExpectations("ws://localhost/", NoSubProtocols(),
+                                  "http://localhost", socket_data.Pass(),
+                                  timer.PassAs<base::Timer>());
+  ASSERT_TRUE(weak_timer.get());
+  EXPECT_TRUE(weak_timer->IsRunning());
+
+  RunUntilIdle();
+  EXPECT_TRUE(has_failed());
+  EXPECT_EQ("Error in connection establishment: net::ERR_CONNECTION_REFUSED",
+            failure_message());
+  ASSERT_TRUE(weak_timer.get());
+  EXPECT_FALSE(weak_timer->IsRunning());
+}
+
 // Cancellation during connect works.
 TEST_F(WebSocketStreamCreateTest, CancellationDuringConnect) {
   scoped_ptr<DeterministicSocketData> socket_data(BuildNullSocketData());
   socket_data->set_connect_data(MockConnect(SYNCHRONOUS, ERR_IO_PENDING));
   CreateAndConnectRawExpectations("ws://localhost/",
                                   NoSubProtocols(),
                                   "http://localhost",
                                   socket_data.Pass());
   stream_request_.reset();
   RunUntilIdle();
@@ skipping 256 matching lines @@
   if (original) {
     samples->Subtract(*original);  // Cancel the original values.
   }
   EXPECT_EQ(1, samples->GetCount(INCOMPLETE));
   EXPECT_EQ(0, samples->GetCount(CONNECTED));
   EXPECT_EQ(0, samples->GetCount(FAILED));
 }
 
 }  // namespace
 }  // namespace net