Add blink::WebSocketHandshakeThrottle

third_party/WebKit/Source/modules/websockets/DocumentWebSocketChannelTest.cpp

 // Copyright 2014 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 "modules/websockets/WebSocketChannel.h"
+#include "modules/websockets/DocumentWebSocketChannel.h"
 
 #include <stdint.h>
 #include <memory>
 #include "core/dom/DOMArrayBuffer.h"
 #include "core/dom/Document.h"
 #include "core/fileapi/Blob.h"
 #include "core/testing/DummyPageHolder.h"
-#include "modules/websockets/DocumentWebSocketChannel.h"
+#include "modules/websockets/WebSocketChannel.h"
 #include "modules/websockets/WebSocketChannelClient.h"
 #include "modules/websockets/WebSocketHandle.h"
 #include "modules/websockets/WebSocketHandleClient.h"
 #include "platform/heap/Handle.h"
 #include "platform/weborigin/KURL.h"
 #include "platform/wtf/PtrUtil.h"
 #include "platform/wtf/Vector.h"
 #include "platform/wtf/text/WTFString.h"
+#include "public/platform/WebCallbacks.h"
+#include "public/platform/WebSocketHandshakeThrottle.h"
+#include "public/platform/WebURL.h"
+#include "public/web/WebLocalFrame.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 using testing::_;
 using testing::InSequence;
 using testing::PrintToString;
 using testing::AnyNumber;
 using testing::SaveArg;
 
 namespace blink {
@@ skipping 50 matching lines @@
                     SecurityOrigin*,
                     const KURL&,
                     const String&,
                     WebSocketHandleClient*));
   MOCK_METHOD4(Send,
                void(bool, WebSocketHandle::MessageType, const char*, size_t));
   MOCK_METHOD1(FlowControl, void(int64_t));
   MOCK_METHOD2(Close, void(unsigned short, const String&));
 };
 
+class MockWebSocketHandshakeThrottle : public WebSocketHandshakeThrottle {
+ public:
+  static MockWebSocketHandshakeThrottle* Create() {
+    return new testing::StrictMock<MockWebSocketHandshakeThrottle>();
+  }
+  MockWebSocketHandshakeThrottle() {}
+  ~MockWebSocketHandshakeThrottle() override { Destructor(); }
+
+  MOCK_METHOD3(ThrottleHandshake,
+               void(const WebURL&,
+                    WebLocalFrame*,
+                    WebCallbacks<void, const WebString&>*));
+
+  // This method is used to allow us to require that the destructor is called at
+  // a particular time.
+  MOCK_METHOD0(Destructor, void());
+};
+
 class DocumentWebSocketChannelTest : public ::testing::Test {
  public:
   DocumentWebSocketChannelTest()
       : page_holder_(DummyPageHolder::Create()),
         channel_client_(MockWebSocketChannelClient::Create()),
         handle_(MockWebSocketHandle::Create()),
-        channel_(DocumentWebSocketChannel::Create(&page_holder_->GetDocument(),
-                                                  channel_client_.Get(),
-                                                  SourceLocation::Capture(),
-                                                  Handle())),
+        handshake_throttle_(nullptr),
         sum_of_consumed_buffered_amount_(0) {
     ON_CALL(*ChannelClient(), DidConsumeBufferedAmount(_))
         .WillByDefault(Invoke(
             this, &DocumentWebSocketChannelTest::DidConsumeBufferedAmount));
   }
 
   ~DocumentWebSocketChannelTest() { Channel()->Disconnect(); }
 
+  void SetUp() override {
+    channel_ = DocumentWebSocketChannel::CreateForTesting(
+        &page_holder_->GetDocument(), channel_client_.Get(),
+        SourceLocation::Capture(), Handle(),
+        WTF::WrapUnique(handshake_throttle_));
+  }
+
   MockWebSocketChannelClient* ChannelClient() { return channel_client_.Get(); }
 
   WebSocketChannel* Channel() {
     return static_cast<WebSocketChannel*>(channel_.Get());
   }
 
   WebSocketHandleClient* HandleClient() {
     return static_cast<WebSocketHandleClient*>(channel_.Get());
   }
 
+  WebCallbacks<void, const WebString&>* WebCallbacks() {
+    return channel_.Get();
+  }
+
   MockWebSocketHandle* Handle() { return handle_; }
 
   void DidConsumeBufferedAmount(unsigned long a) {
     sum_of_consumed_buffered_amount_ += a;
   }
 
   void Connect() {
     {
       InSequence s;
       EXPECT_CALL(*Handle(), Initialize(_));
       EXPECT_CALL(*Handle(), Connect(KURL(KURL(), "ws://localhost/"), _, _, _,
                                      _, HandleClient()));
       EXPECT_CALL(*Handle(), FlowControl(65536));
       EXPECT_CALL(*ChannelClient(), DidConnect(String("a"), String("b")));
     }
     EXPECT_TRUE(Channel()->Connect(KURL(KURL(), "ws://localhost/"), "x"));
     HandleClient()->DidConnect(Handle(), String("a"), String("b"));
     ::testing::Mock::VerifyAndClearExpectations(this);
   }
 
   std::unique_ptr<DummyPageHolder> page_holder_;
   Persistent<MockWebSocketChannelClient> channel_client_;
   MockWebSocketHandle* handle_;
+  // |handshake_throttle_| is owned by |channel_| once SetUp() has been called.
+  MockWebSocketHandshakeThrottle* handshake_throttle_;
   Persistent<DocumentWebSocketChannel> channel_;
   unsigned long sum_of_consumed_buffered_amount_;
 };
 
 MATCHER_P2(MemEq,
            p,
            len,
            std::string("pointing to memory") + (negation ? " not" : "") +
                " equal to \"" +
                std::string(p, len) +
@@ skipping 615 matching lines @@
     EXPECT_CALL(
         *ChannelClient(),
         DidClose(WebSocketChannelClient::kClosingHandshakeIncomplete,
                  WebSocketChannel::kCloseEventCodeAbnormalClosure, String()));
   }
 
   Channel()->Fail("fail message from WebSocket", kErrorMessageLevel,
                   SourceLocation::Create(String(), 0, 0, nullptr));
 }
 
