|
| 1 | +import { Module, register } from 'node:module' |
| 2 | +import { pathToFileURL } from 'node:url' |
| 3 | +import { MessageChannel, type MessagePort } from 'node:worker_threads' |
| 4 | +import * as fs from 'node:fs/promises' |
| 5 | + |
| 6 | +// NOTE: If the minimum supported Node version gets above v23.6.0+ then this |
| 7 | +// entire file can be massively simplified by using `registerHooks(…)` |
| 8 | +interface LoaderState { |
| 9 | + /** |
| 10 | + * Whether or not the loader is enabled |
| 11 | + */ |
| 12 | + enabled: boolean |
| 13 | + |
| 14 | + /** |
| 15 | + * The list of "hooked" module IDs |
| 16 | + */ |
| 17 | + modules: string[] |
| 18 | + |
| 19 | + /** |
| 20 | + * The port used to communicate with the main thread |
| 21 | + */ |
| 22 | + port?: MessagePort | null |
| 23 | +} |
| 24 | + |
| 25 | +export interface ModuleHook { |
| 26 | + enable: () => Promise<void> |
| 27 | + disable: () => Promise<void> |
| 28 | + during: <T>(fn: () => Promise<T>) => Promise<T> |
| 29 | +} |
| 30 | + |
| 31 | +/** |
| 32 | + * Hooks `require(…)`, `await import(…)`, and `import from "…"` so that the |
| 33 | + * given modules are returned instead of being loaded from disk. |
| 34 | + */ |
| 35 | +export function defineModules(modules: Record<string, unknown>): ModuleHook { |
| 36 | + // The channel used to communicate between the main and loader threads |
| 37 | + // - port1: used by the main thread |
| 38 | + // - port2: used by the loader thread |
| 39 | + let channel = new MessageChannel() |
| 40 | + |
| 41 | + // The current state of the loader |
| 42 | + // A copy of this is kept in and used by the loader thread |
| 43 | + let state: LoaderState = { |
| 44 | + enabled: false, |
| 45 | + modules: Object.keys(modules), |
| 46 | + } |
| 47 | + |
| 48 | + function update(partial: Partial<LoaderState>) { |
| 49 | + Object.assign(state, partial) |
| 50 | + channel.port1.postMessage({ state }) |
| 51 | + } |
| 52 | + |
| 53 | + // Define a global function that can be used to load bundled modules |
| 54 | + // This is used by both the require() replacement and the ESM loader |
| 55 | + globalThis.__tw_load__ = (id: string) => modules[id] |
| 56 | + |
| 57 | + // Hook into require() and createRequire() so they can load the given modules |
| 58 | + function wrapRequire(original: NodeJS.Require) { |
| 59 | + function customRequire(id: string) { |
| 60 | + fs.appendFile('loader.log', 'loader require(' + id + ')\n') |
| 61 | + |
| 62 | + if (!state.enabled) return original.call(this, id) |
| 63 | + if (!state.modules.includes(id)) return original.call(this, id) |
| 64 | + return globalThis.__tw_load__(id) |
| 65 | + } |
| 66 | + |
| 67 | + function customresolve(id: string) { |
| 68 | + if (!state.enabled) return original.resolve.apply(this, arguments) |
| 69 | + if (!state.modules.includes(id)) return original.resolve.apply(this, arguments) |
| 70 | + return id |
| 71 | + } |
| 72 | + |
| 73 | + return Object.assign( |
| 74 | + customRequire, |
| 75 | + // Make sure we carry over other properties of the original require(…) |
| 76 | + original, |
| 77 | + // Replace `require.resolve(…)` with our custom resolver |
| 78 | + { resolve: customresolve }, |
| 79 | + ) |
| 80 | + } |
| 81 | + |
| 82 | + let origRequire = Module.prototype.require |
| 83 | + let origCreateRequire = Module.createRequire |
| 84 | + |
| 85 | + // Augment the default "require" available in every CJSS module |
| 86 | + Module.prototype.require = wrapRequire(origRequire) |
| 87 | + |
| 88 | + // Augment any "require" created by the "createRequire" method so require |
| 89 | + // calls used by ES modules are also intercepted. |
| 90 | + Module.createRequire = function () { |
| 91 | + return wrapRequire(origCreateRequire.apply(this, arguments)) |
| 92 | + } |
| 93 | + |
| 94 | + // Hook into the static and dynamic ESM imports so that they can load bundled modules |
| 95 | + let uri = `data:text/javascript;base64,${btoa(loader)}` |
| 96 | + channel.port2.unref() |
| 97 | + register(uri, { |
| 98 | + parentURL: pathToFileURL(__filename), |
| 99 | + transferList: [channel.port2], |
| 100 | + data: { |
| 101 | + state: { |
| 102 | + ...state, |
| 103 | + port: channel.port2, |
| 104 | + }, |
| 105 | + }, |
| 106 | + }) |
| 107 | + |
| 108 | + let enable = async () => { |
| 109 | + await fs.appendFile('loader.log', 'loader enable' + '\n') |
| 110 | + update({ enabled: true }) |
| 111 | + } |
| 112 | + |
| 113 | + let disable = async () => { |
| 114 | + update({ enabled: false }) |
| 115 | + await fs.appendFile('loader.log', 'loader disable' + '\n') |
| 116 | + } |
| 117 | + |
| 118 | + let during = async <T>(fn: () => Promise<T>) => { |
| 119 | + await enable() |
| 120 | + try { |
| 121 | + return await fn() |
| 122 | + } finally { |
| 123 | + await disable() |
| 124 | + } |
| 125 | + } |
| 126 | + |
| 127 | + return { enable, disable, during } |
| 128 | +} |
| 129 | + |
| 130 | +/** |
| 131 | + * The loader here is embedded as a string rather than a separate JS file because that complicates |
| 132 | + * the build process. We can turn this into a data URI and use it directly. It lets us keep this |
| 133 | + * file entirely self-contained feature-wise. |
| 134 | + */ |
| 135 | +const js = String.raw |
| 136 | +const loader = js` |
| 137 | + import { Module } from "node:module"; |
| 138 | + import * as fs from "node:fs/promises"; |
| 139 | +
|
| 140 | + /** @type {LoaderState} */ |
| 141 | + const state = { |
| 142 | + enabled: false, |
| 143 | + modules: [], |
| 144 | + port: null, |
| 145 | + }; |
| 146 | +
|
| 147 | + /** Updates the current state of the loader */ |
| 148 | + function sync(data) { |
| 149 | + Object.assign(state, data.state ?? {}) |
| 150 | + } |
| 151 | +
|
| 152 | + /** Set up communication with the main thread */ |
| 153 | + export async function initialize(data) { |
| 154 | + sync(data); |
| 155 | + state.port?.on("message", sync); |
| 156 | + } |
| 157 | +
|
| 158 | + /** Returns the a special ID for known, bundled modules */ |
| 159 | + export async function resolve(id, context, next) { |
| 160 | + await fs.appendFile("loader.log", "loader resolve " + id + "\n"); |
| 161 | +
|
| 162 | + if (!state.enabled) return next(id, context); |
| 163 | + if (!state.modules.includes(id)) return next(id, context); |
| 164 | +
|
| 165 | + return { |
| 166 | + shortCircuit: true, |
| 167 | + url: 'bundled:' + id, |
| 168 | + }; |
| 169 | + } |
| 170 | +
|
| 171 | + /* Loads a bundled module using a global handler */ |
| 172 | + export async function load(url, context, next) { |
| 173 | + await fs.appendFile("loader.log", "loader load " + url + "\n"); |
| 174 | +
|
| 175 | + if (!state.enabled) return next(url, context); |
| 176 | + if (!url.startsWith("bundled:")) return next(url, context); |
| 177 | +
|
| 178 | + let id = url.slice(8); |
| 179 | + if (!state.modules.includes(id)) return next(url, context); |
| 180 | +
|
| 181 | + let source = 'export default globalThis.__tw_load__('; |
| 182 | + source += JSON.stringify(id); |
| 183 | + source += ')'; |
| 184 | +
|
| 185 | + return { |
| 186 | + shortCircuit: true, |
| 187 | + format: "module", |
| 188 | + source, |
| 189 | + }; |
| 190 | + } |
| 191 | +` |
0 commit comments