Bug 1804969 - Rewrite chrome:// JS imports in Storybook r=mconley,hjones

This patch will rewrite all chrome:// URLs in .mjs files, but it isn't
emitting proper URLs for assets. This means that JS imports will map
correctly, but any img/css references won't have a valid path outside of
local development and CSS files that use @import will not resolve imports
correctly.

To reference images and CSS files you will still need to ensure those files
are in the Storybook static path and use a separate URL to reference them
in the `window.IS_STORYBOOK` case.

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

.hgignore

@@ -260,6 +260,7 @@ toolkit/components/certviewer/content/package-lock.json
 # Ignore Storybook generated files
 ^browser/components/storybook/node_modules/
 ^browser/components/storybook/storybook-static/
+^browser/components/storybook/.storybook/rewrites.js
 
 # Ignore jscodeshift installed by mach esmify on windows
 ^tools/esmify/jscodeshift

browser/components/storybook/.storybook/chrome-uri-loader.js

@@ -0,0 +1,98 @@
+/* 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/. */
+/* eslint-env node */
+
+/**
+ * This file contains a webpack loader which has the goal of rewriting chrome://
+ * URIs to local paths. This allows JS files loaded in Storybook to load JS
+ * files using their chrome:// URI. Using the chrome:// URI is avoidable in many
+ * cases, however in some cases such as importing the lit.all.mjs file from
+ * browser/components/ there is no way to avoid it on the Firefox side.
+ *
+ * This loader depends on the `./mach storybook manifest` step to generate the
+ * rewrites.js file. That file exports an object with the files we know how to
+ * rewrite chrome:// URIs for.
+ *
+ * This loader allows code like this to work with storybook:
+ *
+ *     import { html } from "chrome://global/content/vendor/lit.all.mjs";
+ *     import "chrome://global/content/elements/moz-button-group.mjs";
+ *
+ * In this example the file would be rewritten in the webpack bundle as:
+ *
+ *     import { html } from "toolkit/content/widgets/vendor/lit.all.mjs";
+ *     import "toolkit/content/widgets/moz-button-group/moz-button-group.mjs";
+ */
+
+const path = require("path");
+
+// Object<ChromeURI, LocalPath> - This is generated by `./mach storybook manifest`.
+const rewrites = require("./rewrites.js");
+
+const projectRoot = path.join(process.cwd(), "../../..");
+
+/**
+ * Return an array of the unique chrome:// URIs referenced in this file.
+ *
+ * @param {string} source - The source file to scan.
+ * @returns {string[]} Unique list of chrome:// URIs
+ */
+function getReferencedChromeUris(source) {
+  // We can only rewrite files that get imported. Which means currently we only
+  // support .js and .mjs. In the future we hope to rewrite .css and .svg.
+  const chromeRegex = /chrome:\/\/.*?\.(js|mjs)/g;
+  const matches = new Set();
+  for (let match of source.matchAll(chromeRegex)) {
+    // Add the full URI to the set of matches.
+    matches.add(match[0]);
+  }
+  return [...matches];
+}
+
+/**
+ * Replace references to chrome:// URIs with the relative path on disk from the
+ * project root.
+ *
+ * @this {WebpackLoader} https://webpack.js.org/api/loaders/
+ * @param {string} source - The source file to update.
+ * @returns {string} The updated source.
+ */
+async function rewriteChromeUris(source) {
+  const chromeUriToLocalPath = new Map();
+  // We're going to rewrite the chrome:// URIs, find all referenced URIs.
+  let chromeDependencies = getReferencedChromeUris(source);
+  for (let chromeUri of chromeDependencies) {
+    let localRelativePath = rewrites[chromeUri];
+    if (localRelativePath) {
+      localRelativePath = localRelativePath.replaceAll("\\", "/");
+      // Store the mapping to a local path for this chrome URI.
+      chromeUriToLocalPath.set(chromeUri, localRelativePath);
+      // Tell webpack the file being handled depends on the referenced file.
+      this.addDependency(path.join(projectRoot, localRelativePath));
+    }
+  }
+  // Rewrite the source file with mapped chrome:// URIs.
+  let rewrittenSource = source;
+  for (let [chromeUri, localPath] of chromeUriToLocalPath.entries()) {
+    rewrittenSource = rewrittenSource.replaceAll(chromeUri, localPath);
+  }
+  return rewrittenSource;
+}
+
+/**
+ * The WebpackLoader export. Runs async since apparently that's preferred.
+ *
+ * @param {string} source - The source to rewrite.
+ * @param {Map} sourceMap - Source map data, unused.
+ * @param {Object} meta - Metadata, unused.
+ */
+module.exports = async function chromeUriLoader(source) {
+  // Get a callback to tell webpack when we're done.
+  const callback = this.async();
+  // Rewrite the source async since that appears to be preferred (and will be
+  // necessary once we support rewriting CSS/SVG/etc).
+  const newSource = await rewriteChromeUris.call(this, source);
+  // Give webpack the rewritten content.
+  callback(null, newSource);
+};

