Enable websocket filtering via SubresourceFilter

chrome/browser/subresource_filter/subresource_filter_browsertest.cc

Index: chrome/browser/subresource_filter/subresource_filter_browsertest.cc
diff --git a/chrome/browser/subresource_filter/subresource_filter_browsertest.cc b/chrome/browser/subresource_filter/subresource_filter_browsertest.cc
index 9724eae20670359081a1762dfa4dd5a2dcbc01d0..ece0fc308a59a076ba9b3d9ee181117ee943e7df 100644
--- a/chrome/browser/subresource_filter/subresource_filter_browsertest.cc
+++ b/chrome/browser/subresource_filter/subresource_filter_browsertest.cc
@@ -14,6 +14,7 @@
 #include "base/macros.h"
 #include "base/path_service.h"
 #include "base/strings/string_piece.h"
+#include "base/strings/stringprintf.h"
 #include "base/test/histogram_tester.h"
 #include "chrome/browser/browser_process.h"
 #include "chrome/browser/metrics/subprocess_metrics_provider.h"
@@ -46,6 +47,8 @@
 #include "content/public/test/test_utils.h"
 #include "net/dns/mock_host_resolver.h"
 #include "net/test/embedded_test_server/embedded_test_server.h"
+#include "net/test/spawned_test_server/spawned_test_server.h"
+#include "net/test/test_data_directory.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
@@ -250,12 +253,27 @@ class SubresourceFilterBrowserTestImpl : public InProcessBrowserTest {
     host_resolver()->AddRule("*", "127.0.0.1");
     content::SetupCrossSiteRedirector(embedded_test_server());
     ASSERT_TRUE(embedded_test_server()->Start());
+
+    websocket_test_server_ = base::MakeUnique<net::SpawnedTestServer>(
+        net::SpawnedTestServer::TYPE_WS, net::SpawnedTestServer::kLocalhost,
+        net::GetWebSocketTestDataDirectory());
+    ASSERT_TRUE(websocket_test_server_->Start());
+  }
+
+  net::SpawnedTestServer* websocket_test_server() {
+    return websocket_test_server_.get();
   }
 
   GURL GetTestUrl(const std::string& path) {
     return embedded_test_server()->base_url().Resolve(path);
   }
 
+  GURL GetWebSocketUrl(const std::string& path) {
+    GURL::Replacements replacements;
+    replacements.SetSchemeStr("ws");
+    return websocket_test_server_->GetURL(path).ReplaceComponents(replacements);
+  }
+
   void ConfigureAsPhishingURL(const GURL& url) {
     fake_safe_browsing_database_->AddBlacklistedURL(
         url, safe_browsing::SB_THREAT_TYPE_URL_PHISHING);
@@ -337,6 +355,7 @@ class SubresourceFilterBrowserTestImpl : public InProcessBrowserTest {
   safe_browsing::TestSafeBrowsingServiceFactory test_safe_browsing_factory_;
   scoped_refptr<FakeSafeBrowsingDatabaseManager> fake_safe_browsing_database_;
 
+  std::unique_ptr<net::SpawnedTestServer> websocket_test_server_;
   std::unique_ptr<ScopedSubresourceFilterFeatureToggle> scoped_feature_toggle_;
   TestRulesetPublisher test_ruleset_publisher_;
   const bool measure_performance_;
@@ -370,6 +389,18 @@ class SubresourceFilterWhitelistSiteOnReloadBrowserTest
       : SubresourceFilterBrowserTestImpl(false, true) {}
 };
 
+enum WebSocketCreationPolicy {
+  IN_MAIN_FRAME,
+  IN_WORKER,
+};
+class SubresourceFilterWebSocketBrowserTest
+    : public SubresourceFilterBrowserTestImpl,
+      public ::testing::WithParamInterface<WebSocketCreationPolicy> {
+ public:
+  SubresourceFilterWebSocketBrowserTest()
+      : SubresourceFilterBrowserTestImpl(false, false) {}
+};
+
 // Tests -----------------------------------------------------------------------
 
 IN_PROC_BROWSER_TEST_F(SubresourceFilterBrowserTest, MainFrameActivation) {
@@ -640,6 +671,50 @@ IN_PROC_BROWSER_TEST_F(SubresourceFilterBrowserTest,
   ExpectParsedScriptElementLoadedStatusInFrames({"b", "d"}, {false, true});
 }
 
+IN_PROC_BROWSER_TEST_P(SubresourceFilterWebSocketBrowserTest, BlockWebSocket) {
+  GURL url(GetTestUrl(
+      base::StringPrintf("subresource_filter/page_with_websocket.html?%s",
+                         GetParam() == IN_WORKER ? "inWorker" : "")));
+  GURL websocket_url(GetWebSocketUrl("echo-with-no-extension"));
+  ConfigureAsPhishingURL(url);
+  ASSERT_NO_FATAL_FAILURE(
+      SetRulesetToDisallowURLsWithPathSuffix("echo-with-no-extension"));
+  ui_test_utils::NavigateToURL(browser(), url);
+
+  bool websocket_received_data = false;
+  EXPECT_TRUE(content::ExecuteScriptAndExtractBool(
+      browser()->tab_strip_model()->GetActiveWebContents(),
+      base::StringPrintf("connectWebSocket('%s');",
+                         websocket_url.spec().c_str()),
+      &websocket_received_data));
+  EXPECT_FALSE(websocket_received_data);
+}
+
+IN_PROC_BROWSER_TEST_P(SubresourceFilterWebSocketBrowserTest,
+                       DoNotBlockWebSocket) {
+  GURL url(GetTestUrl(
+      base::StringPrintf("subresource_filter/page_with_websocket.html?%s",
+                         GetParam() == IN_WORKER ? "inWorker" : "")));
+  GURL websocket_url(GetWebSocketUrl("echo-with-no-extension"));
+  ASSERT_NO_FATAL_FAILURE(
+      SetRulesetToDisallowURLsWithPathSuffix("echo-with-no-extension"));
+  ui_test_utils::NavigateToURL(browser(), url);
+
+  bool websocket_received_data = false;
+  EXPECT_TRUE(content::ExecuteScriptAndExtractBool(
+      browser()->tab_strip_model()->GetActiveWebContents(),
+      base::StringPrintf("connectWebSocket('%s');",
+                         websocket_url.spec().c_str()),
+      &websocket_received_data));
+  EXPECT_TRUE(websocket_received_data);
+}
+
+INSTANTIATE_TEST_CASE_P(
+    /* no prefix */,
+    SubresourceFilterWebSocketBrowserTest,
+    ::testing::Values(WebSocketCreationPolicy::IN_WORKER,
+                      WebSocketCreationPolicy::IN_MAIN_FRAME));
+
 // Tests checking how histograms are recorded. ---------------------------------
 
 namespace {