Add closing handshake timeout.

If either the Close message or socket close fails to arrive within 4
minutes, we should close the connection outselves.

Also add tests for the closing handshake timeout (using a 1ms timeout).

BUG=230756
TEST=net_unittests --gtest_filter=WebSocketChannel*

Review URL: https://codereview.chromium.org/31813007

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

3 files changed, 255 insertions, 5 deletions

diff --git a/net/websockets/websocket_channel.cc b/net/websockets/websocket_channel.cc
index c370bd7..5e8d8d6 100644
--- a/net/websockets/websocket_channel.cc
+++ b/net/websockets/websocket_channel.cc
@@ -10,6 +10,7 @@
 #include "base/bind.h"
 #include "base/safe_numerics.h"
 #include "base/strings/string_util.h"
+#include "base/time/time.h"
 #include "net/base/big_endian.h"
 #include "net/base/io_buffer.h"
 #include "net/base/net_log.h"
@@ -26,6 +27,9 @@ namespace {
 const int kDefaultSendQuotaLowWaterMark = 1 << 16;
 const int kDefaultSendQuotaHighWaterMark = 1 << 17;
 const size_t kWebSocketCloseCodeLength = 2;
+// This timeout is based on TCPMaximumSegmentLifetime * 2 from
+// MainThreadWebSocketChannel.cpp in Blink.
+const int kClosingHandshakeTimeoutSeconds = 2 * 2 * 60;
 
 }  // namespace
 
@@ -89,6 +93,7 @@ WebSocketChannel::WebSocketChannel(
       send_quota_low_water_mark_(kDefaultSendQuotaLowWaterMark),
       send_quota_high_water_mark_(kDefaultSendQuotaHighWaterMark),
       current_send_quota_(0),
+      timeout_(base::TimeDelta::FromSeconds(kClosingHandshakeTimeoutSeconds)),
       closing_code_(0),
       state_(FRESHLY_CONSTRUCTED) {}
 
@@ -96,6 +101,9 @@ WebSocketChannel::~WebSocketChannel() {
   // The stream may hold a pointer to read_frames_, and so it needs to be
   // destroyed first.
   stream_.reset();
+  // The timer may have a callback pointing back to us, so stop it just in case
+  // someone decides to run the event loop from their destructor.
+  timer_.Stop();
 }
 
 void WebSocketChannel::SendAddChannelRequest(
@@ -181,7 +189,6 @@ void WebSocketChannel::StartClosingHandshake(uint16 code,
     return;
   }
   // TODO(ricea): Validate |code|? Check that |reason| is valid UTF-8?
-  // TODO(ricea): There should be a timeout for the closing handshake.
   SendClose(code, reason);  // Sets state_ to SEND_CLOSED
 }
 
@@ -190,10 +197,13 @@ void WebSocketChannel::SendAddChannelRequestForTesting(
     const std::vector<std::string>& requested_subprotocols,
     const GURL& origin,
     const WebSocketStreamFactory& factory) {
-  SendAddChannelRequestWithFactory(socket_url,
-                                   requested_subprotocols,
-                                   origin,
-                                   factory);
+  SendAddChannelRequestWithFactory(
+      socket_url, requested_subprotocols, origin, factory);
+}
+
+void WebSocketChannel::SetClosingHandshakeTimeoutForTesting(
+    base::TimeDelta delay) {
+  timeout_ = delay;
 }
 
 void WebSocketChannel::SendAddChannelRequestWithFactory(
@@ -482,6 +492,10 @@ void WebSocketChannel::HandleFrame(const WebSocketFrameHeader::OpCode opcode,
           break;
 
         case SEND_CLOSED:
+          // Give them another few minutes to actually close the connection.
+          DCHECK(timer_.IsRunning());
+          timer_.Reset();
+
           state_ = CLOSE_WAIT;
           // From RFC6455 section 7.1.5: "Each endpoint
           // will see the status code sent by the other end as _The WebSocket
@@ -577,6 +591,12 @@ void WebSocketChannel::SendClose(uint16 code, const std::string& reason) {
     std::copy(
         reason.begin(), reason.end(), body->data() + kWebSocketCloseCodeLength);
   }
+  // This use of base::Unretained() is safe because we stop the timer in the
+  // destructor.
+  timer_.Start(
+      FROM_HERE,
+      timeout_,
+      base::Bind(&WebSocketChannel::CloseTimeout, base::Unretained(this)));
   SendIOBuffer(true, WebSocketFrameHeader::kOpCodeClose, body, size);
   state_ = (state_ == CONNECTED) ? SEND_CLOSED : CLOSE_WAIT;
 }
@@ -618,4 +638,15 @@ void WebSocketChannel::ParseClose(const scoped_refptr<IOBuffer>& buffer,
   }
 }
 
+void WebSocketChannel::CloseTimeout() {
+  stream_->Close();
+  // TODO(ricea): This becomes a DCHECK() once
+  // https://codereview.chromium.org/26544003/ has landed.
+  if (state_ != CLOSED) {
+    state_ = CLOSED;
+    event_interface_->OnDropChannel(kWebSocketErrorAbnormalClosure,
+                                    "Abnormal Closure");
+  }
+}
+
 }  // namespace net
