Merge pull request #12597 from brave/issues/21661

Don't impose WebSockets limits on Extensions

browser/brave_shields/websockets_pool_limit_browsertest.cc

@@ -15,12 +15,19 @@
 #include "content/public/test/browser_test.h"
 #include "content/public/test/browser_test_utils.h"
 #include "content/public/test/content_mock_cert_verifier.h"
+#include "extensions/buildflags/buildflags.h"
 #include "net/dns/mock_host_resolver.h"
 #include "net/test/spawned_test_server/spawned_test_server.h"
 #include "net/test/test_data_directory.h"
 #include "third_party/blink/public/common/features.h"
 #include "url/gurl.h"
 
+#if BUILDFLAG(ENABLE_EXTENSIONS)
+#include "chrome/browser/extensions/chrome_test_extension_loader.h"
+#include "extensions/common/extension.h"
+#include "extensions/test/test_extension_dir.h"
+#endif  // BUILDFLAG(ENABLE_EXTENSIONS)
+
 namespace {
 
 const int kWebSocketsPoolLimit = 10;
@@ -302,6 +309,31 @@ IN_PROC_BROWSER_TEST_F(WebSocketsPoolLimitBrowserTest,
   OpenWebSockets(a_com_rfh, kWsOpenInSwScript, kWebSocketsPoolLimit + 5);
 }
 
+#if BUILDFLAG(ENABLE_EXTENSIONS)
+IN_PROC_BROWSER_TEST_F(WebSocketsPoolLimitBrowserTest,
+                       PoolIsNotLimitedForExtensions) {
+  extensions::TestExtensionDir test_extension_dir;
+  test_extension_dir.WriteManifest(R"({
+    "name": "Test",
+    "manifest_version": 2,
+    "version": "0.1",
+    "permissions": ["webRequest", "webRequestBlocking", "*://a.com/*"],
+    "content_security_policy": "script-src 'self' 'unsafe-eval'; object-src 'self'"
+  })");
+  test_extension_dir.WriteFile(FILE_PATH_LITERAL("empty.html"), "");
+
+  extensions::ChromeTestExtensionLoader extension_loader(browser()->profile());
+  scoped_refptr<const extensions::Extension> extension =
+      extension_loader.LoadExtension(test_extension_dir.UnpackedPath());
+  const GURL url = extension->GetResourceURL("/empty.html");
+  auto* extension_rfh = ui_test_utils::NavigateToURLWithDisposition(
+      browser(), url, WindowOpenDisposition::NEW_FOREGROUND_TAB,
+      ui_test_utils::BROWSER_TEST_WAIT_FOR_LOAD_STOP);
+  ASSERT_TRUE(extension_rfh);
+  OpenWebSockets(extension_rfh, kWsOpenScript, kWebSocketsPoolLimit + 5);
+}
+#endif  // BUILDFLAG(ENABLE_EXTENSIONS)
+
 class WebSocketsPoolLimitDisabledBrowserTest
     : public WebSocketsPoolLimitBrowserTest {
  public:

chromium_src/third_party/blink/renderer/modules/websockets/websocket_channel_impl.cc

@@ -6,6 +6,7 @@
 #include "third_party/blink/renderer/modules/websockets/websocket_channel_impl.h"
 
 #include "third_party/blink/public/common/features.h"
+#include "third_party/blink/public/common/scheme_registry.h"
 #include "third_party/blink/public/platform/web_content_settings_client.h"
 
 #define WebSocketChannelImpl WebSocketChannelImpl_ChromiumImpl
@@ -39,7 +40,10 @@ bool WebSocketChannelImpl::ShouldDisallowConnection(const KURL& url) {
   if (base::FeatureList::IsEnabled(blink::features::kRestrictWebSocketsPool)) {
     if (blink::WebContentSettingsClient* settings =
             brave::GetContentSettingsClientFor(execution_context_)) {
-      if (settings->GetBraveFarblingLevel() != BraveFarblingLevel::OFF) {
+      const bool is_extension = CommonSchemeRegistry::IsExtensionScheme(
+          execution_context_->GetSecurityOrigin()->Protocol().Ascii());
+      if (!is_extension &&
+          settings->GetBraveFarblingLevel() != BraveFarblingLevel::OFF) {
         websocket_in_use_tracker_ =
             ResourcePoolLimiter::GetInstance().IssueResourceInUseTracker(
                 execution_context_,

chromium_src/third_party/blink/renderer/modules/websockets/websocket_channel_impl.h

@@ -30,7 +30,7 @@ using WebSocketChannelImpl_BraveImpl = WebSocketChannelImpl;
 
 namespace blink {
 
-class MODULES_EXPORT WebSocketChannelImpl
+class MODULES_EXPORT WebSocketChannelImpl final
     : public WebSocketChannelImpl_ChromiumImpl {
  public:
   using WebSocketChannelImpl_ChromiumImpl::WebSocketChannelImpl_ChromiumImpl;

third_party/blink/renderer/core/resource_pool_limiter/resource_pool_limiter.cc

@@ -58,8 +58,8 @@ int GetResourceLimit(ResourcePoolLimiter::ResourceType resource_type) {
 }  // namespace
 
 ResourcePoolLimiter::ResourceInUseTracker::ResourceInUseTracker(
-    String resource_id_in_use)
-    : resource_id_in_use_(std::move(resource_id_in_use)) {}
+    String resource_id)
+    : resource_id_(std::move(resource_id)) {}
 
 ResourcePoolLimiter::ResourceInUseTracker::~ResourceInUseTracker() {
   ResourcePoolLimiter::GetInstance().DropResourceInUse(this);
@@ -87,13 +87,13 @@ ResourcePoolLimiter::IssueResourceInUseTracker(
 
   MutexLocker locker(resources_in_use_lock_);
   // `insert` doesn't change the value if it already exists.
-  int& resource_in_use_value =
+  int& resource_in_use_count =
       resources_in_use_.insert(resource_id, 0).stored_value->value;
-  if (resource_in_use_value >= GetResourceLimit(resource_type)) {
+  if (resource_in_use_count >= GetResourceLimit(resource_type)) {
     return nullptr;
   }
 
-  ++resource_in_use_value;
+  ++resource_in_use_count;
   return std::make_unique<ResourceInUseTracker>(resource_id.IsolatedCopy());
 }
 
@@ -101,7 +101,7 @@ void ResourcePoolLimiter::DropResourceInUse(
     const ResourceInUseTracker* resource_in_use_tracker) {
   MutexLocker locker(resources_in_use_lock_);
   auto resource_in_use_it =
-      resources_in_use_.find(resource_in_use_tracker->resource_id_in_use());
+      resources_in_use_.find(resource_in_use_tracker->resource_id());
   DCHECK(resource_in_use_it != resources_in_use_.end());
   if (--resource_in_use_it->value == 0) {
     resources_in_use_.erase(resource_in_use_it);

third_party/blink/renderer/core/resource_pool_limiter/resource_pool_limiter.h

@@ -27,13 +27,13 @@ class CORE_EXPORT ResourcePoolLimiter {
 
   class CORE_EXPORT ResourceInUseTracker {
    public:
-    explicit ResourceInUseTracker(String resource_id_in_use);
+    explicit ResourceInUseTracker(String resource_id);
     ~ResourceInUseTracker();
 
-    const String& resource_id_in_use() const { return resource_id_in_use_; }
+    const String& resource_id() const { return resource_id_; }
 
    private:
-    String resource_id_in_use_;
+    String resource_id_;
   };
 
   static ResourcePoolLimiter& GetInstance();