Bug 1821051 - Enforce rule limit for session rules r=rpl

Differential Revision: https://phabricator.services.mozilla.com/D172252

toolkit/components/extensions/parent/ext-declarativeNetRequest.js

@@ -67,7 +67,10 @@ this.declarativeNetRequest = class extends ExtensionAPI {
           if (failures.length) {
             throw new ExtensionError(failures[0].message);
           }
-          ruleManager.setSessionRules(ruleValidator.getValidatedRules());
+          let validatedRules = ruleValidator.getValidatedRules();
+          let ruleQuotaCounter = new ExtensionDNR.RuleQuotaCounter();
+          ruleQuotaCounter.tryAddRules("_session", validatedRules);
+          ruleManager.setSessionRules(validatedRules);
         },
 
         async getEnabledRulesets() {

toolkit/components/extensions/test/xpcshell/test_ext_dnr_dynamic_rules.js

@@ -257,6 +257,10 @@ add_task(async function test_dynamic_rules_count_limits() {
         "Expect the number of dynamic rules to be still allowed, despite the session rule added"
       );
 
+      // NOTE: In this test we do not exceed the quota of session rules. The
+      // updateSessionRules call here is to verify that the quota of session and
+      // dynamic rules are separate. The limits for session rules are tested
+      // by session_rules_total_rule_limit in test_ext_dnr_session_rules.js.
       browser.test.assertDeepEq(
         updateSessionResult,
         { status: "fulfilled", value: undefined },

toolkit/components/extensions/test/xpcshell/test_ext_dnr_session_rules.js

@@ -1031,3 +1031,85 @@ add_task(async function validate_action_redirect_transform() {
     },
   });
 });
+
+add_task(async function session_rules_total_rule_limit() {
+  await runAsDNRExtension({
+    background: async dnrTestUtils => {
+      const {
+        MAX_NUMBER_OF_DYNAMIC_AND_SESSION_RULES,
+      } = browser.declarativeNetRequest;
+
+      let inputRules = [];
+      let nextRuleId = 1;
+      for (let i = 0; i < MAX_NUMBER_OF_DYNAMIC_AND_SESSION_RULES; ++i) {
+        inputRules.push(dnrTestUtils.makeRuleInput(nextRuleId++));
+      }
+      let excessRule = dnrTestUtils.makeRuleInput(nextRuleId++);
+
+      browser.test.log(`Should be able to add ${inputRules.length} rules.`);
+      await browser.declarativeNetRequest.updateSessionRules({
+        addRules: inputRules,
+      });
+
+      browser.test.assertEq(
+        inputRules.length,
+        (await browser.declarativeNetRequest.getSessionRules()).length,
+        "Added up to MAX_NUMBER_OF_DYNAMIC_AND_SESSION_RULES session rules"
+      );
+
+      await browser.test.assertRejects(
+        browser.declarativeNetRequest.updateSessionRules({
+          addRules: [excessRule],
+        }),
+        `Number of rules in ruleset "_session" exceeds MAX_NUMBER_OF_DYNAMIC_AND_SESSION_RULES.`,
+        "Should not accept more than MAX_NUMBER_OF_DYNAMIC_AND_SESSION_RULES rules"
+      );
+
+      await browser.test.assertRejects(
+        browser.declarativeNetRequest.updateSessionRules({
+          removeRuleIds: [inputRules[0].id],
+          addRules: [inputRules[0], excessRule],
+        }),
+        `Number of rules in ruleset "_session" exceeds MAX_NUMBER_OF_DYNAMIC_AND_SESSION_RULES.`,
+        "Removing one rule is not enough to make space for two rules"
+      );
+
+      browser.test.log("Should be able to replace one rule while at the limit");
+      await browser.declarativeNetRequest.updateSessionRules({
+        removeRuleIds: [inputRules[0].id],
+        addRules: [excessRule],
+      });
+
+      browser.test.log("Should be able to remove many rules, even at quota");
+      await browser.declarativeNetRequest.updateSessionRules({
+        // Note: inputRules[0].id was already removed, but that's fine.
+        removeRuleIds: inputRules.map(r => r.id),
+      });
+
+      browser.test.assertDeepEq(
+        [dnrTestUtils.makeRuleOutput(excessRule.id)],
+        await browser.declarativeNetRequest.getSessionRules(),
+        "Expected one rule after removing all-but-one-rule"
+      );
+
+      await browser.test.assertRejects(
+        browser.declarativeNetRequest.updateSessionRules({
+          addRules: inputRules,
+        }),
+        `Number of rules in ruleset "_session" exceeds MAX_NUMBER_OF_DYNAMIC_AND_SESSION_RULES.`,
+        "Should not be able to add MAX_NUMBER_OF_DYNAMIC_AND_SESSION_RULES when there is already a rule"
+      );
+
+      await browser.declarativeNetRequest.updateSessionRules({
+        removeRuleIds: [excessRule.id],
+      });
+      browser.test.assertDeepEq(
+        [],
+        await browser.declarativeNetRequest.getSessionRules(),
+        "Removed last remaining rule"
+      );
+
+      browser.test.notifyPass();
+    },
+  });
+});