diff --git a/net/websockets/websocket_channel.h b/net/websockets/websocket_channel.h
index d391f0e..5016fb2 100644
--- a/net/websockets/websocket_channel.h
+++ b/net/websockets/websocket_channel.h
@@ -12,6 +12,8 @@
 #include "base/callback.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/memory/scoped_vector.h"
+#include "base/time/time.h"
+#include "base/timer/timer.h"
 #include "net/base/net_export.h"
 #include "net/websockets/websocket_frame.h"
 #include "net/websockets/websocket_stream.h"
@@ -93,6 +95,11 @@ class NET_EXPORT WebSocketChannel {
       const GURL& origin,
       const WebSocketStreamFactory& factory);
 
+  // The default timout for the closing handshake is a sensible value (see
+  // kClosingHandshakeTimeoutSeconds in websocket_channel.cc). However, we can
+  // set it to a very small value for testing purposes.
+  void SetClosingHandshakeTimeoutForTesting(base::TimeDelta delay);
+
  private:
   // The object passes through a linear progression of states from
   // FRESHLY_CONSTRUCTED to CLOSED, except that the SEND_CLOSED and RECV_CLOSED
@@ -209,6 +216,10 @@ class NET_EXPORT WebSocketChannel {
                   uint16* code,
                   std::string* reason);
 
+  // Called if the closing handshake times out. Closes the connection and
+  // informs the |event_interface_| if appropriate.
+  void CloseTimeout();
+
   // The URL of the remote server.
   GURL socket_url_;
 