browser/components/storybook/.storybook/main.js

@@ -40,6 +40,11 @@ module.exports = {
       type: "asset/source",
     });
 
+    config.module.rules.push({
+      test: /\.mjs/,
+      loader: path.resolve(__dirname, "./chrome-uri-loader.js"),
+    });
+
     config.optimization = {
       splitChunks: false,
       runtimeChunk: false,

browser/components/storybook/mach_commands.py

@@ -2,7 +2,12 @@
 # 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/.
 
+import json
+import os
+
+import mozpack.path as mozpath
 from mach.decorators import Command, SubCommand
+from mozpack.manifests import InstallManifest
 
 
 def run_mach(command_context, cmd, **kwargs):
@@ -23,6 +28,7 @@ def run_npm(command_context, args):
     description="Start the Storybook server",
 )
 def storybook_run(command_context):
+    storybook_manifest(command_context)
     return run_npm(command_context, args=["run", "storybook"])
 
 
@@ -46,4 +52,92 @@ def storybook_install(command_context):
     description="Build the Storybook for export.",
 )
 def storybook_build(command_context):
+    storybook_manifest(command_context)
     return run_npm(command_context, args=["run", "build-storybook"])
+
+
+@SubCommand(
+    "storybook",
+    "manifest",
+    description="Create rewrites.js which has mappings from chrome:// URIs to local paths. "
+    "Requires a ./mach build faster build. Required for our chrome-uri-loader.js webpack loader.",
+)
+def storybook_manifest(command_context):
+    config_environment = command_context.config_environment
+    # The InstallManifest object will have mappings of JAR entries to paths on disk.
+    unified_manifest = InstallManifest(
+        mozpath.join(config_environment.topobjdir, "faster", "unified_install_dist_bin")
+    )
+    paths = {}
+
+    for dest, entry in unified_manifest._dests.items():
+        # dest in the JAR path
+        # entry can be many things, but we care about the [1, file_path] form
+        #   1 essentially means this is a file
+        if (
+            entry[0] == 1
+            and (dest.endswith(".js") or dest.endswith(".mjs"))
+            and (
+                dest.startswith("chrome/toolkit/") or dest.startswith("browser/chrome/")
+            )
+        ):
+            try:
+                # Try to map the dest to a chrome URI. This could fail for some weird cases that
+                # don't seem like they're worth handling.
+                chrome_uri = _parse_dest_to_chrome_uri(dest)
+                # Since we run through mach we're relative to the project root here.
+                paths[chrome_uri] = os.path.relpath(entry[1])
+            except Exception as e:
+                # Log the files that failed, this could get noisy but the list is short for now.
+                print('Error rewriting to chrome:// URI "{}" [{}]'.format(dest, e))
+                pass
+
+    with open("browser/components/storybook/.storybook/rewrites.js", "w") as f:
+        f.write("module.exports = ")
+        json.dump(paths, f, indent=2)
+
+
+def _parse_dest_to_chrome_uri(dest):
+    """Turn a jar destination into a chrome:// URI. Will raise an error on unknown input."""
+
+    global_start = dest.find("global/")
+    content_start = dest.find("content/")
+    skin_classic_browser = "skin/classic/browser/"
+    browser_skin_start = dest.find(skin_classic_browser)
+    package, provider, path = "", "", ""
+
+    if global_start != -1:
+        # e.g. chrome/toolkit/content/global/vendor/lit.all.mjs
+        #      chrome://global/content/vendor/lit.all.mjs
+        # If the jar location has global in it, then:
+        #   * the package is global,
+        #   * the portion after global should be the path,
+        #   * the provider is in the path somewhere (we want skin or content).
+        package = "global"
+        provider = "skin" if "/skin/" in dest else "content"
+        path = dest[global_start + len("global/") :]
+    elif content_start != -1:
+        # e.g. browser/chrome/browser/content/browser/aboutDialog.js
+        #      chrome://browser/content/aboutDialog.js
+        # e.g. chrome/toolkit/content/mozapps/extensions/shortcuts.js
+        #      chrome://mozapps/content/extensions/shortcuts.js
+        # If the jar location has content/ in it, then:
+        #   * starting from "content/" split on slashes and,
+        #   * the provider is "content",
+        #   * the package is the next part,
+        #   * the path is the remainder.
+        provider, package, path = dest[content_start:].split("/", 2)
+    elif browser_skin_start != -1:
+        # e.g. browser/chrome/browser/skin/classic/browser/browser.css
+        #      chrome://browser/skin/browser.css
+        # If the jar location has skin/classic/browser/ in it, then:
+        #   * the package is browser,
+        #   * the provider is skin,
+        #   * the path is what remains after sking/classic/browser.
+        package = "browser"
+        provider = "skin"
+        path = dest[browser_skin_start + len(skin_classic_browser) :]
+
+    return "chrome://{package}/{provider}/{path}".format(
+        package=package, provider=provider, path=path
+    )