[DNR] Implement storing/setting actions matched badge text

This CL implements the logic to store the number of actions matched for
an extension's ruleset and surface this count as a extension's badge
text if the extension has called setActionCountAsBadgeText(true).

Bug: 973211
Change-Id: Icc438d7be5f4e33bd8e0e77c4034bd2f69391a20
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1656970
Reviewed-by: Devlin <rdevlin.cronin@chromium.org>
Reviewed-by: Karan Bhatia <karandeepb@chromium.org>
Commit-Queue: Kelvin Jiang <kelvinjiang@chromium.org>
Cr-Commit-Position: refs/heads/master@{#686579}

chrome/browser/extensions/api/chrome_extensions_api_client.cc

@@ -16,6 +16,7 @@
 #include "chrome/browser/extensions/api/chrome_device_permissions_prompt.h"
 #include "chrome/browser/extensions/api/declarative_content/chrome_content_rules_registry.h"
 #include "chrome/browser/extensions/api/declarative_content/default_content_predicate_evaluators.h"
+#include "chrome/browser/extensions/api/extension_action/extension_action_api.h"
 #include "chrome/browser/extensions/api/feedback_private/chrome_feedback_private_delegate.h"
 #include "chrome/browser/extensions/api/file_system/chrome_file_system_delegate.h"
 #include "chrome/browser/extensions/api/management/chrome_management_api_delegate.h"
@@ -25,7 +26,10 @@
 #include "chrome/browser/extensions/api/storage/managed_value_store_cache.h"
 #include "chrome/browser/extensions/api/storage/sync_value_store_cache.h"
 #include "chrome/browser/extensions/chrome_extension_web_contents_observer.h"
+#include "chrome/browser/extensions/extension_action.h"
+#include "chrome/browser/extensions/extension_action_manager.h"
 #include "chrome/browser/extensions/extension_action_runner.h"
+#include "chrome/browser/extensions/extension_tab_util.h"
 #include "chrome/browser/extensions/system_display/display_info_provider.h"
 #include "chrome/browser/favicon/favicon_utils.h"
 #include "chrome/browser/guest_view/app_view/chrome_app_view_guest_delegate.h"
@@ -213,6 +217,30 @@ void ChromeExtensionsAPIClient::NotifyWebRequestWithheld(
   runner->OnWebRequestBlocked(extension);
 }
 
+void ChromeExtensionsAPIClient::UpdateActionCount(
+    content::BrowserContext* context,
+    const ExtensionId& extension_id,
+    int tab_id,
+    int action_count) {
+  const Extension* extension =
+      ExtensionRegistry::Get(context)->enabled_extensions().GetByID(
+          extension_id);
+  DCHECK(extension);
+
+  ExtensionAction* action =
+      ExtensionActionManager::Get(context)->GetExtensionAction(*extension);
+  DCHECK(action);
+
+  action->SetDNRActionCount(tab_id, action_count);
+  content::WebContents* tab_contents = nullptr;
+  if (ExtensionTabUtil::GetTabById(
+          tab_id, context, true /* include_incognito */, &tab_contents) &&
+      tab_contents) {
+    ExtensionActionAPI::Get(context)->NotifyChange(action, tab_contents,
+                                                   context);
+  }
+}
+
 AppViewGuestDelegate* ChromeExtensionsAPIClient::CreateAppViewGuestDelegate()
     const {
   return new ChromeAppViewGuestDelegate();
@@ -241,22 +269,20 @@ WebViewGuestDelegate* ChromeExtensionsAPIClient::CreateWebViewGuestDelegate(
   return new ChromeWebViewGuestDelegate(web_view_guest);
 }
 
-WebViewPermissionHelperDelegate* ChromeExtensionsAPIClient::
-    CreateWebViewPermissionHelperDelegate(
-        WebViewPermissionHelper* web_view_permission_helper) const {
+WebViewPermissionHelperDelegate*
+ChromeExtensionsAPIClient::CreateWebViewPermissionHelperDelegate(
+    WebViewPermissionHelper* web_view_permission_helper) const {
   return new ChromeWebViewPermissionHelperDelegate(web_view_permission_helper);
 }
 
 scoped_refptr<ContentRulesRegistry>
 ChromeExtensionsAPIClient::CreateContentRulesRegistry(
     content::BrowserContext* browser_context,
     RulesCacheDelegate* cache_delegate) const {
-  return scoped_refptr<ContentRulesRegistry>(
-      new ChromeContentRulesRegistry(
-          browser_context,
-          cache_delegate,
-          base::Bind(&CreateDefaultContentPredicateEvaluators,
-                     base::Unretained(browser_context))));
+  return scoped_refptr<ContentRulesRegistry>(new ChromeContentRulesRegistry(
+      browser_context, cache_delegate,
+      base::Bind(&CreateDefaultContentPredicateEvaluators,
+                 base::Unretained(browser_context))));
 }
 
 std::unique_ptr<DevicePermissionsPrompt>

chrome/browser/extensions/api/chrome_extensions_api_client.h

@@ -39,6 +39,10 @@ class ChromeExtensionsAPIClient : public ExtensionsAPIClient {
   void NotifyWebRequestWithheld(int render_process_id,
                                 int render_frame_id,
                                 const ExtensionId& extension_id) override;
+  void UpdateActionCount(content::BrowserContext* context,
+                         const ExtensionId& extension_id,
+                         int tab_id,
+                         int action_count) override;
   AppViewGuestDelegate* CreateAppViewGuestDelegate() const override;
   ExtensionOptionsGuestDelegate* CreateExtensionOptionsGuestDelegate(
       ExtensionOptionsGuest* guest) const override;

chrome/browser/extensions/api/declarative_net_request/declarative_net_request_browsertest.cc

@@ -29,16 +29,21 @@
 #include "base/threading/thread_restrictions.h"
 #include "base/values.h"
 #include "build/build_config.h"
+#include "chrome/browser/extensions/extension_action.h"
+#include "chrome/browser/extensions/extension_action_manager.h"
 #include "chrome/browser/extensions/extension_action_runner.h"
 #include "chrome/browser/extensions/extension_browsertest.h"
+#include "chrome/browser/extensions/extension_tab_util.h"
 #include "chrome/browser/extensions/extension_util.h"
 #include "chrome/browser/extensions/load_error_reporter.h"
 #include "chrome/browser/extensions/scripting_permissions_modifier.h"
 #include "chrome/browser/net/profile_network_context_service.h"
 #include "chrome/browser/net/profile_network_context_service_factory.h"
 #include "chrome/browser/profiles/profile.h"
+#include "chrome/browser/ui/browser_tabstrip.h"
 #include "chrome/browser/ui/tabs/tab_strip_model.h"
 #include "chrome/common/chrome_paths.h"
+#include "chrome/common/extensions/api/extension_action/action_info.h"
 #include "chrome/common/webui_url_constants.h"
 #include "chrome/test/base/ui_test_utils.h"
 #include "components/prefs/pref_service.h"
@@ -2636,6 +2641,263 @@ IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest, Redirect) {
   }
 }
 
+// Test that the badge text for an extension will update to reflect the number
+// of actions taken on requests matching the extension's ruleset.
+IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
+                       ActionsMatchedCountAsBadgeText) {
+  auto get_url_for_host = [this](std::string hostname) {
+    return embedded_test_server()->GetURL(hostname,
+                                          "/pages_with_script/index.html");
+  };
+
+  // This page simulates a user clicking on a link, so that the next page it
+  // navigates to has a Referrer header.
+  auto get_url_with_referrer = [this](std::string hostname) {
+    return embedded_test_server()->GetURL(hostname, "/simulate_click.html");
+  };
+
+  // Navigates frame with name |frame_name| to |url|.
+  auto navigate_frame = [this](const std::string& frame_name, const GURL& url) {
+    content::TestNavigationObserver navigation_observer(
+        web_contents(), 1 /*number_of_navigations*/);
+
+    // Before navigation, We set the referrer policy of the iframe to
+    // 'no-referrer' to prevent a referer header from being added for the iframe
+    // navigation.
+    ASSERT_TRUE(content::ExecuteScript(
+        GetMainFrame(),
+        base::StringPrintf(R"(
+          document.getElementsByName('%s')[0].referrerPolicy = 'no-referrer';
+          document.getElementsByName('%s')[0].src = '%s';)",
+                           frame_name.c_str(), frame_name.c_str(),
+                           url.spec().c_str())));
+    navigation_observer.Wait();
+  };
+
+  const std::string kFrameName1 = "frame1";
+  const GURL page_url = embedded_test_server()->GetURL(
+      "norulesmatched.com", "/page_with_two_frames.html");
+
+  struct {
+    std::string url_filter;
+    int id;
+    int priority;
+    std::string action_type;
+    base::Optional<std::string> redirect_url;
+    base::Optional<std::vector<std::string>> remove_headers_list;
+  } rules_data[] = {
+      {"abc.com", 1, 1, "block", base::nullopt, base::nullopt},
+      {"def.com", 2, 1, "redirect", "http://zzz.com", base::nullopt},
+      {"jkl.com", 3, 1, "removeHeaders", base::nullopt,
+       std::vector<std::string>({"referer"})},
+      {"abcd.com", 4, 1, "block", base::nullopt, base::nullopt},
+      {"abcd", 5, 1, "allow", base::nullopt, base::nullopt},
+  };
+
+  // Load the extension.
+  std::vector<TestRule> rules;
+  for (const auto& rule_data : rules_data) {
+    TestRule rule = CreateGenericRule();
+    rule.condition->url_filter = rule_data.url_filter;
+    rule.id = rule_data.id;
+    rule.priority = rule_data.priority;
+    rule.condition->resource_types =
+        std::vector<std::string>({"main_frame", "sub_frame"});
+    rule.action->type = rule_data.action_type;
+    rule.action->redirect.emplace();
+    rule.action->redirect->url = rule_data.redirect_url;
+    rule.action->remove_headers_list = rule_data.remove_headers_list;
+    rules.push_back(rule);
+  }
+
+  ASSERT_NO_FATAL_FAILURE(LoadExtensionWithRules(
+      rules, "test_extension", {URLPattern::kAllUrlsPattern}));
+
+  const Extension* dnr_extension = extension_registry()->GetExtensionById(
+      last_loaded_extension_id(), extensions::ExtensionRegistry::ENABLED);
+
+  ExtensionPrefs::Get(profile())->SetDNRUseActionCountAsBadgeText(
+      last_loaded_extension_id(), true);
+
+  ExtensionAction* action =
+      ExtensionActionManager::Get(web_contents()->GetBrowserContext())
+          ->GetExtensionAction(*dnr_extension);
+
+  struct {
+    std::string frame_hostname;
+    std::string expected_badge_text;
+    bool has_referrer_header;
+  } test_cases[] = {
+      // zzz.com does not match any rules, but we should still display 0 as the
+      // badge text as the preference is on.
+      {"zzz.com", "0", false},
+      // abc.com is blocked by a matching rule and should increment the badge
+      // text.
+      {"abc.com", "1", false},
+      // def.com is redirected by a matching rule and should increment the badge
+      // text.
+      {"def.com", "2", false},
+      // jkl.com matches with a removeHeaders rule, but has no headers.
+      // Therefore no action is taken and the badge text stays the same.
+      {"jkl.com", "2", false},
+      // jkl.com matches with a removeHeaders rule and has a referrer header.
+      // Therefore the badge text should be incremented.
+      {"jkl.com", "3", true},
+      // abcd.com matches both a block rule and an allow rule. Since the allow
+      // rule overrides the block rule, no action is taken and the badge text
+      // stays the same,
+      {"abcd.com", "3", false},
+  };
+
+  ui_test_utils::NavigateToURL(browser(), page_url);
+  ASSERT_TRUE(WasFrameWithScriptLoaded(GetMainFrame()));
+
+  // Verify that the badge text is 0 when navigation finishes.
+  int first_tab_id = ExtensionTabUtil::GetTabId(web_contents());
+  EXPECT_EQ("0", action->GetDisplayBadgeText(first_tab_id));
+
+  for (const auto& test_case : test_cases) {
+    GURL url = test_case.has_referrer_header
+                   ? get_url_with_referrer(test_case.frame_hostname)
+                   : get_url_for_host(test_case.frame_hostname);
+    SCOPED_TRACE(base::StringPrintf("Testing %s", url.spec().c_str()));
+
+    navigate_frame(kFrameName1, url);
+    EXPECT_EQ(test_case.expected_badge_text,
+              action->GetDisplayBadgeText(first_tab_id));
+  }
+
+  std::string first_tab_badge_text = action->GetDisplayBadgeText(first_tab_id);
+
+  const GURL second_tab_url = get_url_for_host("nomatch.com");
+  ui_test_utils::NavigateToURLWithDisposition(
+      browser(), second_tab_url, WindowOpenDisposition::NEW_FOREGROUND_TAB,
+      ui_test_utils::BROWSER_TEST_WAIT_FOR_NAVIGATION);
+
+  ASSERT_EQ(2, browser()->tab_strip_model()->count());
+  ASSERT_TRUE(browser()->tab_strip_model()->IsTabSelected(1));
+
+  int second_tab_id = ExtensionTabUtil::GetTabId(web_contents());
+  EXPECT_EQ("0", action->GetDisplayBadgeText(second_tab_id));
+
+  // Verify that the badge text for the first tab is unaffected.
+  EXPECT_EQ(first_tab_badge_text, action->GetDisplayBadgeText(first_tab_id));
+}
+
+// Test that the action matched badge text for an extension is visible in an
+// incognito context if the extension is incognito enabled.
+IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
+                       ActionsMatchedCountAsBadgeTextIncognito) {
+  TestRule rule = CreateGenericRule();
+  rule.condition->url_filter = "abc.com";
+  rule.id = kMinValidID;
+  rule.condition->resource_types = std::vector<std::string>({"main_frame"});
+  rule.action->type = "block";
+
+  std::vector<TestRule> rules({rule});
+  ASSERT_NO_FATAL_FAILURE(LoadExtensionWithRules(
+      {rules}, "test_extension", {URLPattern::kAllUrlsPattern}));
+
+  const ExtensionId extension_id = last_loaded_extension_id();
+  util::SetIsIncognitoEnabled(extension_id, profile(), true /*enabled*/);
+
+  ExtensionPrefs::Get(profile())->SetDNRUseActionCountAsBadgeText(extension_id,
+                                                                  true);
+
+  Browser* incognito_browser = CreateIncognitoBrowser();
+  ui_test_utils::NavigateToURL(incognito_browser, GURL("http://abc.com"));
+
+  content::WebContents* incognito_contents =
+      incognito_browser->tab_strip_model()->GetActiveWebContents();
+
+  const Extension* dnr_extension = extension_registry()->GetExtensionById(
+      extension_id, extensions::ExtensionRegistry::ENABLED);
+  ExtensionAction* incognito_action =
+      ExtensionActionManager::Get(incognito_contents->GetBrowserContext())
+          ->GetExtensionAction(*dnr_extension);
+
+  // TODO(crbug.com/992251): This should be a "1" after the main-frame
+  // navigation case is fixed.
+  EXPECT_EQ("0", incognito_action->GetDisplayBadgeText(
+                     ExtensionTabUtil::GetTabId(incognito_contents)));
+}
+
+// Test that the actions matched badge text for an extension will be reset
+// when a main-frame navigation finishes.
+// TODO(crbug.com/992251): Edit this test to add more main-frame cases.
+IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
+                       ActionsMatchedCountAsBadgeTextMainFrame) {
+  auto get_url_for_host = [this](std::string hostname) {
+    return embedded_test_server()->GetURL(hostname,
+                                          "/pages_with_script/index.html");
+  };
+
+  struct {
+    std::string url_filter;
+    int id;
+    int priority;
+    std::string action_type;
+    std::vector<std::string> resource_types;
+    base::Optional<std::string> redirect_url;
+  } rules_data[] = {
+      {"abc.com", 1, 1, "block", std::vector<std::string>({"script"}),
+       base::nullopt},
+      {"def.com", 2, 1, "redirect", std::vector<std::string>({"main_frame"}),
+       get_url_for_host("abc.com").spec()},
+  };
+
+  // Load the extension.
+  std::vector<TestRule> rules;
+  for (const auto& rule_data : rules_data) {
+    TestRule rule = CreateGenericRule();
+    rule.condition->url_filter = rule_data.url_filter;
+    rule.id = rule_data.id;
+    rule.priority = rule_data.priority;
+    rule.condition->resource_types = rule_data.resource_types;
+    rule.action->type = rule_data.action_type;
+    rule.action->redirect.emplace();
+    rule.action->redirect->url = rule_data.redirect_url;
+    rules.push_back(rule);
+  }
+
+  ASSERT_NO_FATAL_FAILURE(LoadExtensionWithRules(
+      rules, "test_extension", {URLPattern::kAllUrlsPattern}));
+
+  const Extension* dnr_extension = extension_registry()->GetExtensionById(
+      last_loaded_extension_id(), extensions::ExtensionRegistry::ENABLED);
+
+  ExtensionPrefs::Get(profile())->SetDNRUseActionCountAsBadgeText(
+      last_loaded_extension_id(), true);
+
+  ExtensionAction* action =
+      ExtensionActionManager::Get(web_contents()->GetBrowserContext())
+          ->GetExtensionAction(*dnr_extension);
+
+  struct {
+    std::string frame_hostname;
+    std::string expected_badge_text;
+  } test_cases[] = {
+      // The script on get_url_for_host("abc.com") matches with a rule and
+      // should increment the badge text.
+      {"abc.com", "1"},
+      // No rules match, so the badge text should be 0 once navigation finishes.
+      {"nomatch.com", "0"},
+      // The request to def.com will redirect to get_url_for_host("abc.com") and
+      // the script on abc.com should match with a rule.
+      {"def.com", "1"},
+  };
+
+  int first_tab_id = ExtensionTabUtil::GetTabId(web_contents());
+  for (const auto& test_case : test_cases) {
+    GURL url = get_url_for_host(test_case.frame_hostname);
+    SCOPED_TRACE(base::StringPrintf("Testing %s", url.spec().c_str()));
+
+    ui_test_utils::NavigateToURL(browser(), url);
+    EXPECT_EQ(test_case.expected_badge_text,
+              action->GetDisplayBadgeText(first_tab_id));
+  }
+}
+
 // Test fixture to verify that host permissions for the request url and the
 // request initiator are properly checked when redirecting requests. Loads an
 // example.com url with four sub-frames named frame_[1..4] from hosts