+class DocumentWebSocketChannelHandshakeThrottleTest
+    : public DocumentWebSocketChannelTest {
+ public:
+  DocumentWebSocketChannelHandshakeThrottleTest() {
+    handshake_throttle_ = MockWebSocketHandshakeThrottle::Create();
+  }
+
+  // Expectations for the normal result of calling Channel()->Connect() with a
+  // non-null throttle.
+  void NormalHandshakeExpectations() {
+    EXPECT_CALL(*Handle(), Initialize(_));
+    EXPECT_CALL(*Handle(), Connect(_, _, _, _, _, _));
+    EXPECT_CALL(*Handle(), FlowControl(_));
+    EXPECT_CALL(*handshake_throttle_, ThrottleHandshake(_, _, _));
+  }
+
+  static KURL url() { return KURL(KURL(), "ws://localhost/"); }
+};
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest, ThrottleArguments) {
+  EXPECT_CALL(*Handle(), Initialize(_));
+  EXPECT_CALL(*Handle(), Connect(_, _, _, _, _, _));
+  EXPECT_CALL(*Handle(), FlowControl(_));
+  EXPECT_CALL(*handshake_throttle_,
+              ThrottleHandshake(WebURL(url()), _, WebCallbacks()));
+  EXPECT_CALL(*handshake_throttle_, Destructor());
+  Channel()->Connect(url(), "");
+}
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest, ThrottleSucceedsFirst) {
+  Checkpoint checkpoint;
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(checkpoint, Call(1));
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(checkpoint, Call(2));
+    EXPECT_CALL(*ChannelClient(), DidConnect(String("a"), String("b")));
+  }
+  Channel()->Connect(url(), "");
+  checkpoint.Call(1);
+  WebCallbacks()->OnSuccess();
+  checkpoint.Call(2);
+  HandleClient()->DidConnect(Handle(), String("a"), String("b"));
+}
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest, HandshakeSucceedsFirst) {
+  Checkpoint checkpoint;
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(checkpoint, Call(1));
+    EXPECT_CALL(checkpoint, Call(2));
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*ChannelClient(), DidConnect(String("a"), String("b")));
+  }
+  Channel()->Connect(url(), "");
+  checkpoint.Call(1);
+  HandleClient()->DidConnect(Handle(), String("a"), String("b"));
+  checkpoint.Call(2);
+  WebCallbacks()->OnSuccess();
+}
+
+// This happens if JS code calls close() during the handshake.
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest, FailDuringThrottle) {
+  Checkpoint checkpoint;
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*ChannelClient(), DidError());
+    EXPECT_CALL(*ChannelClient(), DidClose(_, _, _));
+    EXPECT_CALL(checkpoint, Call(1));
+  }
+  Channel()->Connect(url(), "");
+  Channel()->Fail("close during handshake", kWarningMessageLevel,
+                  SourceLocation::Create(String(), 0, 0, nullptr));
+  checkpoint.Call(1);
+}
+
+// It makes no difference to the behaviour if the WebSocketHandle has actually
+// connected.
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest,
+       FailDuringThrottleAfterConnect) {
+  Checkpoint checkpoint;
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*ChannelClient(), DidError());
+    EXPECT_CALL(*ChannelClient(), DidClose(_, _, _));
+    EXPECT_CALL(checkpoint, Call(1));
+  }
+  Channel()->Connect(url(), "");
+  HandleClient()->DidConnect(Handle(), String("a"), String("b"));
+  Channel()->Fail("close during handshake", kWarningMessageLevel,
+                  SourceLocation::Create(String(), 0, 0, nullptr));
+  checkpoint.Call(1);
+}
+
+// This happens if the JS context is destroyed during the handshake.
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest, CloseDuringThrottle) {
+  Checkpoint checkpoint;
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*Handle(), Close(_, _));
+    EXPECT_CALL(checkpoint, Call(1));
+  }
+  Channel()->Connect(url(), "");
+  Channel()->Close(DocumentWebSocketChannel::kCloseEventCodeGoingAway, "");
+  checkpoint.Call(1);
+}
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest,
+       CloseDuringThrottleAfterConnect) {
+  Checkpoint checkpoint;
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*Handle(), Close(_, _));
+    EXPECT_CALL(checkpoint, Call(1));
+  }
+  Channel()->Connect(url(), "");
+  HandleClient()->DidConnect(Handle(), String("a"), String("b"));
+  Channel()->Close(DocumentWebSocketChannel::kCloseEventCodeGoingAway, "");
+  checkpoint.Call(1);
+}
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest,
+       DisconnectDuringThrottle) {
+  Checkpoint checkpoint;
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(checkpoint, Call(1));
+  }
+  Channel()->Connect(url(), "");
+  Channel()->Disconnect();
+  checkpoint.Call(1);
+}
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest,
+       DisconnectDuringThrottleAfterConnect) {
+  Checkpoint checkpoint;
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(checkpoint, Call(1));
+  }
+  Channel()->Connect(url(), "");
+  HandleClient()->DidConnect(Handle(), String("a"), String("b"));
+  Channel()->Disconnect();
+  checkpoint.Call(1);
+}
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest,
+       ThrottleReportsErrorBeforeConnect) {
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*ChannelClient(), DidError());
+    EXPECT_CALL(*ChannelClient(), DidClose(_, _, _));
+  }
+  Channel()->Connect(url(), "");
+  WebCallbacks()->OnError("Connection blocked by throttle");
+}
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest,
+       ThrottleReportsErrorAfterConnect) {
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*ChannelClient(), DidError());
+    EXPECT_CALL(*ChannelClient(), DidClose(_, _, _));
+  }
+  Channel()->Connect(url(), "");
+  HandleClient()->DidConnect(Handle(), String("a"), String("b"));
+  WebCallbacks()->OnError("Connection blocked by throttle");
+}
+
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest,
+       ConnectFailBeforeThrottle) {
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*ChannelClient(), DidError());
+    EXPECT_CALL(*ChannelClient(), DidClose(_, _, _));
+  }
+  Channel()->Connect(url(), "");
+  HandleClient()->DidFail(Handle(), "connect failed");
+}
+
+// TODO(ricea): Can this actually happen?
+TEST_F(DocumentWebSocketChannelHandshakeThrottleTest,
+       ConnectCloseBeforeThrottle) {
+  NormalHandshakeExpectations();
+  {
+    InSequence s;
+    EXPECT_CALL(*handshake_throttle_, Destructor());
+    EXPECT_CALL(*ChannelClient(), DidClose(_, _, _));
+  }
+  Channel()->Connect(url(), "");
+  HandleClient()->DidClose(Handle(), false,
+                           WebSocketChannel::kCloseEventCodeProtocolError,
+                           "connect error");
+}
+
 }  // namespace
 
 }  // namespace blink