[DNR] Properly attribute actions for main-frame navigation requests

This CL leverages crrev.com/c/1923611 by using the navigation ID within
WebRequestInfo to correlate a main-frame navigation request originating
from the previously active WebContents to the new WebContents loaded
from that request.

This is accomplished by adding the actions matched with a request to a
pending set of actions for a given navigation ID if the request is
associated main-frame navigation. Once the navigation completes, the
actions matched for the navigation are added to the badge text and
cleared from the pending set.

Bug: 992251
Change-Id: Iac7b3502a3598c9d06ab3261ea33f30d2cb20ee7
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1956110
Commit-Queue: Kelvin Jiang <kelvinjiang@chromium.org>
Reviewed-by: Karan Bhatia <karandeepb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#726572}

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

@@ -3096,15 +3096,12 @@ IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
       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(
+  EXPECT_EQ("1", 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) {
@@ -3112,6 +3109,10 @@ IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
                                           "/pages_with_script/index.html");
   };
 
+  auto get_set_cookie_url = [this](std::string hostname) {
+    return embedded_test_server()->GetURL(hostname, "/set-cookie?a=b");
+  };
+
   struct {
     std::string url_filter;
     int id;
@@ -3119,11 +3120,23 @@ IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
     std::string action_type;
     std::vector<std::string> resource_types;
     base::Optional<std::string> redirect_url;
+    base::Optional<std::vector<std::string>> remove_headers_list;
   } 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()},
+       base::nullopt, base::nullopt},
+      {"||def.com", 2, 1, "redirect", std::vector<std::string>({"main_frame"}),
+       get_url_for_host("abc.com").spec(), base::nullopt},
+      {"gotodef.com", 3, 1, "redirect",
+       std::vector<std::string>({"main_frame"}),
+       get_url_for_host("def.com").spec(), base::nullopt},
+      {"ghi.com", 4, 1, "block", std::vector<std::string>({"main_frame"}),
+       base::nullopt, base::nullopt},
+      {"gotosetcookie.com", 5, 1, "redirect",
+       std::vector<std::string>({"main_frame"}),
+       get_set_cookie_url("setcookie.com").spec(), base::nullopt},
+      {"setcookie.com", 6, 1, "removeHeaders",
+       std::vector<std::string>({"main_frame"}), base::nullopt,
+       std::vector<std::string>({"setCookie"})},
   };
 
   // Load the extension.
@@ -3137,6 +3150,7 @@ IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
     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);
   }
 
