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..806768d3c7e52daec57227c5cfd85dcfacdf69d0
--- /dev/null
+++ b/components/safe_browsing/renderer/websocket_sb_handshake_throttle_unittest.cc
@@ -0,0 +1,135 @@
+// 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 { RESULT_NOT_CALLED, RESULT_SUCCESS, RESULT_ERROR };
+
+  FakeWebCallbacks() : result_(RESULT_NOT_CALLED) {}
+
+  void OnSuccess() override {
+    result_ = RESULT_SUCCESS;
+    run_loop_.Quit();
+  }
+
+  void OnError(const blink::WebString& message) override {
+    result_ = 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() : mojo_binding_(&safe_browsing_) {
+    mojo_binding_.Bind(mojo::MakeRequest(&safe_browsing_ptr_));
+    throttle_ = base::MakeUnique<WebSocketSBHandshakeThrottle>(
+        safe_browsing_ptr_.get());
+  }
+
+  base::test::ScopedTaskEnvironment message_loop_;
+  FakeSafeBrowsing safe_browsing_;
+  mojo::Binding<mojom::SafeBrowsing> mojo_binding_;
+  mojom::SafeBrowsingPtr safe_browsing_ptr_;
+  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::RESULT_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::RESULT_ERROR, fake_callbacks_.result_);
+  EXPECT_EQ(
+      blink::WebString(
+          "WebSocket connection to wss://test/ failed safe browsing check"),
+      fake_callbacks_.message_);
+}
+
+}  // namespace
+
+}  // namespace safe_browsing