Enable websocket filtering via SubresourceFilter

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 16 matching lines @@
  * (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/Document.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/inspector/ConsoleMessage.h"
 #include "core/inspector/InspectorInstrumentation.h"
+#include "core/loader/DocumentLoader.h"
+#include "core/loader/DocumentSubresourceFilterClient.h"
 #include "core/loader/FrameLoader.h"
 #include "core/loader/MixedContentChecker.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/loader/fetch/UniqueIdentifier.h"
 #include "platform/network/NetworkLog.h"
 #include "platform/network/WebSocketHandshakeRequest.h"
@@ skipping 132 matching lines @@
       document()->frame()->interfaceProvider() !=
           InterfaceProvider::getEmptyInterfaceProvider()) {
     // Initialize the WebSocketHandle with the frame's InterfaceProvider to
     // provide the WebSocket implementation with context about this frame.
     // This is important so that the browser can show UI associated with
     // the WebSocket (e.g., for certificate errors).
     m_handle->initialize(document()->frame()->interfaceProvider());
   } else {
     m_handle->initialize(Platform::current()->interfaceProvider());
   }
+
+  // If the connection needs to be filtered, asynchronously fail. Note that

[yhirano, 2017/03/01 06:40:58]

Please move this section before L187.

[Charlie Harrison, 2017/03/01 18:55:57]

Done.

+  // returning "true" just indicates that this was not synchronous security
+  // error.
+  if (shouldFilterConnection(url)) {
+    TaskRunnerHelper::get(TaskType::Networking, document())
+        ->postTask(
+            BLINK_FROM_HERE,
+            WTF::bind(&DocumentWebSocketChannel::failWithClosureCode,
+                      wrapPersistent(this), CloseEventCodePolicyViolation));
+    return true;
+  }
+
   m_handle->connect(url, protocols, document()->getSecurityOrigin(),
                     document()->firstPartyForCookies(), document()->userAgent(),
                     this);
 
   flowControlIfNecessary();
   TRACE_EVENT_INSTANT1("devtools.timeline", "WebSocketCreate",
                        TRACE_EVENT_SCOPE_THREAD, "data",
                        InspectorWebSocketCreateEvent::data(
                            document(), m_identifier, url, protocol));
   InspectorInstrumentation::didCreateWebSocket(document(), m_identifier, url,
@@ skipping 82 matching lines @@
   m_messages.append(new Message(codeToSend, reason));
   processSendQueue();
 }
 
 void DocumentWebSocketChannel::fail(const String& reason,
                                     MessageLevel level,
                                     std::unique_ptr<SourceLocation> location) {
   NETWORK_DVLOG(1) << this << " fail(" << reason << ")";
   // m_handle and m_client can be null here.
 
-  connection_handle_for_scheduler_.reset();
-
   InspectorInstrumentation::didReceiveWebSocketFrameError(document(),
                                                           m_identifier, reason);
   const String message = "WebSocket connection to '" + m_url.elidedString() +
                          "' failed: " + reason;
   document()->addConsoleMessage(ConsoleMessage::create(

[yhirano, 2017/03/01 06:40:58]

Is showing a console message useful for developers when filtered? If so, can you
move this to failWithClosureCode?

[Charlie Harrison, 2017/03/01 18:55:57]

It is useful. Eventually we will move the notification logic to the
SubresourceFilter code, but for now we can keep the generic error. I've
basically replaced fail with failWithClosureCode and sent an empty reason string
here. 

       JSMessageSource, level, message, std::move(location)));
-
-  if (m_client)
-    m_client->didError();
-  // |reason| is only for logging and should not be provided for scripts,
-  // hence close reason must be empty.
-  handleDidClose(false, CloseEventCodeAbnormalClosure, String());
-  // handleDidClose may delete this object.
+  failWithClosureCode(CloseEventCodeAbnormalClosure);
 }
 
 void DocumentWebSocketChannel::disconnect() {
   NETWORK_DVLOG(1) << this << " disconnect()";
   if (m_identifier) {
     TRACE_EVENT_INSTANT1(
         "devtools.timeline", "WebSocketDestroy", TRACE_EVENT_SCOPE_THREAD,
         "data", InspectorWebSocketEvent::data(document(), m_identifier));
     InspectorInstrumentation::didCloseWebSocket(document(), m_identifier);
   }
@@ skipping 333 matching lines @@
   m_blobLoader.clear();
   if (errorCode == FileError::kAbortErr) {
     // The error is caused by cancel().
     return;
   }
   // FIXME: Generate human-friendly reason message.
   failAsError("Failed to load Blob: error code = " + String::number(errorCode));
   // |this| can be deleted here.
 }
 
+void DocumentWebSocketChannel::failWithClosureCode(unsigned short code) {
+  connection_handle_for_scheduler_.reset();
+
+  if (m_client)
+    m_client->didError();
+  // |reason| is only for logging and should not be provided for scripts,
+  // hence close reason must be empty.
+  handleDidClose(false, code, String());
+  // handleDidClose may delete this object.
+}
+
+bool DocumentWebSocketChannel::shouldFilterConnection(const KURL& url) {
+  if (!m_handle)
+    return false;
+  DocumentLoader* loader = document()->loader();
+  if (!loader)
+    return false;
+  DocumentSubresourceFilterClient* subresourceFilter =
+      loader->subresourceFilter();
+  if (!subresourceFilter)
+    return false;
+  return !subresourceFilter->allowWebSocketConnection(url);
+}
+
 DEFINE_TRACE(DocumentWebSocketChannel) {
   visitor->trace(m_blobLoader);
   visitor->trace(m_messages);
   visitor->trace(m_client);
   visitor->trace(m_document);
   WebSocketChannel::trace(visitor);
 }
 
 std::ostream& operator<<(std::ostream& ostream,
                          const DocumentWebSocketChannel* channel) {
   return ostream << "DocumentWebSocketChannel "
                  << static_cast<const void*>(channel);
 }
 
 }  // namespace blink