SafeBrowsing support for WebSocket (post-network-servicification)

chrome/browser/safe_browsing/safe_browsing_service_browsertest.cc

 // Copyright (c) 2012 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.
 //
 // This test creates a safebrowsing service using test safebrowsing database
 // and a test protocol manager. It is used to test logics in safebrowsing
 // service.
 
 #include "chrome/browser/safe_browsing/safe_browsing_service.h"
 
@@ skipping 59 matching lines @@
 #include "crypto/sha2.h"
 #include "net/cookies/cookie_store.h"
 #include "net/cookies/cookie_util.h"
 #include "net/test/embedded_test_server/embedded_test_server.h"
 #include "net/test/embedded_test_server/http_request.h"
 #include "net/test/embedded_test_server/http_response.h"
 #include "sql/connection.h"
 #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"
 #include "chromeos/chromeos_switches.h"
 #endif
 
 #if !defined(SAFE_BROWSING_DB_LOCAL)
 #error This test requires SAFE_BROWSING_DB_LOCAL.
 #endif
 
@@ skipping 10 matching lines @@
 
 const char kBlacklistResource[] = "/blacklisted/script.js";
 const char kEmptyPage[] = "/empty.html";
 const char kMaliciousResource[] = "/malware/script.js";
 const char kMalwareFile[] = "/downloads/dangerous/dangerous.exe";
 const char kMalwarePage[] = "/safe_browsing/malware.html";
 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";
 
 class NeverCompletingHttpResponse : public net::test_server::HttpResponse {
  public:
   ~NeverCompletingHttpResponse() override {}
 
   void SendResponse(
       const net::test_server::SendBytesCallback& send,
       const net::test_server::SendCompleteCallback& done) override {
@@ skipping 1349 matching lines @@
   EXPECT_TRUE(csd_service->enabled());
 
   // Delete the Profile. SBS stops again.
   pref_service2 = nullptr;
   profile2.reset();
   EXPECT_FALSE(sb_service->enabled_by_prefs());
   WaitForIOAndCheckEnabled(sb_service, false);
   EXPECT_FALSE(csd_service->enabled());
 }
 
+IN_PROC_BROWSER_TEST_F(SafeBrowsingServiceTest, MalwareWebSocketBlocked) {
+  // This test currently only passes when the network service is enabled.
+  if (base::CommandLine::ForCurrentProcess()->HasSwitch(

[yhirano, 2017/06/21 09:38:38]

I prefer

if (!HasSwitch())
  return;

[Adam Rice, 2017/06/22 12:35:04]

Done.

+          ::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 {
     // Browser should be fully torn down by now, so we can safely check these
     // counters.
     EXPECT_EQ(1, TestProtocolManager::create_count());
     EXPECT_EQ(1, TestProtocolManager::delete_count());
 
     SafeBrowsingServiceTest::TearDown();
   }
@@ skipping 867 matching lines @@
 }
 
 INSTANTIATE_TEST_CASE_P(
     MaybeSetMetadata,
     V4SafeBrowsingServiceMetadataTest,
     testing::Values(ThreatPatternType::NONE,
                     ThreatPatternType::MALWARE_LANDING,
                     ThreatPatternType::MALWARE_DISTRIBUTION));
 
 }  // namespace safe_browsing