SafeBrowsing support for WebSocket (post-network-servicification)

components/safe_browsing/renderer/websocket_sb_handshake_throttle.cc

Index: components/safe_browsing/renderer/websocket_sb_handshake_throttle.cc
diff --git a/components/safe_browsing/renderer/websocket_sb_handshake_throttle.cc b/components/safe_browsing/renderer/websocket_sb_handshake_throttle.cc
new file mode 100644
index 0000000000000000000000000000000000000000..60b626c4c564824805866a79049c22da67d0957a
--- /dev/null
+++ b/components/safe_browsing/renderer/websocket_sb_handshake_throttle.cc
@@ -0,0 +1,83 @@
+// 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 "components/safe_browsing/renderer/websocket_sb_handshake_throttle.h"
+
+#include <utility>
+
+#include "base/bind.h"
+#include "base/command_line.h"
+#include "base/logging.h"
+#include "base/memory/ptr_util.h"
+#include "base/strings/stringprintf.h"
+#include "content/public/common/content_switches.h"
+#include "content/public/common/resource_type.h"
+#include "content/public/common/service_names.mojom.h"
+#include "content/public/renderer/render_frame.h"
+#include "content/public/renderer/render_thread.h"
+#include "ipc/ipc_message.h"
+#include "mojo/public/cpp/bindings/interface_request.h"
+#include "services/service_manager/public/cpp/connector.h"
+#include "third_party/WebKit/public/platform/WebString.h"
+#include "third_party/WebKit/public/platform/WebURL.h"
+
+namespace safe_browsing {
+
+// static
+std::unique_ptr<WebSocketSBHandshakeThrottle>
+WebSocketSBHandshakeThrottle::MaybeCreate() {
+  // TODO(ricea): Is there some way to avoid doing a string lookup every time
+  // here?
+  if (!base::CommandLine::ForCurrentProcess()->HasSwitch(
+          switches::kEnableNetworkService)) {
+    return nullptr;
+  }
+  mojom::SafeBrowsingPtr safe_browsing;
+  content::RenderThread::Get()->GetConnector()->BindInterface(

[yzshen1, 2017/06/20 19:52:19]

You could consider passing in a (not-owned) SafeBrowsing interface pointer, and
let all the WSSBHT instances use the same SafeBrowing interface pointer. That is
what RendererURLLoaderThrottle does.

Creating a SafeBrowsing interface pointer requires an IPC message to the browser
and creating the corresponding impl object. It is nice to avoid this cost, it is
much bigger than the string lookup above. :)

[Adam Rice, 2017/06/21 06:08:35]

Done. I realised I can share the same interface pointer with
RendererURLLoaderThrottle in //chrome for added efficiency. I've kept the
initialisation logic inside MaybeCreate() so that I don't need to duplicate it
for WebView.

[yzshen1, 2017/06/21 16:55:06]

I think it may make more sense to initialize safe_browsing_service in
*ContentRendererClient: it is shared with RendererURLLoaderThrottle. If we do
lazy-init here, we need to duplicate it in RendererURLLoaderThrottle as well.
It feel more natural to let *ContentRendererClient do the work.

[Adam Rice, 2017/06/22 13:42:44]

Done. I factored the shared code out into a new method. We may need to mess it
up again later if we enable the HTTP and WebSocket implementations at different
times.

+      content::mojom::kBrowserServiceName, &safe_browsing);
+  return base::MakeUnique<WebSocketSBHandshakeThrottle>(
+      std::move(safe_browsing));
+}
+
+WebSocketSBHandshakeThrottle::WebSocketSBHandshakeThrottle(
+    mojom::SafeBrowsingPtr safe_browsing)
+    : callbacks_(nullptr), safe_browsing_(std::move(safe_browsing)) {}
+
+WebSocketSBHandshakeThrottle::~WebSocketSBHandshakeThrottle() {}
+
+void WebSocketSBHandshakeThrottle::ThrottleHandshake(
+    const blink::WebURL& url,
+    blink::WebLocalFrame* web_local_frame,
+    blink::WebCallbacks<void, const blink::WebString&>* callbacks) {
+  DCHECK(!callbacks_);
+  DCHECK(!url_checker_);
+  callbacks_ = callbacks;
+  url_ = url;
+  int render_frame_id = MSG_ROUTING_NONE;
+  if (web_local_frame) {
+    render_frame_id =
+        content::RenderFrame::FromWebFrame(web_local_frame)->GetRoutingID();
+  }
+  int load_flags = 0;
+  // This use of Unretained() is safe because this object will not be called
+  // back after |safe_browsing_| is destroyed.
+  safe_browsing_->CreateCheckerAndCheck(
+      render_frame_id, mojo::MakeRequest(&url_checker_), url, load_flags,
+      content::RESOURCE_TYPE_SUB_RESOURCE,
+      base::BindOnce(&WebSocketSBHandshakeThrottle::OnCheckResult,
+                     base::Unretained(this)));
+}
+
+void WebSocketSBHandshakeThrottle::OnCheckResult(bool safe) {
+  if (safe) {
+    callbacks_->OnSuccess();
+  } else {
+    callbacks_->OnError(blink::WebString::FromUTF8(base::StringPrintf(
+        "WebSocket connection to %s failed safe browsing check",
+        url_.spec().c_str())));
+  }
+  // |this| is destroyed here.
+}
+
+}  // namespace safe_browsing