Enable websocket filtering via SubresourceFilter

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

Index: third_party/WebKit/Source/modules/websockets/DocumentWebSocketChannel.cpp
diff --git a/third_party/WebKit/Source/modules/websockets/DocumentWebSocketChannel.cpp b/third_party/WebKit/Source/modules/websockets/DocumentWebSocketChannel.cpp
index 743f06d89a4c106ab6e7fe419c23c658ca828319..0489a8c488b7f76a3da0a4bf44c517fc143f2cc1 100644
--- a/third_party/WebKit/Source/modules/websockets/DocumentWebSocketChannel.cpp
+++ b/third_party/WebKit/Source/modules/websockets/DocumentWebSocketChannel.cpp
@@ -34,12 +34,15 @@
 #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"
@@ -192,6 +195,19 @@ bool DocumentWebSocketChannel::connect(const KURL& url,
   } 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);
@@ -294,21 +310,13 @@ void DocumentWebSocketChannel::fail(const String& reason,
   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() {
@@ -662,6 +670,30 @@ void DocumentWebSocketChannel::didFailLoadingBlob(
   // |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);