SafeBrowsing support for WebSocket (post-network-servicification)

components/safe_browsing/renderer/websocket_sb_handshake_throttle_unittest.cc

Index: components/safe_browsing/renderer/websocket_sb_handshake_throttle_unittest.cc
diff --git a/components/safe_browsing/renderer/websocket_sb_handshake_throttle_unittest.cc b/components/safe_browsing/renderer/websocket_sb_handshake_throttle_unittest.cc
new file mode 100644
index 0000000000000000000000000000000000000000..84abecb03ef73fff00433c95abd2e60efb523efe
--- /dev/null
+++ b/components/safe_browsing/renderer/websocket_sb_handshake_throttle_unittest.cc
@@ -0,0 +1,137 @@
+// 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/callback.h"
+#include "base/memory/ptr_util.h"
+#include "base/run_loop.h"
+#include "base/test/scoped_task_environment.h"
+#include "components/safe_browsing/common/safe_browsing.mojom.h"
+#include "content/public/common/resource_type.h"
+#include "ipc/ipc_message.h"
+#include "mojo/public/cpp/bindings/binding.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "third_party/WebKit/public/platform/WebString.h"
+#include "third_party/WebKit/public/platform/WebURL.h"
+
+namespace safe_browsing {
+
+namespace {
+
+constexpr char kTestUrl[] = "wss://test/";
+
+class FakeSafeBrowsing : public mojom::SafeBrowsing {
+ public:
+  FakeSafeBrowsing() : render_frame_id_(), load_flags_(-1), resource_type_() {}
+
+  void CreateCheckerAndCheck(int32_t render_frame_id,
+                             mojom::SafeBrowsingUrlCheckerRequest request,
+                             const GURL& url,
+                             int32_t load_flags,
+                             content::ResourceType resource_type,
+                             CreateCheckerAndCheckCallback callback) override {
+    render_frame_id_ = render_frame_id;
+    request_ = std::move(request);
+    url_ = url;
+    load_flags_ = load_flags;
+    resource_type_ = resource_type;
+    callback_ = std::move(callback);
+    run_loop_.Quit();
+  }
+
+  void RunUntilCalled() { run_loop_.Run(); }
+
+  int32_t render_frame_id_;
+  mojom::SafeBrowsingUrlCheckerRequest request_;
+  GURL url_;
+  int32_t load_flags_;
+  content::ResourceType resource_type_;
+  CreateCheckerAndCheckCallback callback_;
+  base::RunLoop run_loop_;
+};
+
+class FakeWebCallbacks
+    : public blink::WebCallbacks<void, const blink::WebString&> {
+ public:
+  enum Result { NOT_CALLED, SUCCESS, ERROR };
+
+  FakeWebCallbacks() : result_(NOT_CALLED) {}
+
+  void OnSuccess() override {
+    result_ = SUCCESS;
+    run_loop_.Quit();
+  }
+
+  void OnError(const blink::WebString& message) override {
+    result_ = ERROR;
+    message_ = message;
+    run_loop_.Quit();
+  }
+
+  void RunUntilCalled() { run_loop_.Run(); }
+
+  Result result_;
+  blink::WebString message_;
+  base::RunLoop run_loop_;
+};
+
+class WebSocketSBHandshakeThrottleTest : public ::testing::Test {
+ protected:
+  WebSocketSBHandshakeThrottleTest()
+      : safe_browsing_(new FakeSafeBrowsing()), mojo_binding_(safe_browsing_) {
+    mojom::SafeBrowsingPtr safe_browsing_ptr;
+    mojo_binding_.Bind(mojo::MakeRequest(&safe_browsing_ptr));
+    throttle_ = base::MakeUnique<WebSocketSBHandshakeThrottle>(
+        std::move(safe_browsing_ptr));
+  }
+
+  base::test::ScopedTaskEnvironment message_loop_;
+  FakeSafeBrowsing* safe_browsing_;

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

There seems no need to use a raw pointer and leak. You could use a
FakeSafeBrowsing instance instead, |mojo_binding_| doesn't take ownership.

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

Thanks! I was confused about the ownership semantics.

+  mojo::Binding<mojom::SafeBrowsing> mojo_binding_;
+  std::unique_ptr<WebSocketSBHandshakeThrottle> throttle_;
+  FakeWebCallbacks fake_callbacks_;
+};
+
+TEST_F(WebSocketSBHandshakeThrottleTest, Construction) {}
+
+TEST_F(WebSocketSBHandshakeThrottleTest, CheckArguments) {
+  throttle_->ThrottleHandshake(GURL(kTestUrl), nullptr, &fake_callbacks_);
+  safe_browsing_->RunUntilCalled();
+  // TODO(ricea): Find a way to create a WebLocalFrame in a unit test so that
+  // the code that looks up the render_frame_id can be tested.
+  EXPECT_EQ(MSG_ROUTING_NONE, safe_browsing_->render_frame_id_);
+  EXPECT_EQ(GURL(kTestUrl), safe_browsing_->url_);
+  EXPECT_EQ(0, safe_browsing_->load_flags_);
+  EXPECT_EQ(content::RESOURCE_TYPE_SUB_RESOURCE,
+            safe_browsing_->resource_type_);
+  EXPECT_TRUE(safe_browsing_->callback_);
+}
+
+TEST_F(WebSocketSBHandshakeThrottleTest, Safe) {
+  throttle_->ThrottleHandshake(GURL(kTestUrl), nullptr, &fake_callbacks_);
+  safe_browsing_->RunUntilCalled();
+  std::move(safe_browsing_->callback_).Run(true);
+  fake_callbacks_.RunUntilCalled();
+  EXPECT_EQ(FakeWebCallbacks::SUCCESS, fake_callbacks_.result_);
+}
+
+TEST_F(WebSocketSBHandshakeThrottleTest, Unsafe) {
+  throttle_->ThrottleHandshake(GURL(kTestUrl), nullptr, &fake_callbacks_);
+  safe_browsing_->RunUntilCalled();
+  std::move(safe_browsing_->callback_).Run(false);
+  fake_callbacks_.RunUntilCalled();
+  EXPECT_EQ(FakeWebCallbacks::ERROR, fake_callbacks_.result_);
+  EXPECT_EQ(
+      blink::WebString(
+          "WebSocket connection to wss://test/ failed safe browsing check"),
+      fake_callbacks_.message_);
+}
+
+}  // namespace
+
+}  // namespace safe_browsing