Enable websocket filtering via SubresourceFilter

third_party/WebKit/Source/core/loader/SubresourceFilter.cpp

Index: third_party/WebKit/Source/core/loader/SubresourceFilter.cpp
diff --git a/third_party/WebKit/Source/core/loader/SubresourceFilter.cpp b/third_party/WebKit/Source/core/loader/SubresourceFilter.cpp
index fe3c77e3abe26747c54b368c3344455a9011c1bf..081ca86edd9c40afebd29f60e014a8a8ce48409f 100644
--- a/third_party/WebKit/Source/core/loader/SubresourceFilter.cpp
+++ b/third_party/WebKit/Source/core/loader/SubresourceFilter.cpp
@@ -4,8 +4,11 @@
 
 #include "core/loader/SubresourceFilter.h"
 
+#include "core/dom/TaskRunnerHelper.h"
 #include "core/loader/DocumentLoader.h"
-#include "public/platform/WebDocumentSubresourceFilter.h"
+#include "platform/WebTaskRunner.h"
+#include "platform/weborigin/KURL.h"
+#include "public/platform/WebTraceLocation.h"
 
 namespace blink {
 
@@ -32,20 +35,44 @@ bool SubresourceFilter::allowLoad(
   // Pair<url string, context> -> LoadPolicy.
   WebDocumentSubresourceFilter::LoadPolicy loadPolicy =
       m_subresourceFilter->getLoadPolicy(resourceUrl, requestContext);
-  if (reportingPolicy == SecurityViolationReportingPolicy::Report) {
-    switch (loadPolicy) {
-      case WebDocumentSubresourceFilter::Allow:
-        break;
-      case WebDocumentSubresourceFilter::Disallow:
-        m_subresourceFilter->reportDisallowedLoad();
-      // fall through
-      case WebDocumentSubresourceFilter::WouldDisallow:
-        m_documentLoader->didObserveLoadingBehavior(
-            WebLoadingBehaviorSubresourceFilterMatch);
-        break;
-    }
-  }
+  if (reportingPolicy == SecurityViolationReportingPolicy::Report)
+    reportLoad(loadPolicy);
   return loadPolicy != WebDocumentSubresourceFilter::Disallow;
 }
 
+bool SubresourceFilter::allowWebSocketConnection(const KURL& url) {
+  // TODO(csharrison): Should probably have a new API for this in
+  // WebDocumentSubresourceFilter rather than sending subresource context.
+  // Alternatively, could augment the filter API so that we send a
+  // WebDocumentSubresourceFilterResourceType that matches
+  // subresource_filter::proto::ElementType.
+  WebDocumentSubresourceFilter::LoadPolicy loadPolicy =
+      m_subresourceFilter->getLoadPolicy(
+          url, WebURLRequest::RequestContextSubresource);
+
+  // Post a task to notify this load to avoid unduly blocking the worker
+  // thread. Note that this unconditionally calls reportLoad unlike allowLoad,
+  // because there aren't developer-invisible connections (like speculative
+  // preloads) happening here.
+  TaskRunnerHelper::get(TaskType::Networking, m_documentLoader->frame())

[engedy, 2017/03/06 16:31:01]

nit: Given that the DSF is totally not thread safe, would it be at all customary
in Blink to declare that:

DCHECK(taskRunner->runsTasksOnCurrentThread());

[Charlie Harrison, 2017/03/06 19:22:12]

Yep! Added.

+      ->postTask(BLINK_FROM_HERE, WTF::bind(&SubresourceFilter::reportLoad,
+                                            wrapPersistent(this), loadPolicy));
+  return loadPolicy != WebDocumentSubresourceFilter::Disallow;
+}
+
+void SubresourceFilter::reportLoad(
+    WebDocumentSubresourceFilter::LoadPolicy loadPolicy) {
+  switch (loadPolicy) {
+    case WebDocumentSubresourceFilter::Allow:
+      break;
+    case WebDocumentSubresourceFilter::Disallow:
+      m_subresourceFilter->reportDisallowedLoad();
+    // fall through
+    case WebDocumentSubresourceFilter::WouldDisallow:
+      m_documentLoader->didObserveLoadingBehavior(
+          WebLoadingBehaviorSubresourceFilterMatch);
+      break;
+  }
+}
+
 }  // namespace blink