Bug 1751436 - Add support for persistAcrossSessions in `scripting.R…

…egisteredContentScript`. r=robwu

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

toolkit/components/extensions/Extension.jsm

@@ -59,6 +59,7 @@ XPCOMUtils.defineLazyModuleGetters(lazy, {
   ExtensionPreferencesManager:
     "resource://gre/modules/ExtensionPreferencesManager.jsm",
   ExtensionProcessScript: "resource://gre/modules/ExtensionProcessScript.jsm",
+  ExtensionScriptingStore: "resource://gre/modules/ExtensionScriptingStore.jsm",
   ExtensionStorage: "resource://gre/modules/ExtensionStorage.jsm",
   ExtensionStorageIDB: "resource://gre/modules/ExtensionStorageIDB.jsm",
   ExtensionTelemetry: "resource://gre/modules/ExtensionTelemetry.jsm",
@@ -482,6 +483,13 @@ var ExtensionAddonObserver = {
       lazy.ServiceWorkerCleanUp.removeFromPrincipal(principal)
     );
 
+    // Clear the persisted dynamic content scripts created with the scripting
+    // API (if any).
+    lazy.AsyncShutdown.profileChangeTeardown.addBlocker(
+      `Clear scripting store for ${addon.id}`,
+      lazy.ExtensionScriptingStore.clear(addon.id)
+    );
+
     if (!Services.prefs.getBoolPref(LEAVE_STORAGE_PREF, false)) {
       // Clear browser.storage.local backends.
       lazy.AsyncShutdown.profileChangeTeardown.addBlocker(
@@ -3053,6 +3061,20 @@ class Extension extends ExtensionData {
 
       GlobalManager.init(this);
 
+      if (this.hasPermission("scripting")) {
+        this.state = "Startup: Initialize scripting store";
+        // We have to await here because `initSharedData` depends on the data
+        // fetched from the scripting store. This has to be done early because
+        // we need the data to run the content scripts in existing pages at
+        // startup.
+        try {
+          await lazy.ExtensionScriptingStore.initExtension(this);
+          this.state = "Startup: Scripting store initialized";
+        } catch (err) {
+          this.logError(`Failed to initialize scripting store: ${err}`);
+        }
+      }
+
       this.initSharedData();
 
       this.policy.active = false;

toolkit/components/extensions/ExtensionScriptingStore.jsm

@@ -0,0 +1,306 @@
+/* -*- Mode: indent-tabs-mode: nil; js-indent-level: 2 -*- */
+/* vim: set sts=2 sw=2 et tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+"use strict";
+
+const { ExtensionUtils } = ChromeUtils.import(
+  "resource://gre/modules/ExtensionUtils.jsm"
+);
+
+const lazy = {};
+
+ChromeUtils.defineModuleGetter(
+  lazy,
+  "FileUtils",
+  "resource://gre/modules/FileUtils.jsm"
+);
+
+ChromeUtils.defineModuleGetter(
+  lazy,
+  "KeyValueService",
+  "resource://gre/modules/kvstore.jsm"
+);
+
+class Store {
+  async _init() {
+    const { path: storePath } = lazy.FileUtils.getDir("ProfD", [
+      "extension-store",
+    ]);
+    // Make sure the folder exists.
+    await IOUtils.makeDirectory(storePath, { ignoreExisting: true });
+    this._store = await lazy.KeyValueService.getOrCreate(
+      storePath,
+      "scripting-contentScripts"
+    );
+  }
+
+  lazyInit() {
+    if (!this._initPromise) {
+      this._initPromise = this._init();
+    }
+
+    return this._initPromise;
+  }
+
+  /**
+   * Returns all the stored scripts for a given extension (ID).
+   *
+   * @param {string} extensionId An extension ID
+   * @returns {Array} An array of scripts
+   */
+  async getAll(extensionId) {
+    await this.lazyInit();
+    const pairs = await this.getByExtensionId(extensionId);
+
+    return pairs.map(([_, script]) => script);
+  }
+
+  /**
+   * Writes all the scripts provided for a given extension (ID) to the internal
+   * store (which is eventually stored on disk).
+   *
+   * We store each script of an extension as a key/value pair where the key is
+   * `<extensionId>/<scriptId>` and the value is the corresponding script
+   * details as a JSON string.
+   *
+   * The format on disk should look like this one:
+   *
+   * ```
+   * {
+   *   "@extension-id/script-1": {"id: "script-1", <other props>},
+   *   "@extension-id/script-2": {"id: "script-2", <other props>}
+   * }
+   * ```
+   *
+   * @param {string} extensionId An extension ID
+   * @param {Array} scripts An array of scripts to store on disk
+   */
+  async writeMany(extensionId, scripts) {
+    await this.lazyInit();
+
+    return this._store.writeMany(
+      scripts.map(script => [
+        `${extensionId}/${script.id}`,
+        JSON.stringify(script),
+      ])
+    );
+  }
+
+  /**
+   * Deletes all the stored scripts for a given extension (ID).
+   *
+   * @param {string} extensionId An extension ID
+   */
+  async deleteAll(extensionId) {
+    await this.lazyInit();
+    const pairs = await this.getByExtensionId(extensionId);
+
+    return Promise.all(pairs.map(([key, _]) => this._store.delete(key)));
+  }
+
+  /**
+   * Returns an array of key/script pairs from the internal store belonging to
+   * the given extension (ID).
+   *
+   * The data returned by this method should look like this (assuming we have
+   * two scripts named `script-1` and `script-2` for the extension with ID
+   * `@extension-id`):
+   *
+   * ```
+   * [
+   *   ["@extension-id/script-1", {"id: "script-1", <other props>}],
+   *   ["@extension-id/script-2", {"id: "script-2", <other props>}]
+   * ]
+   * ```
+   *
+   * @param {string} extensionId An extension ID
+   * @returns {Array} An array of key/script pairs
+   */
+  async getByExtensionId(extensionId) {
+    await this.lazyInit();
+
+    const entries = [];
+    // Retrieve all the scripts registered for the given extension ID by
+    // enumerating all keys that are stored in a lexical order.
+    const enumerator = await this._store.enumerate(
+      `${extensionId}/`, // from_key (inclusive)
+      `${extensionId}0` // to_key (exclusive)
+    );
+
+    while (enumerator.hasMoreElements()) {
+      const { key, value } = enumerator.getNext();
+      entries.push([key, JSON.parse(value)]);
+    }
+
+    return entries;
+  }
+}
+
+const store = new Store();
+
+/**
+ * Given an extension and some content script options, this function returns
+ * the content script representation we use internally, which is an object with
+ * a `scriptId` and a nested object containing `options`. These (internal)
+ * objects are shared with all content processes using IPC/sharedData.
+ *
+ * This function can optionally prepend the extension's base URL to the CSS and
+ * JS paths, which is needed when we load internal scripts from the scripting
+ * store (because the UUID in the base URL changes).
+ *
+ * @param {Extension} extension
+ *        The extension that owns the content script.
+ * @param {object} options
+ *        Content script options.
+ * @param {boolean} prependBaseURL
+ *        Whether to prepend JS and CSS paths with the extension's base URL.
+ *
+ * @returns {object}
+ */
+const makeInternalContentScript = (
+  extension,
+  options,
+  prependBaseURL = false
+) => {
+  let cssPaths = options.css || [];
+  let jsPaths = options.js || [];
+
+  if (prependBaseURL) {
+    cssPaths = cssPaths.map(css => `${extension.baseURL}${css}`);
+    jsPaths = jsPaths.map(js => `${extension.baseURL}${js}`);
+  }
+
+  return {
+    scriptId: ExtensionUtils.getUniqueId(),
+    options: {
+      // We need to store the user-supplied script ID for persisted scripts.
+      id: options.id,
+      allFrames: options.allFrames || false,
+      // Although this flag defaults to true with MV3, it is not with MV2.
+      // Check permissions at runtime since we aren't checking permissions
+      // upfront.
+      checkPermissions: true,
+      cssPaths,
+      excludeMatches: options.excludeMatches,
+      jsPaths,
+      matchAboutBlank: true,
+      matches: options.matches,
+      originAttributesPatterns: null,
+      persistAcrossSessions: options.persistAcrossSessions,
+      runAt: options.runAt || "document_idle",
+    },
+  };
+};
+
+/**
+ * Given an internal content script registered with the "scripting" API (and an
+ * extension), this function returns a new object that matches the public
+ * "scripting" API.
+ *
+ * This function is primarily in `scripting.getRegisteredContentScripts()`.
+ *
+ * @param {Extension} extension
+ *        The extension that owns the content script.
+ * @param {object} internalScript
+ *        An internal script (see also: `makeInternalContentScript()`).
+ *
+ * @returns {object}
+ */
+const makePublicContentScript = (extension, internalScript) => {
+  let script = {
+    id: internalScript.id,
+    allFrames: internalScript.allFrames,
+    matches: internalScript.matches,
+    runAt: internalScript.runAt,
+    persistAcrossSessions: internalScript.persistAcrossSessions,
+  };
+
+  if (internalScript.cssPaths.length) {
+    script.css = internalScript.cssPaths.map(cssPath =>
+      cssPath.replace(extension.baseURL, "")
+    );
+  }
+
+  if (internalScript.excludeMatches?.length) {
+    script.excludeMatches = internalScript.excludeMatches;
+  }
+
+  if (internalScript.jsPaths.length) {
+    script.js = internalScript.jsPaths.map(jsPath =>
+      jsPath.replace(extension.baseURL, "")
+    );
+  }
+
+  return script;
+};
+
+const ExtensionScriptingStore = {
+  async initExtension(extension) {
+    // Load the scripts from the storage, then convert them to their internal
+    // representation and add them to the extension's registered scripts.
+    const scripts = await store.getAll(extension.id);
+
+    scripts.forEach(script => {
+      const { scriptId, options } = makeInternalContentScript(
+        extension,
+        script,
+        true /* prepend the css/js paths with the extension base URL */
+      );
+      extension.registeredContentScripts.set(scriptId, options);
+    });
+  },
+
+  getInitialScriptIdsMap(extension) {
+    // This returns the current map of public script IDs to internal IDs.
+    // `extension.registeredContentScripts` is initialized in `initExtension`,
+    // which may be updated later via the scripting API. In practice, the map
+    // of script IDs is retrieved before any scripting API method is exposed,
+    // so the return value always matches the initial result from
+    // `initExtension`.
+    return new Map(
+      Array.from(
+        extension.registeredContentScripts.entries(),
+        ([scriptId, options]) => [options.id, scriptId]
+      )
+    );
+  },
+
+  async persistAll(extension) {
+    // We only persist the scripts that should be persisted and we convert each
+    // script to their "public" representation before storing them. This is
+    // because we don't want to deal with data migrations if we ever want to
+    // change the internal representation (the "public" representation is less
+    // likely to change because it is bound to the public scripting API).
+    const scripts = Array.from(extension.registeredContentScripts.values())
+      .filter(options => options.persistAcrossSessions)
+      .map(options => makePublicContentScript(extension, options));
+
+    // We want to replace all the scripts for the extension so we should delete
+    // the existing ones first, and then write the new ones.
+    //
+    // TODO: Bug 1783131 - Implement individual updates without requiring all
+    // data to be erased and written.
+    await store.deleteAll(extension.id);
+
+    return store.writeMany(extension.id, scripts);
+  },
+
+  // Delete all the persisted scripts for the given extension (id).
+  async clear(extensionId) {
+    return store.deleteAll(extensionId);
+  },
+
+  // As its name implies, don't use this method for anything but an easy access
+  // to the internal store for testing purposes.
+  _getStoreForTesting() {
+    return store;
+  },
+};
+
+var EXPORTED_SYMBOLS = [
+  "ExtensionScriptingStore",
+  "makeInternalContentScript",
+  "makePublicContentScript",
+];

toolkit/components/extensions/moz.build

@@ -24,6 +24,7 @@ EXTRA_JS_MODULES += [
     "ExtensionPreferencesManager.jsm",
     "ExtensionProcessScript.jsm",
     "extensionProcessScriptLoader.js",
+    "ExtensionScriptingStore.jsm",
     "ExtensionSettingsStore.jsm",
     "ExtensionShortcuts.jsm",
     "ExtensionStorage.jsm",