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 2d0242d0a331b7998ee45e2f4a90fef6a86dae27..898cf9e152a7c08439683643a15af5953628267d 100644
--- a/third_party/WebKit/Source/modules/websockets/DocumentWebSocketChannel.cpp
+++ b/third_party/WebKit/Source/modules/websockets/DocumentWebSocketChannel.cpp
@@ -34,11 +34,13 @@
 #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/inspector/ConsoleMessage.h"
 #include "core/inspector/InspectorInstrumentation.h"
+#include "core/loader/DocumentLoader.h"
 #include "core/loader/FrameLoader.h"
 #include "core/loader/FrameLoaderClient.h"
 #include "core/loader/MixedContentChecker.h"
@@ -53,6 +55,7 @@
 #include "platform/weborigin/SecurityOrigin.h"
 #include "public/platform/InterfaceProvider.h"
 #include "public/platform/Platform.h"
+#include "public/platform/WebDocumentSubresourceFilter.h"
 #include "wtf/PtrUtil.h"
 
 namespace blink {
@@ -203,6 +206,10 @@ bool DocumentWebSocketChannel::connect(const KURL& url,
                            document(), m_identifier, url, protocol));
   InspectorInstrumentation::didCreateWebSocket(document(), m_identifier, url,
                                                protocol);
+  TaskRunnerHelper::get(TaskType::Networking, document())
+      ->postTask(BLINK_FROM_HERE,
+                 WTF::bind(&DocumentWebSocketChannel::maybeFilterConnection,
+                           wrapPersistent(this)));
   return true;
 }
 
@@ -662,6 +669,26 @@ void DocumentWebSocketChannel::didFailLoadingBlob(
   // |this| can be deleted here.
 }
 
+void DocumentWebSocketChannel::maybeFilterConnection() {
+  if (!m_handle)
+    return;
+  DocumentLoader* loader = document()->loader();
+  if (!loader)
+    return;
+  WebDocumentSubresourceFilter* subresourceFilter = loader->subresourceFilter();
+  if (!subresourceFilter)
+    return;
+
+  WebDocumentSubresourceFilter::LoadPolicy loadPolicy =
+      subresourceFilter->getLoadPolicy(m_url);
+  if (loadPolicy == WebDocumentSubresourceFilter::Disallow)
+    subresourceFilter->reportDisallowedLoad();
+  if (loadPolicy != WebDocumentSubresourceFilter::Allow)
+    loader->didObserveLoadingBehavior(WebLoadingBehaviorSubresourceFilterMatch);
+  // TODO(csharrison): Figure out the ideal event code / reason to send here.
+  close(CloseEventCodeNotSpecified, "");

[yhirano, 2017/02/24 03:55:29]

if (loadPolicy == Allow)
 return;

?

[yhirano, 2017/02/24 03:55:30]

Is 1008 good for this case?

[Charlie Harrison, 2017/02/24 21:13:02]

Done. Oops :P

[Charlie Harrison, 2017/02/24 21:13:02]

Done.

+}
+
 DEFINE_TRACE(DocumentWebSocketChannel) {
   visitor->trace(m_blobLoader);
   visitor->trace(m_messages);