Add blink::WebSocketHandshakeThrottle

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

 /*
  * Copyright (C) 2013 Google Inc. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
  * met:
  *
  *     * Redistributions of source code must retain the above copyright
  * notice, this list of conditions and the following disclaimer.
  *     * Redistributions in binary form must reproduce the above
@@ skipping 12 matching lines @@
  * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
 #include "modules/websockets/DocumentWebSocketChannel.h"
 
-#include <memory>
 #include "core/dom/DOMArrayBuffer.h"
 #include "core/dom/ExecutionContext.h"
 #include "core/dom/TaskRunnerHelper.h"
 #include "core/fileapi/FileReaderLoader.h"
 #include "core/fileapi/FileReaderLoaderClient.h"
 #include "core/frame/LocalFrame.h"
 #include "core/frame/LocalFrameClient.h"
+#include "core/frame/WebLocalFrameBase.h"
 #include "core/inspector/ConsoleMessage.h"
 #include "core/loader/DocumentLoader.h"
 #include "core/loader/FrameLoader.h"
 #include "core/loader/MixedContentChecker.h"
 #include "core/loader/SubresourceFilter.h"
 #include "core/loader/ThreadableLoadingContext.h"
+#include "core/page/ChromeClient.h"
+#include "core/page/Page.h"
 #include "core/probe/CoreProbes.h"
 #include "modules/websockets/InspectorWebSocketEvents.h"
 #include "modules/websockets/WebSocketChannelClient.h"
 #include "modules/websockets/WebSocketFrame.h"
 #include "modules/websockets/WebSocketHandleImpl.h"
 #include "platform/WebFrameScheduler.h"
 #include "platform/WebTaskRunner.h"
 #include "platform/loader/fetch/UniqueIdentifier.h"
 #include "platform/network/NetworkLog.h"
 #include "platform/network/WebSocketHandshakeRequest.h"
 #include "platform/weborigin/SecurityOrigin.h"
 #include "platform/wtf/Functional.h"
 #include "platform/wtf/PtrUtil.h"
 #include "public/platform/InterfaceProvider.h"
 #include "public/platform/Platform.h"
+#include "public/platform/WebSocketHandshakeThrottle.h"
 #include "public/platform/WebTraceLocation.h"
+#include "public/platform/WebURL.h"
 
 namespace blink {
 
 class DocumentWebSocketChannel::BlobLoader final
     : public GarbageCollectedFinalized<DocumentWebSocketChannel::BlobLoader>,
       public FileReaderLoaderClient {
  public:
   BlobLoader(PassRefPtr<BlobDataHandle>, DocumentWebSocketChannel*);
   ~BlobLoader() override {}
 
@@ skipping 55 matching lines @@
   channel_->DidFinishLoadingBlob(loader_->ArrayBufferResult());
   // |this| is deleted here.
 }
 
 void DocumentWebSocketChannel::BlobLoader::DidFail(
     FileError::ErrorCode error_code) {
   channel_->DidFailLoadingBlob(error_code);
   // |this| is deleted here.
 }
 
+struct DocumentWebSocketChannel::ConnectInfo {
+  ConnectInfo(const String& selected_protocol, const String& extensions)
+      : selected_protocol(selected_protocol), extensions(extensions) {}
+
+  String selected_protocol;

[yhirano, 2017/06/13 09:14:39]

const

[Adam Rice, 2017/06/13 12:53:42]

Done.

+  String extensions;

[yhirano, 2017/06/13 09:14:39]

ditto

[Adam Rice, 2017/06/13 12:53:42]

Done.

+};
+
+// static
+DocumentWebSocketChannel* DocumentWebSocketChannel::CreateForTesting(
+    Document* document,
+    WebSocketChannelClient* client,
+    std::unique_ptr<SourceLocation> location,
+    WebSocketHandle* handle,
+    std::unique_ptr<WebSocketHandshakeThrottle> handshake_throttle) {
+  return new DocumentWebSocketChannel(
+      ThreadableLoadingContext::Create(*document), client, std::move(location),
+      handle, std::move(handshake_throttle));
+}
+
+// static
+DocumentWebSocketChannel* DocumentWebSocketChannel::Create(
+    ThreadableLoadingContext* loading_context,
+    WebSocketChannelClient* client,
+    std::unique_ptr<SourceLocation> location,
+    WebSocketHandle* handle) {
+  return new DocumentWebSocketChannel(
+      loading_context, client, std::move(location), handle,
+      Platform::Current()->CreateWebSocketHandshakeThrottle());
+}
+
 DocumentWebSocketChannel::DocumentWebSocketChannel(
     ThreadableLoadingContext* loading_context,
     WebSocketChannelClient* client,
     std::unique_ptr<SourceLocation> location,
-    WebSocketHandle* handle)
+    WebSocketHandle* handle,
+    std::unique_ptr<WebSocketHandshakeThrottle> handshake_throttle)
     : handle_(WTF::WrapUnique(handle ? handle : new WebSocketHandleImpl())),
       client_(client),
       identifier_(CreateUniqueIdentifier()),
       loading_context_(loading_context),
       sending_quota_(0),
       received_data_size_for_flow_control_(
           kReceivedDataSizeForFlowControlHighWaterMark * 2),  // initial quota
       sent_size_of_top_message_(0),
-      location_at_construction_(std::move(location)) {}
+      location_at_construction_(std::move(location)),
+      handshake_throttle_(std::move(handshake_throttle)),
+      throttle_passed_(false) {}
 
 DocumentWebSocketChannel::~DocumentWebSocketChannel() {
   DCHECK(!blob_loader_);
 }
 
 bool DocumentWebSocketChannel::Connect(const KURL& url,
                                        const String& protocol) {
   NETWORK_DVLOG(1) << this << " connect()";
   if (!handle_)
     return false;
@@ skipping 53 matching lines @@
     // This is important so that the browser can show UI associated with
     // the WebSocket (e.g., for certificate errors).
     handle_->Initialize(GetDocument()->GetFrame()->GetInterfaceProvider());
   } else {
     handle_->Initialize(Platform::Current()->GetInterfaceProvider());
   }
   handle_->Connect(url, protocols, loading_context_->GetSecurityOrigin(),
                    loading_context_->FirstPartyForCookies(),
                    loading_context_->UserAgent(), this);
 
+  // TODO(ricea): Maybe lookup GetDocument()->GetFrame() less often?
+  if (handshake_throttle_ && GetDocument() && GetDocument()->GetFrame() &&
+      GetDocument()->GetFrame()->GetPage()) {
+    // TODO(ricea): We may need to do something special here for SharedWorkers
+    // and ServiceWorkers
+    // TODO(ricea): Figure out who owns this WebFrame object and how long it can
+    // be expected to live.
+    LocalFrame* frame = GetDocument()->GetFrame();
+    WebLocalFrame* web_frame =
+        frame->GetPage()->GetChromeClient().GetWebLocalFrameBase(frame);
+    handshake_throttle_->ThrottleHandshake(url, web_frame, this);
+  } else {
+    // Treat no throttle as success.
+    throttle_passed_ = true;
+  }
+
   FlowControlIfNecessary();
   TRACE_EVENT_INSTANT1("devtools.timeline", "WebSocketCreate",
                        TRACE_EVENT_SCOPE_THREAD, "data",
                        InspectorWebSocketCreateEvent::Data(
                            GetDocument(), identifier_, url, protocol));
   probe::didCreateWebSocket(GetDocument(), identifier_, url, protocol);
   return true;
 }
 
 void DocumentWebSocketChannel::Send(const CString& message) {
@@ skipping 99 matching lines @@
 void DocumentWebSocketChannel::Disconnect() {
   NETWORK_DVLOG(1) << this << " disconnect()";
   if (identifier_) {
     TRACE_EVENT_INSTANT1(
         "devtools.timeline", "WebSocketDestroy", TRACE_EVENT_SCOPE_THREAD,
         "data", InspectorWebSocketEvent::Data(GetDocument(), identifier_));
     probe::didCloseWebSocket(GetDocument(), identifier_);
   }
   connection_handle_for_scheduler_.reset();
   AbortAsyncOperations();
+  handshake_throttle_.reset();
   handle_.reset();
   client_ = nullptr;
   identifier_ = 0;
 }
 
 DocumentWebSocketChannel::Message::Message(const CString& text)
     : type(kMessageTypeText), text(text) {}
 
 DocumentWebSocketChannel::Message::Message(
     PassRefPtr<BlobDataHandle> blob_data_handle)
@@ skipping 72 matching lines @@
         break;
       case kMessageTypeBinaryAsCharVector:
         SendInternal(WebSocketHandle::kMessageTypeBinary,
                      message->vector_data->data(), message->vector_data->size(),
                      &consumed_buffered_amount);
         break;
       case kMessageTypeClose: {
         // No message should be sent from now on.
         DCHECK_EQ(messages_.size(), 1u);
         DCHECK_EQ(sent_size_of_top_message_, 0u);
+        handshake_throttle_.reset();
         handle_->Close(message->code, message->reason);
         messages_.pop_front();
         break;
       }
     }
   }
   if (client_ && consumed_buffered_amount > 0)
     client_->DidConsumeBufferedAmount(consumed_buffered_amount);
 }
 
 void DocumentWebSocketChannel::FlowControlIfNecessary() {
   if (!handle_ || received_data_size_for_flow_control_ <
                       kReceivedDataSizeForFlowControlHighWaterMark) {
     return;
   }
   handle_->FlowControl(received_data_size_for_flow_control_);
   received_data_size_for_flow_control_ = 0;
 }
 
 void DocumentWebSocketChannel::AbortAsyncOperations() {
   if (blob_loader_) {
     blob_loader_->Cancel();
     blob_loader_.Clear();
   }
 }
 
 void DocumentWebSocketChannel::HandleDidClose(bool was_clean,
                                               unsigned short code,
                                               const String& reason) {
+  handshake_throttle_.reset();
   handle_.reset();
   AbortAsyncOperations();
   if (!client_) {
     return;
   }
   WebSocketChannelClient* client = client_;
   client_ = nullptr;
   WebSocketChannelClient::ClosingHandshakeCompletionStatus status =
       was_clean ? WebSocketChannelClient::kClosingHandshakeComplete
                 : WebSocketChannelClient::kClosingHandshakeIncomplete;
@@ skipping 13 matching lines @@
                                           const String& selected_protocol,
                                           const String& extensions) {
   NETWORK_DVLOG(1) << this << " didConnect(" << handle << ", "
                    << String(selected_protocol) << ", " << String(extensions)
                    << ")";
 
   DCHECK(handle_);
   DCHECK_EQ(handle, handle_.get());
   DCHECK(client_);
 
+  if (!throttle_passed_) {
+    connect_info_ = WTF::MakeUnique<ConnectInfo>(selected_protocol, extensions);
+    return;
+  }
+
+  handshake_throttle_.reset();
+
   client_->DidConnect(selected_protocol, extensions);
 }
 
 void DocumentWebSocketChannel::DidStartOpeningHandshake(
     WebSocketHandle* handle,
     PassRefPtr<WebSocketHandshakeRequest> request) {
   NETWORK_DVLOG(1) << this << " didStartOpeningHandshake(" << handle << ")";
 
   DCHECK(handle_);
   DCHECK_EQ(handle, handle_.get());
@@ skipping 155 matching lines @@
     WebSocketHandle* handle) {
   NETWORK_DVLOG(1) << this << " didStartClosingHandshake(" << handle << ")";
 
   DCHECK(handle_);
   DCHECK_EQ(handle, handle_.get());
 
   if (client_)
     client_->DidStartClosingHandshake();
 }
 
+void DocumentWebSocketChannel::OnSuccess() {

[yhirano, 2017/06/13 09:14:39]

DCHECK(!throttle_passed);

[Adam Rice, 2017/06/13 12:53:42]

Done. I also added a check that handshake_throttle_ is non-null.

+  throttle_passed_ = true;
+  handshake_throttle_ = nullptr;
+  if (connect_info_) {
+    client_->DidConnect(std::move(connect_info_->selected_protocol),
+                        std::move(connect_info_->extensions));
+    connect_info_.reset();
+  }
+}
+
+void DocumentWebSocketChannel::OnError(const WebString& console_message) {

[yhirano, 2017/06/13 09:14:39]

DCHECK(!throttle_passed);

[Adam Rice, 2017/06/13 12:53:42]

Done.

+  handshake_throttle_ = nullptr;
+  FailAsError(console_message);
+}
+
 void DocumentWebSocketChannel::DidFinishLoadingBlob(DOMArrayBuffer* buffer) {
   blob_loader_.Clear();
   DCHECK(handle_);
   // The loaded blob is always placed on m_messages[0].
   DCHECK_GT(messages_.size(), 0u);
   DCHECK_EQ(messages_.front()->type, kMessageTypeBlob);
   // We replace it with the loaded blob.
   messages_.front() = new Message(buffer);
   ProcessSendQueue();
 }
 
 void DocumentWebSocketChannel::DidFailLoadingBlob(
     FileError::ErrorCode error_code) {
   blob_loader_.Clear();
   if (error_code == FileError::kAbortErr) {
     // The error is caused by cancel().
     return;
   }
   // FIXME: Generate human-friendly reason message.
   FailAsError("Failed to load Blob: error code = " +
               String::Number(error_code));
   // |this| can be deleted here.
 }
 
 void DocumentWebSocketChannel::TearDownFailedConnection() {
   // m_handle and m_client can be null here.
   connection_handle_for_scheduler_.reset();
+  handshake_throttle_.reset();
 
   if (client_)
     client_->DidError();
 
   HandleDidClose(false, kCloseEventCodeAbnormalClosure, String());
   // handleDidClose may delete this object.
 }
 
 bool DocumentWebSocketChannel::ShouldDisallowConnection(const KURL& url) {
   DCHECK(handle_);
@@ skipping 14 matching lines @@
   WebSocketChannel::Trace(visitor);
 }
 
 std::ostream& operator<<(std::ostream& ostream,
                          const DocumentWebSocketChannel* channel) {
   return ostream << "DocumentWebSocketChannel "
                  << static_cast<const void*>(channel);
 }
 
 }  // namespace blink