extensions/browser/api/declarative_net_request/global_rules_tracker.cc

// Copyright 2020 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "extensions/browser/api/declarative_net_request/global_rules_tracker.h"

#include "extensions/browser/api/declarative_net_request/utils.h"
#include "extensions/browser/extension_prefs.h"
#include "extensions/browser/extension_registry.h"
#include "extensions/common/api/declarative_net_request.h"
#include "extensions/common/extension.h"

namespace extensions {
namespace declarative_net_request {

namespace {

namespace dnr_api = api::declarative_net_request;

// Returns the total allocated global rule count, as maintained in extension
// prefs from the set of installed extensions in the registry.
size_t CalculateAllocatedGlobalRuleCount(
    const ExtensionPrefs* extension_prefs,
    const ExtensionRegistry* extension_registry) {
  std::unique_ptr<ExtensionSet> installed_extensions =
      extension_registry->GenerateInstalledExtensionsSet();

  // For each extension, fetch its allocated rules count and add it to
  // |allocated_global_rule_count_|.
  size_t allocated_rule_count;
  size_t allocated_global_rule_count = 0;
  for (const auto& extension : *installed_extensions) {
    if (extension_prefs->GetDNRAllocatedGlobalRuleCount(
            extension->id(), &allocated_rule_count)) {
      allocated_global_rule_count += allocated_rule_count;
    }
  }

  return allocated_global_rule_count;
}

// Returns the total allocated global rule count, as maintained in extension
// prefs from the set of installed extensions from prefs. This should be called
// only if the extension registry has not been populated yet (e.g. when a
// browser session has just started).
size_t CalculateInitialAllocatedGlobalRuleCount(
    const ExtensionPrefs* extension_prefs) {
  std::unique_ptr<ExtensionPrefs::ExtensionsInfo> extensions_info(
      extension_prefs->GetInstalledExtensionsInfo());

  // For each extension, fetch its allocated rules count and add it to
  // |allocated_global_rule_count_|.
  size_t allocated_rule_count = 0;
  size_t allocated_global_rule_count = 0;
  for (const auto& info : *extensions_info) {
    // Skip extensions that were loaded from the command-line because we don't
    // want those to persist across browser restart.
    if (info->extension_location == mojom::ManifestLocation::kCommandLine)
      continue;

    if (extension_prefs->GetDNRAllocatedGlobalRuleCount(
            info->extension_id, &allocated_rule_count)) {
      allocated_global_rule_count += allocated_rule_count;
    }
  }

  return allocated_global_rule_count;
}

}  // namespace

GlobalRulesTracker::GlobalRulesTracker(ExtensionPrefs* extension_prefs,
                                       ExtensionRegistry* extension_registry)
    : allocated_global_rule_count_(
          CalculateInitialAllocatedGlobalRuleCount(extension_prefs)),
      extension_prefs_(extension_prefs),
      extension_registry_(extension_registry) {}

GlobalRulesTracker::~GlobalRulesTracker() = default;

size_t GlobalRulesTracker::GetAllocatedGlobalRuleCountForTesting() const {
  DCHECK_EQ(
      allocated_global_rule_count_,
      CalculateAllocatedGlobalRuleCount(extension_prefs_, extension_registry_));
  return allocated_global_rule_count_;
}

bool GlobalRulesTracker::OnExtensionRuleCountUpdated(
    const ExtensionId& extension_id,
    size_t new_rule_count) {
  // Each extension ruleset is allowed to have up to
  // |GetMaximumRulesPerRuleset()| rules during indexing.
  DCHECK_LE(new_rule_count, static_cast<size_t>(GetMaximumRulesPerRuleset()) *
                                dnr_api::MAX_NUMBER_OF_STATIC_RULESETS);

  bool keep_excess_allocation =
      extension_prefs_->GetDNRKeepExcessAllocation(extension_id);

  if (new_rule_count <=
      static_cast<size_t>(GetStaticGuaranteedMinimumRuleCount())) {
    if (!keep_excess_allocation)
      ClearExtensionAllocation(extension_id);

    return true;
  }

  size_t old_allocated_rule_count = GetAllocationInPrefs(extension_id);
  DCHECK_GE(allocated_global_rule_count_, old_allocated_rule_count);

  size_t new_allocated_rule_count =
      new_rule_count - GetStaticGuaranteedMinimumRuleCount();

  if (new_allocated_rule_count == old_allocated_rule_count)
    return true;

  if (keep_excess_allocation &&
      old_allocated_rule_count > new_allocated_rule_count) {
    // Retain the extension's current excess allocation and allow the update.
    return true;
  }

  size_t new_global_rule_count =
      (allocated_global_rule_count_ - old_allocated_rule_count) +
      new_allocated_rule_count;

  // If updating this extension's rule count would cause the global rule count
  // to be exceeded, don't commit the update and return false.
  if (new_global_rule_count > static_cast<size_t>(GetGlobalStaticRuleLimit()))
    return false;

  if (keep_excess_allocation) {
    DCHECK_GT(new_allocated_rule_count, old_allocated_rule_count);
    // The extension is now using more than it's pre-update rule allocation.
    // Remove it's ability to keep the excess allocation.
    extension_prefs_->SetDNRKeepExcessAllocation(extension_id, false);
  }

  allocated_global_rule_count_ = new_global_rule_count;
  extension_prefs_->SetDNRAllocatedGlobalRuleCount(extension_id,
                                                   new_allocated_rule_count);

  return true;
}

size_t GlobalRulesTracker::GetUnallocatedRuleCount() const {
  return GetGlobalStaticRuleLimit() - allocated_global_rule_count_;
}

size_t GlobalRulesTracker::GetAvailableAllocation(
    const ExtensionId& extension_id) const {
  return GetUnallocatedRuleCount() + GetAllocationInPrefs(extension_id);
}

void GlobalRulesTracker::ClearExtensionAllocation(
    const ExtensionId& extension_id) {
  size_t allocated_rule_count = 0;
  if (!extension_prefs_->GetDNRAllocatedGlobalRuleCount(
          extension_id, &allocated_rule_count)) {
    return;
  }

  DCHECK_GE(allocated_global_rule_count_, allocated_rule_count);

  allocated_global_rule_count_ -= allocated_rule_count;
  extension_prefs_->SetDNRAllocatedGlobalRuleCount(extension_id,
                                                   0 /* rule_count */);
}

size_t GlobalRulesTracker::GetAllocationInPrefs(
    const ExtensionId& extension_id) const {
  size_t allocated_rule_count = 0;
  extension_prefs_->GetDNRAllocatedGlobalRuleCount(extension_id,
                                                   &allocated_rule_count);
  return allocated_rule_count;
}

}  // namespace declarative_net_request
}  // namespace extensions