extensions/browser/api/declarative_net_request/action_tracker.h

// Copyright 2019 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.

#ifndef EXTENSIONS_BROWSER_API_DECLARATIVE_NET_REQUEST_ACTION_TRACKER_H_
#define EXTENSIONS_BROWSER_API_DECLARATIVE_NET_REQUEST_ACTION_TRACKER_H_

#include <list>
#include <map>
#include <vector>

#include "base/memory/raw_ptr.h"
#include "base/time/time.h"
#include "base/timer/timer.h"
#include "extensions/common/api/declarative_net_request.h"
#include "extensions/common/api/declarative_net_request/constants.h"
#include "extensions/common/extension_id.h"
#include "third_party/abseil-cpp/absl/types/optional.h"

namespace base {
class Clock;
class RetainingOneShotTimer;
}

namespace content {
class BrowserContext;
}

namespace extensions {

class Extension;
class ExtensionPrefs;
struct WebRequestInfo;

namespace declarative_net_request {
struct RequestAction;

class ActionTracker {
 public:
  // The lifespan of matched rules in |rules_matched_| not associated with an
  // active document,
  static constexpr base::TimeDelta kNonActiveTabRuleLifespan = base::Minutes(5);

  explicit ActionTracker(content::BrowserContext* browser_context);
  ~ActionTracker();
  ActionTracker(const ActionTracker& other) = delete;
  ActionTracker& operator=(const ActionTracker& other) = delete;

  // TODO(crbug.com/1043367): Use a task environment to avoid having to set
  // clocks just for tests.

  // Sets a custom Clock to use in tests. |clock| should be owned by the caller
  // of this function.
  void SetClockForTests(const base::Clock* clock);

  // Sets a custom RetainingOneShotTimer to use in tests, which replaces
  // |trim_rules_timer_|.
  void SetTimerForTest(
      std::unique_ptr<base::RetainingOneShotTimer> injected_trim_rules_timer);

  // Disables checking whether a tab ID corresponds to an existing tab when a
  // rule is matched. Used for unit tests where WebContents/actual tabs do not
  // exist.
  void SetCheckTabIdOnRuleMatchForTest(bool check_tab_id);

  // Called whenever a request matches with a rule.
  void OnRuleMatched(const RequestAction& request_action,
                     const WebRequestInfo& request_info);

  // Updates the action count for all tabs for the specified |extension_id|'s
  // extension action. Called when the extension calls setExtensionActionOptions
  // to enable setting the action count as badge text.
  void OnActionCountAsBadgeTextPreferenceEnabled(
      const ExtensionId& extension_id) const;

  // Clears the TrackedInfo for the specified |extension_id| for all tabs.
  // Called when an extension's ruleset is removed.
  void ClearExtensionData(const ExtensionId& extension_id);

  // Clears the TrackedInfo for every extension for the specified |tab_id|.
  // Called when the tab has been closed.
  void ClearTabData(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 TrackedInfo for all extensions for the given |tab_id|.
  void ResetTrackedInfoForTab(int tab_id, int64_t navigation_id);

  // Returns all matched rules for |extension|. If |tab_id| is provided, only
  // rules matched for |tab_id| will be returned.
  std::vector<api::declarative_net_request::MatchedRuleInfo> GetMatchedRules(
      const Extension& extension,
      const absl::optional<int>& tab_id,
      const base::Time& min_time_stamp);

  // Returns the number of matched rules in |rules_tracked_| for the given
  // |extension_id| and |tab_id|. If |trim_non_active_rules| is true,
  // TrimRulesFromNonActiveTabs is invoked before returning the matched rule
  // count, similar to GetMatchedRules. Should only be used for tests.
  int GetMatchedRuleCountForTest(const ExtensionId& extension_id,
                                 int tab_id,
                                 bool trim_non_active_rules);

  // Returns the number of matched rules in |pending_navigation_actions_| for
  // the given |extension_id| and |navigation_id|. Should only be used for
  // tests.
  int GetPendingRuleCountForTest(const ExtensionId& extension_id,
                                 int64_t navigation_id);

  // Increments the action count for the given |extension_id| and |tab_id|.
  // A negative value for |increment| will decrement the action count, but the
  // action count will never be less than 0.
  void IncrementActionCountForTab(const ExtensionId& extension_id,
                                  int tab_id,
                                  int increment);

 private:
  // 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;
    T secondary_id;

    bool operator<(const TrackedInfoContextKey& other) const;
  };

