SafeBrowsing support for WebSocket (post-network-servicification)

chrome/browser/safe_browsing/safe_browsing_service_browsertest.cc

Index: chrome/browser/safe_browsing/safe_browsing_service_browsertest.cc
diff --git a/chrome/browser/safe_browsing/safe_browsing_service_browsertest.cc b/chrome/browser/safe_browsing/safe_browsing_service_browsertest.cc
index ed2c88a1c3549960a4c6dfd6e8d894faca3cb5b9..d0e5184287391be36eebf001bd59cf7aeb0d4401 100644
--- a/chrome/browser/safe_browsing/safe_browsing_service_browsertest.cc
+++ b/chrome/browser/safe_browsing/safe_browsing_service_browsertest.cc
@@ -77,6 +77,7 @@
 #include "sql/statement.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "url/gurl.h"
+#include "url/url_canon.h"
 
 #if defined(OS_CHROMEOS)
 #include "chrome/browser/chromeos/profiles/profile_helper.h"
@@ -107,6 +108,7 @@ const char kMalwareDelayedLoadsPage[] =
     "/safe_browsing/malware_delayed_loads.html";
 const char kMalwareIFrame[] = "/safe_browsing/malware_iframe.html";
 const char kMalwareImg[] = "/safe_browsing/malware_image.png";
+const char kMalwareWebSocketPage[] = "/safe_browsing/malware_websocket.html";
 const char kNeverCompletesPath[] = "/never_completes";
 const char kPrefetchMalwarePage[] = "/safe_browsing/prefetch_malware.html";
 
@@ -1476,6 +1478,62 @@ IN_PROC_BROWSER_TEST_F(SafeBrowsingServiceTest, StartAndStop) {
   EXPECT_FALSE(csd_service->enabled());
 }
 
+IN_PROC_BROWSER_TEST_F(SafeBrowsingServiceTest, MalwareWebSocketBlocked) {

[vakh (use Gerrit instead), 2017/06/20 21:40:27]

Please add the corresponding tests for V4SafeBrowsingServiceTest class also.
We're going to deprecate SafeBrowsingServiceTest class very very soon.

[Adam Rice, 2017/06/21 09:18:47]

Done. They're almost exact copies, but if the old tests are going to be deleted
I don't think that matters.

+  // This test currently only passes when the network service is enabled.
+  if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+          ::switches::kEnableNetworkService)) {
+    GURL main_url = embedded_test_server()->GetURL(kMalwareWebSocketPage);
+    // This constructs the URL with the same logic as malware_websocket.html.
+    GURL resolved = main_url.Resolve("/safe_browsing/malware-ws");
+    GURL::Replacements replace_scheme;
+    replace_scheme.SetScheme("ws", url::Component(0, strlen("ws")));
+    GURL websocket_url = resolved.ReplaceComponents(replace_scheme);
+
+    // Add the WebSocket url as malware.
+    SBFullHashResult uws_full_hash;
+    GenUrlFullHashResult(websocket_url, MALWARE, &uws_full_hash);
+    SetupResponseForUrl(websocket_url, uws_full_hash);
+
+    // Brute force method for waiting for the interstitial to be displayed.
+    content::WindowedNotificationObserver load_stop_observer(
+        content::NOTIFICATION_ALL,
+        base::Bind(
+            [](SafeBrowsingServiceTest* self,
+               const content::NotificationSource& source,
+               const content::NotificationDetails& details) {
+              return self->ShowingInterstitialPage();
+            },
+            base::Unretained(this)));
+
+    EXPECT_CALL(observer_,
+                OnSafeBrowsingHit(IsUnsafeResourceFor(websocket_url)));
+    ui_test_utils::NavigateToURL(browser(), main_url);
+
+    // If the interstitial fails to be displayed, the test will hang here.
+    load_stop_observer.Wait();
+
+    EXPECT_TRUE(ShowingInterstitialPage());
+    EXPECT_TRUE(got_hit_report());
+  }
+}
+
+IN_PROC_BROWSER_TEST_F(SafeBrowsingServiceTest, UnknownWebSocketNotBlocked) {
+  GURL main_url = embedded_test_server()->GetURL(kMalwareWebSocketPage);
+
+  auto expected_title = base::ASCIIToUTF16("COMPLETED");
+  content::TitleWatcher title_watcher(
+      browser()->tab_strip_model()->GetActiveWebContents(), expected_title);
+
+  // Load the parent page without marking the WebSocket as malware.
+  ui_test_utils::NavigateToURL(browser(), main_url);
+
+  // Wait for the WebSocket connection attempt to complete.
+  auto new_title = title_watcher.WaitAndGetTitle();
+  EXPECT_EQ(expected_title, new_title);
+  EXPECT_FALSE(ShowingInterstitialPage());
+  EXPECT_FALSE(got_hit_report());
+}
+
 class SafeBrowsingServiceShutdownTest : public SafeBrowsingServiceTest {
  public:
   void TearDown() override {