Enable websocket filtering via SubresourceFilter

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

 // Copyright 2017 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "core/loader/DocumentSubresourceFilterClient.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 {
 
 // static
 DocumentSubresourceFilterClient* DocumentSubresourceFilterClient::create(
     DocumentLoader* loader,
     std::unique_ptr<WebDocumentSubresourceFilter> filter) {
   return new DocumentSubresourceFilterClient(loader, std::move(filter));
 }
 
 DocumentSubresourceFilterClient::DocumentSubresourceFilterClient(
     DocumentLoader* documentLoader,
     std::unique_ptr<WebDocumentSubresourceFilter> subresourceFilter)
     : m_documentLoader(documentLoader),
       m_subresourceFilter(std::move(subresourceFilter)) {}
 
 DocumentSubresourceFilterClient::~DocumentSubresourceFilterClient() {}
 
 bool DocumentSubresourceFilterClient::allowLoad(
     const KURL& resourceUrl,
     WebURLRequest::RequestContext requestContext,
     ReportingPolicy reportingPolicy) {
   // TODO(csharrison): Implement a caching layer here which is a HashMap of
   // Pair<url string, context> -> LoadPolicy.
   WebDocumentSubresourceFilter::LoadPolicy loadPolicy =
       m_subresourceFilter->getLoadPolicy(resourceUrl, requestContext);
-  if (reportingPolicy == 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 == Report)
+    notifyLoad(resourceUrl, loadPolicy);
   return loadPolicy != WebDocumentSubresourceFilter::Disallow;
 }
 
+bool DocumentSubresourceFilterClient::allowWebSocketConnection(
+    const KURL& resourceUrl) {
+  // TODO(csharrison): Should probably have a new API for this in
+  // WebDocumentSubresourceFilter rather than sending the unspecified context.

[Charlie Harrison, 2017/02/28 15:11:00]

s/unspecified/subresource

+  WebDocumentSubresourceFilter::LoadPolicy loadPolicy =
+      m_subresourceFilter->getLoadPolicy(
+          resourceUrl, WebURLRequest::RequestContextSubresource);
+
+  // Post a task to notify this load to avoid unduly blocking the worker
+  // thread.
+  TaskRunnerHelper::get(TaskType::Networking, m_documentLoader->frame())
+      ->postTask(BLINK_FROM_HERE,
+                 WTF::bind(&DocumentSubresourceFilterClient::notifyLoad,
+                           wrapPersistent(this), resourceUrl, loadPolicy));
+  return loadPolicy != WebDocumentSubresourceFilter::Disallow;
+}
+
+void DocumentSubresourceFilterClient::notifyLoad(
+    const KURL& resourceUrl,
+    WebDocumentSubresourceFilter::LoadPolicy loadPolicy) {
+  // TODO(csharrison): Add devtools logging here.
+  switch (loadPolicy) {
+    case WebDocumentSubresourceFilter::Allow:
+      break;
+    case WebDocumentSubresourceFilter::Disallow:
+      m_subresourceFilter->reportDisallowedLoad();
+    // fall through
+    case WebDocumentSubresourceFilter::WouldDisallow:
+      m_documentLoader->didObserveLoadingBehavior(
+          WebLoadingBehaviorSubresourceFilterMatch);
+      break;
+  }
+}
+
 }  // namespace blink