  using ExtensionTabIdKey = TrackedInfoContextKey<int>;
  using ExtensionNavigationIdKey = TrackedInfoContextKey<int64_t>;

  // Represents a matched rule. This is used as a lightweight version of
  // api::declarative_net_request::MatchedRuleInfo.
  struct TrackedRule {
    TrackedRule(int rule_id, RulesetID ruleset_id);
    TrackedRule(const TrackedRule& other) = delete;
    TrackedRule& operator=(const TrackedRule& other) = delete;

    const int rule_id;
    const RulesetID ruleset_id;

    // The timestamp for when the rule was matched.
    const base::Time time_stamp;
  };

  // Info tracked for each ExtensionTabIdKey or ExtensionNavigationIdKey.
  struct TrackedInfo {
    TrackedInfo();
    ~TrackedInfo();
    TrackedInfo(const TrackedInfo& other) = delete;
    TrackedInfo& operator=(const TrackedInfo& other) = delete;
    TrackedInfo(TrackedInfo&&);
    TrackedInfo& operator=(TrackedInfo&&);

    // The number of actions matched. Invalid when the corresponding
    // TrackedInfoContextKey has a tab_id of -1. Does not include allow rules.
    size_t action_count = 0;

    // The list of rules matched. Includes allow rules.
    std::list<TrackedRule> matched_rules;
  };

  // Called from OnRuleMatched. Dispatches a OnRuleMatchedDebug event to the
  // observer for the extension specified by |request_action.extension_id|.
  void DispatchOnRuleMatchedDebugIfNeeded(
      const RequestAction& request_action,
      api::declarative_net_request::RequestDetails request_details);

  // For all matched rules attributed to |tab_id|, set their tab ID to the
  // unknown tab ID (-1). Called by ClearTabData and ResetActionCountForTab.
  void TransferRulesOnTabInvalid(int tab_id);

  // Removes all rules in |rules_tracked_| older than
  // |kNonActiveTabRuleLifespan| from non active tabs (i.e. |tab_id| = -1).
  // Called periodically to ensure no rules attributed to the unknown tab ID in
  // |rules_tracked_| are older than |kNonActiveTabRuleLifespan|.
  void TrimRulesFromNonActiveTabs();

  // Schedules TrimRulesFromNonActiveTabs to be run after
  // |kNonActiveTabRuleLifespan|. Called in the constructor and whenever
  // |trim_rules_timer_| gets set.
  void StartTrimRulesTask();

  // Converts an internally represented |tracked_rule| to a MatchedRuleInfo.
  api::declarative_net_request::MatchedRuleInfo CreateMatchedRuleInfo(
      const Extension& extension,
      const TrackedRule& tracked_rule,
      int tab_id) const;

  // A timer to call TrimRulesFromNonActiveTabs with an interval of
  // |kNonActiveTabRuleLifespan|.
  std::unique_ptr<base::RetainingOneShotTimer> trim_rules_timer_ =
      std::make_unique<base::RetainingOneShotTimer>();

  // Maps a pair of (extension ID, tab ID) to the TrackedInfo for that pair.
  std::map<ExtensionTabIdKey, TrackedInfo> rules_tracked_;

  // Maps a pair of (extension ID, navigation ID) to the TrackedInfo matched for
  // the main-frame request associated with the navigation ID in the key. The
  // TrackedInfo is added to |rules_tracked_| once the navigation commits.
  std::map<ExtensionNavigationIdKey, TrackedInfo> pending_navigation_actions_;

  raw_ptr<content::BrowserContext> browser_context_;

  raw_ptr<ExtensionPrefs> extension_prefs_;
};

}  // namespace declarative_net_request
}  // namespace extensions

#endif  // EXTENSIONS_BROWSER_API_DECLARATIVE_NET_REQUEST_ACTION_TRACKER_H_