@@ -3157,6 +3171,9 @@ IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
     std::string frame_hostname;
     std::string expected_badge_text;
   } test_cases[] = {
+      // The request to ghi.com should be blocked, so the badge text should be 1
+      // once navigation finishes.
+      {"ghi.com", "1"},
       // The script on get_url_for_host("abc.com") matches with a rule and
       // should increment the badge text.
       {"abc.com", "1"},
@@ -3164,7 +3181,15 @@ IN_PROC_BROWSER_TEST_P(DeclarativeNetRequestBrowserTest,
       {"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"},
+      {"def.com", "2"},
+      // Same as the above test, except with an additional redirect from
+      // gotodef.com to def.com caused by a rule match. Therefore the badge text
+      // should be 3.
+      {"gotodef.com", "3"},
+      // The request to gotosetcookie.com will match with a rule and redirect to
+      // setcookie.com. The Set-Cookie header on setcookie.com will also match
+      // with a rule and get removed. Therefore the badge text should be 2.
+      {"gotosetcookie.com", "2"},
   };
 
   int first_tab_id = ExtensionTabUtil::GetTabId(web_contents());

chrome/browser/extensions/extension_action_runner.cc

@@ -507,12 +507,22 @@ bool ExtensionActionRunner::OnMessageReceived(
 
 void ExtensionActionRunner::DidFinishNavigation(
     content::NavigationHandle* navigation_handle) {
-  if (!navigation_handle->IsInMainFrame() ||
-      !navigation_handle->HasCommitted() ||
-      navigation_handle->IsSameDocument()) {
-    return;
+  declarative_net_request::RulesMonitorService* rules_monitor_service =
+      declarative_net_request::RulesMonitorService::Get(browser_context_);
+
+  const bool is_main_frame = navigation_handle->IsInMainFrame();
+  const bool has_committed = navigation_handle->HasCommitted();
+
+  if (is_main_frame && !has_committed && rules_monitor_service) {
+    // Clean up any pending actions recorded in the action tracker for this
+    // navigation.
+    rules_monitor_service->action_tracker().ClearPendingNavigation(
+        navigation_handle->GetNavigationId());
   }
 
+  if (!is_main_frame || !has_committed || navigation_handle->IsSameDocument())
+    return;
+
   LogUMA();
   num_page_requests_ = 0;
   permitted_extensions_.clear();
@@ -526,17 +536,13 @@ void ExtensionActionRunner::DidFinishNavigation(
   // run".
   ExtensionActionAPI::Get(browser_context_)
       ->ClearAllValuesForTab(web_contents());
-
-  declarative_net_request::RulesMonitorService* rules_monitor_service =
-      declarative_net_request::RulesMonitorService::Get(browser_context_);
-
   // |rules_monitor_service| can be null for some unit tests.
   if (rules_monitor_service) {
+    int tab_id = ExtensionTabUtil::GetTabId(web_contents());
     declarative_net_request::ActionTracker& action_tracker =
         rules_monitor_service->action_tracker();
-
-    int tab_id = ExtensionTabUtil::GetTabId(web_contents());
-    action_tracker.ResetActionCountForTab(tab_id);
+    action_tracker.ResetActionCountForTab(tab_id,
+                                          navigation_handle->GetNavigationId());
   }
 }
 

extensions/browser/api/declarative_net_request/action_tracker.cc

@@ -33,6 +33,7 @@ ActionTracker::ActionTracker(content::BrowserContext* browser_context)
 
 ActionTracker::~ActionTracker() {
   DCHECK(actions_matched_.empty());
+  DCHECK(pending_navigation_actions_.empty());
 }
 
 void ActionTracker::OnRuleMatched(const RequestAction& request_action,
@@ -50,9 +51,20 @@ void ActionTracker::OnRuleMatched(const RequestAction& request_action,
   }
 
   const ExtensionId& extension_id = request_action.extension_id;
-  ExtensionTabIdKey key(extension_id, tab_id);
 
-  size_t action_count = ++actions_matched_[std::move(key)].action_count;
+  // Increment action count for |pending_navigation_actions_| if the request is
+  // a main-frame navigation request.
+  if (request_info.is_navigation_request &&
+      request_info.type == content::ResourceType::kMainFrame) {
+    DCHECK(request_info.navigation_id);
+    pending_navigation_actions_[{extension_id, *request_info.navigation_id}]
+        .action_count++;
+    return;
+  }
+
+  // Otherwise, increment action count for |actions_matched_| and update the
+  // badge text for the current tab.
+  size_t action_count = ++actions_matched_[{extension_id, tab_id}].action_count;
   if (extension_prefs_->GetDNRUseActionCountAsBadgeText(extension_id)) {
     DCHECK(ExtensionsAPIClient::Get());
     ExtensionsAPIClient::Get()->UpdateActionCount(
@@ -72,31 +84,44 @@ void ActionTracker::OnPreferenceEnabled(const ExtensionId& extension_id) const {
       continue;
 
     ExtensionsAPIClient::Get()->UpdateActionCount(
-        browser_context_, extension_id, key.tab_id, value.action_count,
-        true /* clear_badge_text */);
+        browser_context_, extension_id, key.secondary_id /* tab_id */,
+        value.action_count, true /* clear_badge_text */);
   }
 }
 
 void ActionTracker::ClearExtensionData(const ExtensionId& extension_id) {
-  auto compare_by_extension_id =
-      [&extension_id](
-          const std::pair<const ExtensionTabIdKey, TrackedInfo>& it) {
-        return it.first.extension_id == extension_id;
-      };
+  auto compare_by_extension_id = [&extension_id](const auto& it) {
+    return it.first.extension_id == extension_id;
+  };
 
   base::EraseIf(actions_matched_, compare_by_extension_id);
+  base::EraseIf(pending_navigation_actions_, compare_by_extension_id);
 }
 
 void ActionTracker::ClearTabData(int tab_id) {
   auto compare_by_tab_id =
       [&tab_id](const std::pair<const ExtensionTabIdKey, TrackedInfo>& it) {
-        return it.first.tab_id == tab_id;
+        return it.first.secondary_id == tab_id;
       };
 
   base::EraseIf(actions_matched_, compare_by_tab_id);
 }
 
-void ActionTracker::ResetActionCountForTab(int tab_id) {
+void ActionTracker::ClearPendingNavigation(int64_t navigation_id) {
+  RulesMonitorService* rules_monitor_service =
+      RulesMonitorService::Get(browser_context_);
+
+  DCHECK(rules_monitor_service);
+  auto compare_by_navigation_id =
+      [navigation_id](
+          const std::pair<const ExtensionNavigationIdKey, TrackedInfo>& it) {
+        return it.first.secondary_id == navigation_id;
+      };
+
+  base::EraseIf(pending_navigation_actions_, compare_by_navigation_id);
+}
+
+void ActionTracker::ResetActionCountForTab(int tab_id, int64_t navigation_id) {
   DCHECK_NE(tab_id, extension_misc::kUnknownTabId);
 
   RulesMonitorService* rules_monitor_service =
@@ -105,32 +130,51 @@ void ActionTracker::ResetActionCountForTab(int tab_id) {
   DCHECK(rules_monitor_service);
   for (const auto& extension_id :
        rules_monitor_service->extensions_with_rulesets()) {
-    ExtensionTabIdKey key(extension_id, tab_id);
-    actions_matched_[std::move(key)].action_count = 0;
+    ExtensionNavigationIdKey navigation_key(extension_id, navigation_id);
+
+    size_t actions_matched_for_navigation = 0;
+    auto iter = pending_navigation_actions_.find(navigation_key);
+    if (iter != pending_navigation_actions_.end()) {
+      actions_matched_for_navigation = iter->second.action_count;
+      pending_navigation_actions_.erase(iter);
+    }
+
+    actions_matched_[{extension_id, tab_id}].action_count =
+        actions_matched_for_navigation;
 
     if (extension_prefs_->GetDNRUseActionCountAsBadgeText(extension_id)) {
       DCHECK(ExtensionsAPIClient::Get());
       ExtensionsAPIClient::Get()->UpdateActionCount(
-          browser_context_, extension_id, tab_id, 0 /* action_count */,
-          false /* clear_badge_text */);
+          browser_context_, extension_id, tab_id,
+          actions_matched_for_navigation, false /* clear_badge_text */);
     }
   }
-}
-
-ActionTracker::ExtensionTabIdKey::ExtensionTabIdKey(ExtensionId extension_id,
-                                                    int tab_id)
-    : extension_id(std::move(extension_id)), tab_id(tab_id) {}
 
-ActionTracker::ExtensionTabIdKey::ExtensionTabIdKey(
-    ActionTracker::ExtensionTabIdKey&&) = default;
-
-ActionTracker::ExtensionTabIdKey& ActionTracker::ExtensionTabIdKey::operator=(
-    ActionTracker::ExtensionTabIdKey&&) = default;
+  // Double check to make sure the pending counts for |navigation_id| are really
+  // cleared from |pending_navigation_actions_|.
+  ClearPendingNavigation(navigation_id);
+}
 
-bool ActionTracker::ExtensionTabIdKey::operator<(
-    const ExtensionTabIdKey& other) const {
-  return std::tie(tab_id, extension_id) <
-         std::tie(other.tab_id, other.extension_id);
+template <typename T>
+ActionTracker::TrackedInfoContextKey<T>::TrackedInfoContextKey(
+    ExtensionId extension_id,
+    T secondary_id)
+    : extension_id(std::move(extension_id)), secondary_id(secondary_id) {}
+
+template <typename T>
+ActionTracker::TrackedInfoContextKey<T>::TrackedInfoContextKey(
+    ActionTracker::TrackedInfoContextKey<T>&&) = default;
+
+template <typename T>
+ActionTracker::TrackedInfoContextKey<T>&
+ActionTracker::TrackedInfoContextKey<T>::operator=(
+    ActionTracker::TrackedInfoContextKey<T>&&) = default;
+
+template <typename T>
+bool ActionTracker::TrackedInfoContextKey<T>::operator<(
+    const TrackedInfoContextKey<T>& other) const {
+  return std::tie(secondary_id, extension_id) <
+         std::tie(other.secondary_id, other.extension_id);
 }
 
 void ActionTracker::DispatchOnRuleMatchedDebugIfNeeded(

extensions/browser/api/declarative_net_request/action_tracker.h

@@ -47,27 +47,36 @@ class ActionTracker {
   // Called when the tab has been closed.
   void ClearTabData(int tab_id);
 
-  // Sets the action count for every extension for the specified |tab_id| to 0
-  // and notifies the extension action to set the badge text to 0 for that tab.
-  // Called when the a main-frame navigation to a different document finishes on
-  // the tab.
-  void ResetActionCountForTab(int tab_id);
+  // Clears the pending action count for every extension in
+  // |pending_navigation_actions_| for the specified |navigation_id|.
+  void ClearPendingNavigation(int64_t navigation_id);
+
+  // Called when a main-frame navigation to a different document commits.
+  // Updates the badge count for all extensions for the given |tab_id|.
+  void ResetActionCountForTab(int tab_id, int64_t navigation_id);
 
  private:
-  struct ExtensionTabIdKey {
-    ExtensionTabIdKey(ExtensionId extension_id, int tab_id);
-    ExtensionTabIdKey(const ExtensionTabIdKey& other) = delete;
-    ExtensionTabIdKey& operator=(const ExtensionTabIdKey& other) = delete;
-    ExtensionTabIdKey(ExtensionTabIdKey&&);
-    ExtensionTabIdKey& operator=(ExtensionTabIdKey&&);
+  // Template key type used for TrackedInfo, specified by an extension_id and
+  // another ID.
+  template <typename T>
+  struct TrackedInfoContextKey {
+    TrackedInfoContextKey(ExtensionId extension_id, T secondary_id);
+    TrackedInfoContextKey(const TrackedInfoContextKey& other) = delete;
+    TrackedInfoContextKey& operator=(const TrackedInfoContextKey& other) =
+        delete;
+    TrackedInfoContextKey(TrackedInfoContextKey&&);
+    TrackedInfoContextKey& operator=(TrackedInfoContextKey&&);
 
     ExtensionId extension_id;
-    int tab_id;
+    T secondary_id;
 
-    bool operator<(const ExtensionTabIdKey& other) const;
+    bool operator<(const TrackedInfoContextKey& other) const;
   };
 
-  // Info tracked for each ExtensionTabIdKey.
+  using ExtensionTabIdKey = TrackedInfoContextKey<int>;
+  using ExtensionNavigationIdKey = TrackedInfoContextKey<int64_t>;
+
+  // Info tracked for each ExtensionTabIdKey or ExtensionNavigationIdKey.
   struct TrackedInfo {
     size_t action_count = 0;
   };
@@ -82,6 +91,12 @@ class ActionTracker {
   // the extension and tab specified.
   std::map<ExtensionTabIdKey, TrackedInfo> actions_matched_;
 
+  // Maps a pair of (extension ID, navigation ID) to the number of actions
+  // matched for the main-frame request associated with the navigation ID in the
+  // key. These actions are added to |actions_matched_| once the navigation
+  // commits.
+  std::map<ExtensionNavigationIdKey, TrackedInfo> pending_navigation_actions_;
+
   content::BrowserContext* browser_context_;
 
   ExtensionPrefs* extension_prefs_;