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/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"
 #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"
 #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 {
 
 class DocumentWebSocketChannel::BlobLoader final
     : public GarbageCollectedFinalized<DocumentWebSocketChannel::BlobLoader>,
       public FileReaderLoaderClient {
  public:
   BlobLoader(PassRefPtr<BlobDataHandle>, DocumentWebSocketChannel*);
   ~BlobLoader() override {}
@@ skipping 130 matching lines @@
                     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,
                                                protocol);
+  TaskRunnerHelper::get(TaskType::Networking, document())
+      ->postTask(BLINK_FROM_HERE,
+                 WTF::bind(&DocumentWebSocketChannel::maybeFilterConnection,
+                           wrapPersistent(this)));
   return true;
 }
 
 void DocumentWebSocketChannel::send(const CString& message) {
   NETWORK_DVLOG(1) << this << " sendText(" << message << ")";
   // FIXME: Change the inspector API to show the entire message instead
   // of individual frames.
   InspectorInstrumentation::didSendWebSocketFrame(
       document(), m_identifier, WebSocketFrame::OpCodeText, true,
       message.data(), message.length());
@@ skipping 439 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::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);
   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