Bug 1908925 - Enforce storage.session quota in Nightly builds r=robwu,mccr8

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

dom/base/StructuredCloneBlob.h

@@ -48,6 +48,11 @@ class StructuredCloneBlob final : public nsIMemoryReporter {
                    bool aKeepData, JS::MutableHandle<JS::Value> aResult,
                    ErrorResult& aRv);
 
+  uint64_t DataSize() {
+    return mHolder.isSome() && mHolder->HasData() ? mHolder->BufferData().Size()
+                                                  : 0;
+  }
+
   nsISupports* GetParentObject() const { return nullptr; }
   JSObject* GetWrapper() const { return nullptr; }
 

dom/chrome-webidl/StructuredCloneHolder.webidl

@@ -28,6 +28,12 @@ interface StructuredCloneHolder {
   constructor(UTF8String name, UTF8String? anonymizedName,
               any data, optional object? global = null);
 
+  /**
+   * Returns the size of serialized data in bytes. Note that this is smaller
+   * than the actual size of the object in memory, because of buffer sizes etc.
+   */
+  readonly attribute unsigned long long dataSize;
+
   /**
    * Deserializes the structured clone data in the scope of the given global,
    * and returns the result.

modules/libpref/init/all.js

@@ -3703,6 +3703,12 @@ pref("webextensions.tests", false);
 // 16MB default non-parseable upload limit for requestBody.raw.bytes
 pref("webextensions.webRequest.requestBodyMaxRawBytes", 16777216);
 
+#ifdef NIGHTLY_BUILD
+  pref("webextensions.storage.session.enforceQuota", true);
+#else
+  pref("webextensions.storage.session.enforceQuota", false);
+#endif
+
 pref("webextensions.storage.sync.enabled", true);
 // Should we use the old kinto-based implementation of storage.sync? To be removed in bug 1637465.
 pref("webextensions.storage.sync.kinto", false);

toolkit/components/extensions/ExtensionStorage.sys.mjs

@@ -5,6 +5,7 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 import { ExtensionUtils } from "resource://gre/modules/ExtensionUtils.sys.mjs";
+import { XPCOMUtils } from "resource://gre/modules/XPCOMUtils.sys.mjs";
 
 const { DefaultWeakMap, ExtensionError } = ExtensionUtils;
 
@@ -16,6 +17,13 @@ ChromeUtils.defineESModuleGetters(lazy, {
   JSONFile: "resource://gre/modules/JSONFile.sys.mjs",
 });
 
+XPCOMUtils.defineLazyPreferenceGetter(
+  lazy,
+  "enforceSessionQuota",
+  "webextensions.storage.session.enforceQuota",
+  false
+);
+
 function isStructuredCloneHolder(value) {
   return (
     value &&
@@ -484,13 +492,74 @@ ChromeUtils.defineLazyGetter(ExtensionStorage, "extensionDir", () =>
 
 ExtensionStorage.init();
 
+class QuotaMap extends Map {
+  static QUOTA_BYTES = 10485760;
+
+  bytesUsed = 0;
+
+  /**
+   * @param {string} key
+   * @param {StructuredCloneHolder} holder
+   */
+  dataSize(key, holder) {
+    // Using key.length is not really correct, but is probably less surprising
+    // for developers. We don't need an exact count, just ensure it's bounded.
+    return key.length + holder.dataSize;
+  }
+
+  /**
+   * @param {string} key
+   * @param {StructuredCloneHolder} holder
+   */
+  getSizeDelta(key, holder) {
+    let before = this.has(key) ? this.dataSize(key, this.get(key)) : 0;
+    return this.dataSize(key, holder) - before;
+  }
+
+  /** @param {Record<string, StructuredCloneHolder>} items */
+  checkQuota(items) {
+    let after = this.bytesUsed;
+    for (let [key, holder] of Object.entries(items)) {
+      after += this.getSizeDelta(key, holder);
+    }
+    if (lazy.enforceSessionQuota && after > QuotaMap.QUOTA_BYTES) {
+      throw new ExtensionError(
+        "QuotaExceededError: storage.session API call exceeded its quota limitations."
+      );
+    }
+  }
+
+  set(key, holder) {
+    this.checkQuota({ [key]: holder });
+    this.bytesUsed += this.getSizeDelta(key, holder);
+    return super.set(key, holder);
+  }
+
+  delete(key) {
+    if (this.has(key)) {
+      this.bytesUsed -= this.dataSize(key, this.get(key));
+    }
+    return super.delete(key);
+  }
+
+  clear() {
+    this.bytesUsed = 0;
+    super.clear();
+  }
+}
+
 export var extensionStorageSession = {
-  /** @type {WeakMap<Extension, Map<string, any>>} */
-  buckets: new DefaultWeakMap(_extension => new Map()),
+  /** @type {WeakMap<Extension, QuotaMap>} */
+  buckets: new DefaultWeakMap(_extension => new QuotaMap()),
 
   /** @type {WeakMap<Extension, Set<callback>>} */
   listeners: new DefaultWeakMap(_extension => new Set()),
 
+  get QUOTA_BYTES() {
+    // Even if quota is not enforced yet, report the future default of 10MB.
+    return QuotaMap.QUOTA_BYTES;
+  },
+
   /**
    * @param {Extension} extension
    * @param {null | undefined | string | string[] | object} items
@@ -523,6 +592,10 @@ export var extensionStorageSession = {
   set(extension, items) {
     let bucket = this.buckets.get(extension);
 
+    // set() below also checks the quota for each item, but
+    // this check includes all inputs to avoid partial updates.
+    bucket.checkQuota(items);
+
     let changes = {};
     for (let [key, value] of Object.entries(items)) {
       changes[key] = {
@@ -556,6 +629,24 @@ export var extensionStorageSession = {
     this.notifyListeners(extension, changes);
   },
 
+  /**
+   * @param {Extension} extension
+   * @param {null | undefined | string | string[] } keys
+   */
+  getBytesInUse(extension, keys) {
+    let bucket = this.buckets.get(extension);
+    if (keys == null) {
+      return bucket.bytesUsed;
+    }
+    let result = 0;
+    for (let k of [].concat(keys)) {
+      if (bucket.has(k)) {
+        result += bucket.dataSize(k, bucket.get(k));
+      }
+    }
+    return result;
+  },
+
   registerListener(extension, listener) {
     this.listeners.get(extension).add(listener);
     return () => {

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

@@ -283,6 +283,9 @@ this.storage = class extends ExtensionAPIPersistent {
         },
 
         session: {
+          get QUOTA_BYTES() {
+            return extensionStorageSession.QUOTA_BYTES;
+          },
           get(items) {
             return extensionStorageSession.get(extension, items);
           },
@@ -295,6 +298,9 @@ this.storage = class extends ExtensionAPIPersistent {
           clear() {
             extensionStorageSession.clear(extension);
           },
+          getBytesInUse(keys) {
+            return extensionStorageSession.getBytesInUse(extension, keys);
+          },
           onChanged: new EventManager({
             context,
             module: "storage",

toolkit/components/extensions/schemas/storage.json

@@ -168,7 +168,7 @@
         ]
       },
       {
-        "id": "StorageAreaSync",
+        "id": "StorageAreaWithUsage",
         "type": "object",
         "functions": [
           {
@@ -334,7 +334,7 @@
     ],
     "properties": {
       "sync": {
-        "$ref": "StorageAreaSync",
+        "$ref": "StorageAreaWithUsage",
         "description": "Items in the <code>sync</code> storage area are synced by the browser.",
         "properties": {
           "QUOTA_BYTES": {
@@ -386,8 +386,13 @@
       },
       "session": {
         "allowedContexts": ["devtools"],
-        "$ref": "StorageArea",
-        "description": "Items in the <code>session</code> storage area are kept in memory, and only until the either browser or extension is closed or reloaded."
+        "$ref": "StorageAreaWithUsage",
+        "description": "Items in the <code>session</code> storage area are kept in memory, and only until the either browser or extension is closed or reloaded.",
+        "properties": {
+          "QUOTA_BYTES": {
+            "description": "The maximum amount of data (in bytes, currently at 10MB) that can be stored in session storage, as measured by the StructuredCloneHolder of every value plus every key's length."
+          }
+        }
       }
     }
   }

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

@@ -163,3 +163,146 @@ add_task(
     return test_storage_session_after_crash({ persistent: false });
   }
 );
+
+async function test_storage_session_quota({ quotaEnforced }) {
+  let extension = ExtensionTestUtils.loadExtension({
+    manifest: {
+      permissions: ["storage"],
+    },
+    async background() {
+      const MB = 1_000_000;
+
+      // Max overhead per storage item. Currently it's 16 (+ key length),
+      // but that's implementation detail and might easily change.
+      const SLACK = 40;
+
+      let before = 0;
+      let error = null;
+      try {
+        // Try to store 30Mb in session storage.
+        for (let i = 0; i < 30; i++) {
+          await browser.storage.session.set({
+            [`key${i}`]: "x".repeat(MB),
+          });
+
+          let after = await browser.storage.session.getBytesInUse();
+          let delta = after - before;
+          browser.test.assertTrue(
+            delta >= MB && delta <= MB + SLACK,
+            `Expected storage.session.getBytesInUse() delta=${delta}`
+          );
+          before = after;
+        }
+      } catch (e) {
+        error = e.message;
+        browser.test.log(error);
+      }
+
+      browser.test.sendMessage("data", {
+        error,
+        data: await browser.storage.session.get(),
+      });
+
+      await browser.storage.session.remove(["key2", "key3"]);
+      let after = await browser.storage.session.getBytesInUse();
+      let delta = after - before;
+      browser.test.assertTrue(
+        // Note that we're expecting and comparing a negative delta here.
+        -delta >= 2 * MB && -delta <= 2 * (MB + SLACK),
+        `Expected getBytesInUse() after removing 2 items delta=${delta}`
+      );
+
+      error = null;
+      try {
+        await browser.storage.session.set({
+          canary: 13,
+          big: "x".repeat(5 * MB),
+        });
+      } catch (e) {
+        error = e.message;
+        browser.test.log(error);
+      }
+      let data = await browser.storage.session.get();
+
+      await browser.storage.session.clear();
+      let zero = await browser.storage.session.getBytesInUse();
+      browser.test.assertEq(zero, 0, "Zero bytes used after clear.");
+
+      const six = "x".repeat(6 * MB);
+      await browser.storage.session.set({
+        one: "x".repeat(MB),
+        six: six,
+      });
+      before = await browser.storage.session.getBytesInUse();
+      await browser.storage.session.remove("six");
+      await browser.storage.session.set({ 六: six });
+
+      after = await browser.storage.session.getBytesInUse();
+      browser.test.assertEq(
+        after - before,
+        "六".length - "six".length,
+        "Usage increased by key's length difference in js chars (not bytes)."
+      );
+      browser.test.assertEq("六".length, 1, "File encoding sanity check.");
+
+      browser.test.assertEq(
+        after,
+        await browser.storage.session.getBytesInUse(["one", "六"]),
+        "Listing all keys is equivalent to not passing any keys."
+      );
+
+      await browser.storage.session.set({ "": 13 });
+      browser.test.assertEq(
+        await browser.storage.session.getBytesInUse(""),
+        await browser.storage.session.getBytesInUse([""]),
+        `Falsy key "" is correctly interpreted.`
+      );
+
+      browser.test.sendMessage("done", { error, data });
+    },
+  });
+
+  await extension.startup();
+
+  {
+    let { error, data } = await extension.awaitMessage("data");
+
+    if (quotaEnforced) {
+      ok(error.match(/QuotaExceededError/), "Expect error in Nightly builds.");
+      equal(Object.keys(data).length, 10, "10Mb stored in Nightly builds.");
+    } else {
+      equal(error, null, "No error in release builds.");
+      equal(Object.keys(data).length, 30, "30Mb stored in release builds.");
+    }
+  }
+
+  {
+    let { error, data } = await extension.awaitMessage("done");
+    if (quotaEnforced) {
+      ok(error.match(/QuotaExceededError/), "Expect error in Nightly builds.");
+      ok(!data.canary, "No partial updates on error.");
+    } else {
+      equal(data.canary, 13, "Without quota enforcement, canary was set.");
+    }
+  }
+
+  await extension.unload();
+}
+
+add_task(
+  {
+    pref_set: [["webextensions.storage.session.enforceQuota", false]],
+  },
+  async function test_storage_session_quota_no_pref() {
+    await test_storage_session_quota({ quotaEnforced: false });
+  }
+);
+
+add_task(
+  {
+    pref_set: [["webextensions.storage.session.enforceQuota", true]],
+  },
+  async function test_storage_session_quota_nightly() {
+    await test_storage_session_quota({ quotaEnforced: true });
+  }
+);

tools/@types/lib.gecko.dom.d.ts

@@ -18807,6 +18807,7 @@ declare var StreamFilterDataEvent: {
 };
 
 interface StructuredCloneHolder {
+    readonly dataSize: number;
     deserialize(global: any, keepData?: boolean): any;
 }