@@ -248,6 +259,12 @@ class NET_EXPORT WebSocketChannel {
   // on this logical channel (quota units).
   int current_send_quota_;
 
+  // Timer for the closing handshake.
+  base::OneShotTimer<WebSocketChannel> timer_;
+
+  // Timeout for the closing handshake.
+  base::TimeDelta timeout_;
+
   // Storage for the status code and reason from the time the Close frame
   // arrives until the connection is closed and they are passed to
   // OnDropChannel().
diff --git a/net/websockets/websocket_channel_test.cc b/net/websockets/websocket_channel_test.cc
index e584860..6bc2a54 100644
--- a/net/websockets/websocket_channel_test.cc
+++ b/net/websockets/websocket_channel_test.cc
@@ -20,6 +20,7 @@
 #include "base/safe_numerics.h"
 #include "base/strings/string_piece.h"
 #include "net/base/net_errors.h"
+#include "net/base/test_completion_callback.h"
 #include "net/url_request/url_request_context.h"
 #include "net/websockets/websocket_errors.h"
 #include "net/websockets/websocket_event_interface.h"
@@ -84,6 +85,8 @@ std::ostream& operator<<(std::ostream& os,
 
 namespace {
 
+using ::base::TimeDelta;
+
 using ::testing::AnyNumber;
 using ::testing::InSequence;
 using ::testing::MockFunction;
@@ -116,6 +119,10 @@ const size_t kDefaultInitialQuota = 1 << 17;
 // kDefaultSendQuotaLowWaterMark change.
 const size_t kDefaultQuotaRefreshTrigger = (1 << 16) + 1;
 
+// TestTimeouts::tiny_timeout() is 100ms! I could run halfway around the world
+// in that time! I would like my tests to run a bit quicker.
+const int kVeryTinyTimeoutMillis = 1;
+
 // This mock is for testing expectations about how the EventInterface is used.
 class MockWebSocketEventInterface : public WebSocketEventInterface {
  public:
@@ -353,6 +360,21 @@ template <size_t N>
   return ::testing::MakeMatcher(new EqualsFramesMatcher<N>(&frames));
 }
 
+// TestClosure works like TestCompletionCallback, but doesn't take an argument.
+class TestClosure {
+ public:
+  base::Closure closure() { return base::Bind(callback_.callback(), OK); }
+
+  void WaitForResult() { callback_.WaitForResult(); }
+
+ private:
+  // Delegate to TestCompletionCallback for the implementation.
+  TestCompletionCallback callback_;
+};
+
+// A GoogleMock action to run a Closure.
+ACTION_P(InvokeClosure, closure) { closure.Run(); }
+
 // A FakeWebSocketStream whose ReadFrames() function returns data.
 class ReadableFakeWebSocketStream : public FakeWebSocketStream {
  public:
@@ -1366,6 +1388,69 @@ TEST_F(WebSocketChannelEventInterfaceTest, AsyncProtocolErrorGivesStatus1002) {
   base::MessageLoop::current()->RunUntilIdle();
 }
 
+// The closing handshake times out and sends an OnDropChannel event if no
+// response to the client Close message is received.
+TEST_F(WebSocketChannelEventInterfaceTest,
+       ClientInitiatedClosingHandshakeTimesOut) {
+  scoped_ptr<ReadableFakeWebSocketStream> stream(
+      new ReadableFakeWebSocketStream);
+  stream->PrepareReadFramesError(ReadableFakeWebSocketStream::SYNC,
+                                 ERR_IO_PENDING);
+  set_stream(stream.Pass());
+  EXPECT_CALL(*event_interface_, OnAddChannelResponse(false, _));
+  EXPECT_CALL(*event_interface_, OnFlowControl(_));
+  // This checkpoint object verifies that the OnDropChannel message comes after
+  // the timeout.
+  MockFunction<void(int)> checkpoint;
+  TestClosure completion;
+  {
+    InSequence s;
+    EXPECT_CALL(checkpoint, Call(1));
+    EXPECT_CALL(*event_interface_,
+                OnDropChannel(kWebSocketErrorAbnormalClosure, _))
+        .WillOnce(InvokeClosure(completion.closure()));
+  }
+  CreateChannelAndConnectSuccessfully();
+  // OneShotTimer is not very friendly to testing; there is no apparent way to
+  // set an expectation on it. Instead the tests need to infer that the timeout
+  // was fired by the behaviour of the WebSocketChannel object.
+  channel_->SetClosingHandshakeTimeoutForTesting(
+      TimeDelta::FromMilliseconds(kVeryTinyTimeoutMillis));
+  channel_->StartClosingHandshake(kWebSocketNormalClosure, "");
+  checkpoint.Call(1);
+  completion.WaitForResult();
+}
+
+// The closing handshake times out and sends an OnDropChannel event if a Close
+// message is received but the connection isn't closed by the remote host.
+TEST_F(WebSocketChannelEventInterfaceTest,
+       ServerInitiatedClosingHandshakeTimesOut) {
+  scoped_ptr<ReadableFakeWebSocketStream> stream(
+      new ReadableFakeWebSocketStream);
+  static const InitFrame frames[] = {
+      {FINAL_FRAME, WebSocketFrameHeader::kOpCodeClose,
+       NOT_MASKED,  CLOSE_DATA(NORMAL_CLOSURE, "OK")}};
+  stream->PrepareReadFrames(ReadableFakeWebSocketStream::ASYNC, OK, frames);
+  set_stream(stream.Pass());
+  EXPECT_CALL(*event_interface_, OnAddChannelResponse(false, _));
+  EXPECT_CALL(*event_interface_, OnFlowControl(_));
+  MockFunction<void(int)> checkpoint;
+  TestClosure completion;
+  {
+    InSequence s;
+    EXPECT_CALL(checkpoint, Call(1));
+    EXPECT_CALL(*event_interface_, OnClosingHandshake());
+    EXPECT_CALL(*event_interface_,
+                OnDropChannel(kWebSocketErrorAbnormalClosure, _))
+        .WillOnce(InvokeClosure(completion.closure()));
+  }
+  CreateChannelAndConnectSuccessfully();
+  channel_->SetClosingHandshakeTimeoutForTesting(
+      TimeDelta::FromMilliseconds(kVeryTinyTimeoutMillis));
+  checkpoint.Call(1);
+  completion.WaitForResult();
+}
+
 // RFC6455 5.1 "a client MUST mask all frames that it sends to the server".
 // WebSocketChannel actually only sets the mask bit in the header, it doesn't
 // perform masking itself (not all transports actually use masking).
@@ -1733,5 +1818,122 @@ TEST_F(WebSocketChannelStreamTest, ProtocolError) {
   CreateChannelAndConnectSuccessfully();
 }
 
+// Set the closing handshake timeout to a very tiny value before connecting.
+class WebSocketChannelStreamTimeoutTest : public WebSocketChannelStreamTest {
+ protected:
+  WebSocketChannelStreamTimeoutTest() {}
+
+  virtual void CreateChannelAndConnectSuccessfully() OVERRIDE {
+    set_stream(mock_stream_.Pass());
+    CreateChannelAndConnect();
+    channel_->SetClosingHandshakeTimeoutForTesting(
+        TimeDelta::FromMilliseconds(kVeryTinyTimeoutMillis));
+    connect_data_.factory.connect_delegate->OnSuccess(stream_.Pass());
+  }
+};
+
+// In this case the server initiates the closing handshake with a Close
+// message. WebSocketChannel responds with a matching Close message, and waits
+// for the server to close the TCP/IP connection. The server never closes the
+// connection, so the closing handshake times out and WebSocketChannel closes
+// the connection itself.
+TEST_F(WebSocketChannelStreamTimeoutTest, ServerInitiatedCloseTimesOut) {
+  static const InitFrame frames[] = {
+      {FINAL_FRAME, WebSocketFrameHeader::kOpCodeClose,
+       NOT_MASKED,  CLOSE_DATA(NORMAL_CLOSURE, "OK")}};
+  static const InitFrame expected[] = {
+      {FINAL_FRAME, WebSocketFrameHeader::kOpCodeClose,
+       MASKED,      CLOSE_DATA(NORMAL_CLOSURE, "OK")}};
+  EXPECT_CALL(*mock_stream_, GetSubProtocol()).Times(AnyNumber());
+  EXPECT_CALL(*mock_stream_, ReadFrames(_, _))
+      .WillOnce(ReturnFrames(&frames))
+      .WillRepeatedly(Return(ERR_IO_PENDING));
+  MockFunction<void(int)> checkpoint;
+  TestClosure completion;
+  {
+    InSequence s;
+    EXPECT_CALL(*mock_stream_, WriteFrames(EqualsFrames(expected), _))
+        .WillOnce(Return(OK));
+    EXPECT_CALL(checkpoint, Call(1));
+    EXPECT_CALL(*mock_stream_, Close())
+        .WillOnce(InvokeClosure(completion.closure()));
+  }
+
+  CreateChannelAndConnectSuccessfully();
+  checkpoint.Call(1);
+  completion.WaitForResult();
+}
+
+// In this case the client initiates the closing handshake by sending a Close
+// message. WebSocketChannel waits for a Close message in response from the
+// server. The server never responds to the Close message, so the closing
+// handshake times out and WebSocketChannel closes the connection.
+TEST_F(WebSocketChannelStreamTimeoutTest, ClientInitiatedCloseTimesOut) {
+  static const InitFrame expected[] = {
+      {FINAL_FRAME, WebSocketFrameHeader::kOpCodeClose,
+       MASKED,      CLOSE_DATA(NORMAL_CLOSURE, "OK")}};
+  EXPECT_CALL(*mock_stream_, GetSubProtocol()).Times(AnyNumber());
+  EXPECT_CALL(*mock_stream_, ReadFrames(_, _))
+      .WillRepeatedly(Return(ERR_IO_PENDING));
+  TestClosure completion;
+  {
+    InSequence s;
+    EXPECT_CALL(*mock_stream_, WriteFrames(EqualsFrames(expected), _))
+        .WillOnce(Return(OK));
+    EXPECT_CALL(*mock_stream_, Close())
+        .WillOnce(InvokeClosure(completion.closure()));
+  }
+
+  CreateChannelAndConnectSuccessfully();
+  channel_->StartClosingHandshake(kWebSocketNormalClosure, "OK");
+  completion.WaitForResult();
+}
+
+// In this case the client initiates the closing handshake and the server
+// responds with a matching Close message. WebSocketChannel waits for the server
+// to close the TCP/IP connection, but it never does. The closing handshake
+// times out and WebSocketChannel closes the connection.
+TEST_F(WebSocketChannelStreamTimeoutTest, ConnectionCloseTimesOut) {
+  static const InitFrame expected[] = {
+      {FINAL_FRAME, WebSocketFrameHeader::kOpCodeClose,
+       MASKED,      CLOSE_DATA(NORMAL_CLOSURE, "OK")}};
+  static const InitFrame frames[] = {
+      {FINAL_FRAME, WebSocketFrameHeader::kOpCodeClose,
+       NOT_MASKED,  CLOSE_DATA(NORMAL_CLOSURE, "OK")}};
+  EXPECT_CALL(*mock_stream_, GetSubProtocol()).Times(AnyNumber());
+  TestClosure completion;
+  ScopedVector<WebSocketFrame>* read_frames = NULL;
+  CompletionCallback read_callback;
+  {
+    InSequence s;
+    // Copy the arguments to ReadFrames so that the test can call the callback
+    // after it has send the close message.
+    EXPECT_CALL(*mock_stream_, ReadFrames(_, _))
+        .WillOnce(DoAll(SaveArg<0>(&read_frames),
+                        SaveArg<1>(&read_callback),
+                        Return(ERR_IO_PENDING)));
+    // The first real event that happens is the client sending the Close
+    // message.
+    EXPECT_CALL(*mock_stream_, WriteFrames(EqualsFrames(expected), _))
+        .WillOnce(Return(OK));
+    // The |read_frames| callback is called (from this test case) at this
+    // point. ReadFrames is called again by WebSocketChannel, waiting for
+    // ERR_CONNECTION_CLOSED.
+    EXPECT_CALL(*mock_stream_, ReadFrames(_, _))
+        .WillOnce(Return(ERR_IO_PENDING));
+    // The timeout happens and so WebSocketChannel closes the stream.
+    EXPECT_CALL(*mock_stream_, Close())
+        .WillOnce(InvokeClosure(completion.closure()));
+  }
+
+  CreateChannelAndConnectSuccessfully();
+  channel_->StartClosingHandshake(kWebSocketNormalClosure, "OK");
+  ASSERT_TRUE(read_frames);
+  // Provide the "Close" message from the server.
+  *read_frames = CreateFrameVector(frames);
+  read_callback.Run(OK);
+  completion.WaitForResult();
+}
+
 }  // namespace
 }  // namespace net